From 6ac28f4dd3a9a34a2b4abcbe0a815fc59a2d74ad Mon Sep 17 00:00:00 2001
From: Zihao Ye <expye@outlook.com>
Date: Tue, 13 Aug 2024 00:54:27 -0700
Subject: [PATCH] bugfix: fix prefill kernels' lse result for empty kv-cache
 (#440)

Thank @hnyls2002 for spotting this bug.
---
 include/flashinfer/attention/prefill.cuh | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/include/flashinfer/attention/prefill.cuh b/include/flashinfer/attention/prefill.cuh
index 8910e097..8d78c938 100644
--- a/include/flashinfer/attention/prefill.cuh
+++ b/include/flashinfer/attention/prefill.cuh
@@ -1558,7 +1558,7 @@ __launch_bounds__(num_warps_x* num_warps_z* warp_size) void BatchPrefillWithRagg
   // normalize d
   normalize_d<num_frags_x, num_frags_y>(o_frag, m, d);
 
-  const uint32_t num_kv_chunks = ceil_div(kv_len, kv_chunk_size);
+  const uint32_t num_kv_chunks = ceil_div(max(kv_len, 1), kv_chunk_size);
 
   // write back
   write_o_reg_gmem<num_warps_x, num_warps_z, num_frags_x, num_frags_y>(
@@ -1872,7 +1872,7 @@ __launch_bounds__(num_warps_x* num_warps_z* warp_size) void BatchPrefillWithPage
   // normalize d
   normalize_d<num_frags_x, num_frags_y>(o_frag, m, d);
 
-  const uint32_t num_kv_chunks = ceil_div(kv_len, kv_chunk_size);
+  const uint32_t num_kv_chunks = ceil_div(max(kv_len, 1), kv_chunk_size);
 
   // write_back
   write_o_reg_gmem<num_warps_x, num_warps_z, num_frags_x, num_frags_y>(