fix

csrc/pod.cu

@@ -174,9 +174,7 @@ void PODWithPagedKVCacheRun(at::Tensor float_workspace_buffer, at::Tensor int_wo
           params.paged_kv = paged_kv;
 
           if (plan_info.split_kv) {
-            params.partition_kv = true;
-            params.merge_indptr =
-                GetPtrFromBaseOffset<IdType>(int_buffer_ptr, plan_info.merge_indptr_offset);
+            params.partition_kv = true;  // used in prefill kernel
             tmp_v = GetPtrFromBaseOffset<DTypeO>(float_buffer_ptr, plan_info.v_offset);
             tmp_s = GetPtrFromBaseOffset<float>(float_buffer_ptr, plan_info.s_offset);
             if (plan_info.enable_cuda_graph) {
@@ -195,6 +193,13 @@ void PODWithPagedKVCacheRun(at::Tensor float_workspace_buffer, at::Tensor int_wo
         _configureParams(params_p, pod_plan_info.plan_info_p, pod_plan_info.batch_size_vec_p);
         _configureParams(params_d, pod_plan_info.plan_info_d, pod_plan_info.batch_size_vec_d);
 
+        if (pod_plan_info.plan_info_p.split_kv || pod_plan_info.plan_info_d.split_kv) {
+          params_p.merge_indptr = GetPtrFromBaseOffset<IdType>(
+              int_buffer_ptr, pod_plan_info.plan_info_p.merge_indptr_offset);
+          params_d.merge_indptr = GetPtrFromBaseOffset<IdType>(
+              int_buffer_ptr, pod_plan_info.plan_info_d.merge_indptr_offset);
+        }
+
         cudaError_t status = cudaSuccess;
 
         DISPATCH_CTA_TILE_Q(pod_plan_info.plan_info_p.cta_tile_q, CTA_TILE_Q_P, {

flashinfer/pod.py

@@ -665,6 +665,7 @@ def run(
                 sm_scale,
                 rope_scale,
                 rope_theta,
+                get_cuda_stream(q.device),
             ]
 
         self._cached_module.paged_run(*run_args)

include/flashinfer/attention/scheduler.cuh

@@ -1425,8 +1425,13 @@ struct PoDPlanInfo {
     }
     // Append boundary element
     // used in produce_kv_page
-    kv_start_ptr_h_p.push_back(kv_start_ptr_h_p.back() + kv_len_ptr_h_p.back());
-    kv_start_ptr_h_d.push_back(kv_start_ptr_h_d.back() + kv_len_ptr_h_d.back());
+    if (!kv_start_ptr_h_p.empty()) {
+      kv_start_ptr_h_p.push_back(kv_start_ptr_h_p.back() + kv_len_ptr_h_p.back());
+    }
+    if (!kv_start_ptr_h_d.empty()) {
+      kv_start_ptr_h_d.push_back(kv_start_ptr_h_d.back() + kv_len_ptr_h_d.back());
+    }
+
     return std::make_tuple(partition_bitmask, qo_start_ptr_h_p, qo_len_ptr_h_p, kv_start_ptr_h_p,
                            kv_len_ptr_h_p, kv_last_page_len_h_p, qo_start_ptr_h_d, qo_len_ptr_h_d,
                            kv_start_ptr_h_d, kv_len_ptr_h_d, kv_last_page_len_h_d);