compatible with flashinfer v0.2 (#3235)

sgl-project · Jan 31, 2025 · 7811bfd · 7811bfd
1 parent 656f7fc
commit 7811bfd
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py
@@ -800,7 +800,9 @@ def call_begin_forward(
             kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0)
             kv_indptr = kv_indptr[: bs + 1]
             kv_indices = torch.empty(
-                paged_kernel_lens_sum, dtype=torch.int32, device="cuda"
+                paged_kernel_lens_sum + 256,
+                dtype=torch.int32,
+                device=req_pool_indices.device,
             )
             create_flashinfer_kv_indices_triton[(bs,)](
                 self.req_to_token,