[Bugfix][DSv4.1] Preserve NaN-scored candidate block indices (#57454)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Codex <noreply@openai.com>
Woosuk Kwon committed
Sep 17, 2026 at 19:00 UTC
64563d0ec4b761745d5babbc5750b263f0783926
2 files changed
+23
-2
tests/model_executor/layers/test_mla_short_prefill_indexer.py
+21
-1
@@ -365,6 +365,26 @@ def test_select_candidate_blocks_tolerates_empty_rows():
365
assert (out[0] >= 0).all() and (out[2, :2] >= 0).all()
366
367
368
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
369
+@pytest.mark.parametrize("score", [0.0, float("nan")])
370
+@pytest.mark.parametrize("start", [0, 3])
371
+def test_candidate_selection_keeps_valid_blocks_with_nan_scores(score, start):
372
+ """NaN warmup scores must not turn real blocks into repeated sparse padding."""
373
+ logits = torch.full((3, 64), score, device="cuda")
374
+ starts = torch.full((3,), start, device="cuda", dtype=torch.int32)
375
+ lengths = torch.tensor([0, 17, 35], device="cuda", dtype=torch.int32)
376
+ out = torch.empty(3, 8, device="cuda", dtype=torch.int32)
377
+
378
+ sparse_indexer._select_candidate_blocks(logits, starts, starts + lengths, 8, 8, out)
379
+
380
+ for row, count in enumerate([0, 3, 5]):
381
+ valid = out[row][out[row] >= 0].sort().values
382
+ torch.testing.assert_close(
383
+ valid, torch.arange(count, device="cuda", dtype=torch.int32)
384
+ )
385
+ assert (out[row] == -1).sum().item() == 8 - count
386
+
387
+
388
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
389
@pytest.mark.parametrize(
390
"width,block_size,k,decode",
@@ -413,7 +433,7 @@ def test_candidate_kernels_preserve_packed_bounds_and_padding(
433
top = reduced.topk(min(k, nblocks), dim=-1)
434
expected = torch.full((rows, k), -1, device="cuda", dtype=torch.int32)
435
expected[:, : top.indices.shape[1]] = torch.where(
416
- top.values > -torch.inf, top.indices, -1
436
+ top.values != -torch.inf, top.indices, -1
437
).int()
438
actual = torch.empty(rows, k * 2, device="cuda", dtype=torch.int32)[:, ::2]
439
select_candidate_blocks(logits, starts, ends, k, block_size, actual, repeat)
vllm/model_executor/kernels/attention/dsa/candidate_blocks.py
+2
-1
@@ -66,9 +66,10 @@ def _store_candidates_kernel(
66
cols = tl.program_id(1) * TILE + tl.arange(0, TILE)
67
value = tl.load(values + row * k + cols, cols < k, other=-float("inf"))
68
index = tl.load(indices + row * k + cols, cols < k, other=-1)
69
+ # NaN scores can occur during warmup; only -inf denotes padding.
70
tl.store(
71
output + row * out_stride_row + cols * out_stride_col,
71
- tl.where(value > -float("inf"), index, -1),
72
+ tl.where(value != -float("inf"), index, -1),
73
cols < OUT_K,
74
)
75