[Bugfix][DSA] Avoid runtime JIT for offset candidate end buffers (#57667)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Sep 19, 2026 at 00:27 UTC
38537331d3e8b37e480f1bbc174f78888e6c5252
2 files changed
+7
-4
tests/model_executor/layers/test_mla_short_prefill_indexer.py
+6
-3
@@ -628,7 +628,10 @@ def test_candidate_blocks_to_sparse_indices_math():
628
629
@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel")
630
@pytest.mark.parametrize("cbk,sbk", [(8, 8), (16, 8), (16, 16)])
631
-def test_candidate_blocks_to_sparse_indices_matches_reference(cbk: int, sbk: int):
631
+@pytest.mark.parametrize("end_offset", [0, 1])
632
+def test_candidate_blocks_to_sparse_indices_matches_reference(
633
+ cbk: int, sbk: int, end_offset: int
634
+):
635
"""The Triton expansion equals the PyTorch reference on production-shaped
636
random inputs, including strided (column-sliced) candidate rows and
637
strided output buffers."""
@@ -663,7 +666,7 @@ def test_candidate_blocks_to_sparse_indices_matches_reference(cbk: int, sbk: int
666
ke,
667
cbk,
668
sbk,
666
- out=(si_buf[:, : num_candidates * ratio], end_buf[::2]),
669
+ out=(si_buf[:, : num_candidates * ratio], end_buf[end_offset::2]),
670
)
671
assert torch.equal(got_indices, ref_indices)
672
assert torch.equal(got_end, ref_end)
@@ -671,7 +674,7 @@ def test_candidate_blocks_to_sparse_indices_matches_reference(cbk: int, sbk: int
674
cand_buf[:, num_candidates:], torch.full_like(cand_buf[:, num_candidates:], -7)
675
)
676
assert not si_buf[:, num_candidates * ratio :].any()
674
- assert not end_buf[1::2].any()
677
+ assert not end_buf[1 - end_offset :: 2].any()
678
679
680
def _skip_unless_sm100_sparse_kernels():
vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py
+1
-1
@@ -73,7 +73,7 @@ def check_deep_select_layout(num_sparse_cols: int, topk_tokens: int) -> None:
73
_INT32_BIG = tl.constexpr(2**31 - 1)
74
75
76
-@triton.jit
76
+@triton.jit(do_not_specialize_on_alignment=["end_ptr"])
77
def _expand_candidates_kernel(
78
cand_ptr,
79
cand_stride,