[Bugfix][HiSparse] Pad staged prefill KV to avoid a trtllm-gen out-of-bounds read (#61659)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

tests/v1/attention/test_sparse_mla_backends.py

@@ -3659,6 +3659,33 @@ def test_hisparse_gather_prefill_cache_prefers_resident_rows():
         torch.testing.assert_close(staged_flat[i], expected)
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_hisparse_staged_prefill_cache_backs_min_attention_rows():
+    """A staged cache of a few blocks still has 256 allocated rows behind it.
+
+    trtllm-gen sparse MLA reads up to ~250 rows from the start of its KV cache
+    whatever its size, faulting when a smaller cache ends at an unmapped page.
+    """
+    if not _has_hisparse_ops():
+        pytest.skip("hisparse CUDA ops unavailable")
+    device = torch.device("cuda")
+    block_size, row_width = 64, 16
+    block_table = torch.tensor([[3, 1]], dtype=torch.int32, device=device)
+    seq_lens = torch.tensor([100], dtype=torch.int32, device=device)
+    plan = build_hisparse_prefill_staging_plan(
+        block_table, seq_lens, block_size, staging_block_capacity=2
+    )
+    host_cache = torch.zeros(4, block_size, row_width).pin_memory()
+
+    staged = HiSparseRuntime.gather_prefill_cache(None, host_cache, plan)
+
+    assert staged.shape[0] * block_size < 256
+    backing_rows = staged.untyped_storage().nbytes() // (
+        row_width * staged.element_size()
+    )
+    assert backing_rows >= 256
+
+
 def test_hisparse_fp8_decode_resolves_rows_once_then_runs_batched_attention():
     """Verification rows resolve in one swap_in; FP8 attention sees them batched."""
     device = torch.device("cpu")

vllm/v1/hisparse/runtime.py

@@ -39,6 +39,9 @@ if TYPE_CHECKING:
 # fp8_ds_mla KV row: 512 B quantized NoPE + 16 B scales + 128 B RoPE.
 FP8_DS_MLA_ROW_BYTES = 656
 HiSparseTopKResult: TypeAlias = torch.Tensor | tuple[torch.Tensor, torch.Tensor]
+# trtllm-gen sparse MLA reads up to ~250 rows from the start of its KV cache
+# whatever the cache's size: https://github.com/flashinfer-ai/flashinfer/issues/6484
+_MIN_STAGED_PREFILL_KV_ROWS = 256
 
 
 def is_hisparse_decode_batch(attn_metadata: Any | None) -> bool:
@@ -854,15 +857,13 @@ class HiSparseRuntime:
             )
         row_width = kv_cache.shape[-1]
 
+        num_blocks = plan.row_ids.shape[1] // plan.block_size
+        min_blocks = cdiv(_MIN_STAGED_PREFILL_KV_ROWS, plan.block_size)
         staged = torch.empty(
-            (
-                plan.row_ids.shape[1] // plan.block_size,
-                plan.block_size,
-                row_width,
-            ),
+            (max(num_blocks, min_blocks), plan.block_size, row_width),
             dtype=kv_cache.dtype,
             device=plan.block_table.device,
-        )
+        )[:num_blocks]
         if plan.gpu_row_ids is not None and resident_cache is not None:
             self.backend.gather_resident_rows(
                 resident_cache, staged.view(-1, row_width), plan.gpu_row_ids[0]