bugfix: fix prefill kernels' lse result for empty kv-cache (#440)

Thank @hnyls2002 for spotting this bug.
flashinfer-ai · Aug 13, 2024 · 6ac28f4 · 6ac28f4
1 parent c93f647
commit 6ac28f4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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>(