bugfix: fix block sparse wrappers (#556)

The block sparse attention unittests failed as noted in #554, this PR fixes the issue.
flashinfer-ai · Oct 25, 2024 · 2989556 · 2989556
1 parent 89f2c4a
commit 2989556
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/python/flashinfer/sparse.py b/python/flashinfer/sparse.py
@@ -297,6 +297,8 @@ def plan(
         self.R = R
         self.C = C
 
+        kv_indptr_host = indptr.to("cpu", non_blocking=True)
+
         self._cached_module = get_batch_prefill_module(
             q_data_type,
             kv_data_type,
@@ -314,8 +316,8 @@ def plan(
             self._float_workspace_buffer,
             self._int_workspace_buffer,
             self._pin_memory_int_workspace_buffer,
-            qo_indptr,
-            indptr,
+            qo_indptr_host,
+            kv_indptr_host,
             num_blocks_row,
             num_qo_heads,
             num_kv_heads,
@@ -407,7 +409,6 @@ def run(
             self._int_workspace_buffer,
             self._plan_info,
             q,
-            None,
             k,
             v,
             self._packed_mask_buf,