[Bugfix][DSA] Bound DeepSelect sentinel columns in the sparse top-k remap (#58215)
Signed-off-by: Juntian Liu <Juntianl777@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Juntian Liu committed
Sep 23, 2026 at 17:54 UTC
9ef37771beac1c1ce6a8b7ceae1f7ebeb6f51800
2 files changed
+39
-2
tests/model_executor/layers/test_mla_short_prefill_indexer.py
+34
@@ -677,6 +677,40 @@ def test_candidate_blocks_to_sparse_indices_matches_reference(
677
assert not end_buf[1 - end_offset :: 2].any()
678
679
680
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
681
+def test_sparse_topk_remap_bounds_nan_rows():
682
+ """DeepSelect writes an out-of-range sentinel to slot 0 of NaN rows (the
683
+ rest is left untouched); the remap turns it into -1."""
684
+ from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import (
685
+ has_deep_select,
686
+ sparse_topk_remap,
687
+ )
688
+
689
+ if not has_deep_select():
690
+ pytest.skip("DeepSelect extension (vllm._deepselect_C) required")
691
+ rows, num_sparse, sbk, topk = 4, 128, 8, 512
692
+ width = num_sparse * sbk
693
+ logits = torch.randn(rows, width, dtype=torch.bfloat16, device="cuda")
694
+ logits[1:, :16] = float("nan")
695
+ sparse_indices = torch.arange(num_sparse, dtype=torch.int32, device="cuda")
696
+ end = torch.full((rows,), width, dtype=torch.int32, device="cuda")
697
+ out = torch.empty(rows, topk, dtype=torch.int32, device="cuda")
698
+ cols = torch.zeros_like(out)
699
+
700
+ sparse_topk_remap(
701
+ logits,
702
+ sparse_indices.repeat(rows, 1),
703
+ end,
704
+ torch.zeros_like(end),
705
+ sbk,
706
+ topk,
707
+ out,
708
+ col_indices=cols,
709
+ )
710
+ assert (cols[1:, 0] >= width).all() and (out[1:, 0] == -1).all()
711
+ assert (out[0] >= 0).all()
712
+
713
+
714
def _skip_unless_sm100_sparse_kernels():
715
if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10:
716
pytest.skip("SM100 required")
vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py
+5
-2
@@ -146,19 +146,21 @@ def _sparse_topk_remap_kernel(
146
out_ptr,
147
out_stride,
148
k,
149
+ width,
150
SBK: tl.constexpr,
151
K_POW2: tl.constexpr,
152
):
153
"""Remap sparse top-k columns to request-local positions.
154
155
Column ``j * SBK + o`` scores the token at
155
- ``(si[row, j] - ks // SBK) * SBK + o``; -1 columns stay -1.
156
+ ``(si[row, j] - ks // SBK) * SBK + o``; columns outside ``[0, width)``
157
+ (including DeepSelect's sentinel for NaN rows) become -1.
158
"""
159
row = tl.program_id(0)
160
ks_base = tl.load(ks_ptr + row * ks_stride) // SBK
161
cols = tl.arange(0, K_POW2)
162
c = tl.load(col_ptr + row * col_stride + cols, mask=cols < k, other=-1)
161
- valid = c >= 0
163
+ valid = (c >= 0) & (c < width)
164
cc = tl.where(valid, c, 0)
165
blocks = tl.load(si_ptr + row * si_stride + cc // SBK)
166
pos = (blocks - ks_base) * SBK + cc % SBK
@@ -301,6 +303,7 @@ def sparse_topk_remap(
303
topk_indices,
304
topk_indices.stride(0),
305
topk_tokens,
306
+ width,
307
SBK=sparse_block_kv,
308
K_POW2=triton.next_power_of_2(topk_tokens),
309
)