[Bugfix] Route DSv4 sparse-indexer prefill top-k around NaN-broken kernel path on SM12x

tests/kernels/test_top_k_per_row.py

@@ -968,3 +968,59 @@ def test_workspace_topk_padded_stride(top_k: int, backend: str) -> None:
                 f"Row {i}: {backend} with padded stride doesn't match. "
                 f"seq_len={sl}, stride={padded_stride}"
             )
+
+
+@pytest.mark.parametrize("num_rows", NUM_ROWS)
+@pytest.mark.parametrize("top_k", [512, 2048])
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
+@torch.inference_mode()
+def test_top_k_per_row_prefill_torch_fallback(
+    num_rows: int,
+    top_k: int,
+) -> None:
+    """The SM12x torch fallback honors top_k_per_row_prefill's contract:
+    the top-k positions of logits[i, ks_i:ke_i) relative to ks_i, in
+    ascending position order, -1-padded when a row has fewer candidates,
+    with NaN scores never selected."""
+    from vllm.model_executor.layers.sparse_attn_indexer import (
+        _top_k_per_row_prefill_torch,
+    )
+
+    set_random_seed(0)
+    torch.set_default_device("cuda:0")
+
+    vocab_size = 20000
+    row_starts, row_ends = create_row_boundaries(num_rows, vocab_size)
+    logits = create_random_logits(
+        row_starts, row_ends, torch.float32, 42, True, "random"
+    )
+    # Sprinkle NaNs inside some valid spans: the fallback must skip them
+    # (the failure mode on SM12x is NaN-dominated logits).
+    if num_rows > 2:
+        logits[2, : int(row_ends[2])] = float("nan")
+
+    indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
+    # The fallback masks logits in place; give it a scratch copy.
+    _top_k_per_row_prefill_torch(
+        logits.clone(), row_starts, row_ends, indices, top_k
+    )
+
+    for i in range(num_rows):
+        row_start = int(row_starts[i])
+        row_end = int(row_ends[i])
+        row = indices[i].cpu()
+        finite = torch.isfinite(logits[i, row_start:row_end]).sum().item()
+        num_valid = min(top_k, int(finite))
+        selected = row[:num_valid]
+        # Ascending position order, in range, no duplicates.
+        assert (selected.diff() > 0).all(), f"row {i} not ascending"
+        assert (selected >= 0).all() and (selected < row_end - row_start).all()
+        # Everything past the valid entries is -1 padding.
+        assert (row[num_valid:] == -1).all(), f"row {i} padding broken"
+        # Set-equality with the torch.topk reference over the valid span.
+        if num_valid:
+            span = logits[i, row_start:row_end].nan_to_num(float("-inf"))
+            ref = span.topk(num_valid, dim=-1)[1].cpu()
+            assert set(selected.tolist()) == set(ref.tolist()), (
+                f"row {i} selected wrong candidate set"
+            )

vllm/model_executor/layers/sparse_attn_indexer.py

@@ -124,6 +124,53 @@ def _merge_dcp_topk_global(
     )
 
 
+def _top_k_per_row_prefill_torch(
+    logits: torch.Tensor,
+    cu_seqlen_ks: torch.Tensor,
+    cu_seqlen_ke: torch.Tensor,
+    topk_indices: torch.Tensor,
+    topk_tokens: int,
+) -> None:
+    """torch.topk fallback with ``top_k_per_row_prefill``'s output contract.
+
+    On SM12x the CUDA kernel's histogram path (taken by rows with more than
+    ``topk_tokens`` candidates) can leave part of its dynamic-shared-memory
+    output uninitialized and copy it out as indices; downstream,
+    ``compute_global_topk_indices_and_lens`` treats any index ``>= 0`` as
+    valid and dereferences it into the KV block table, crashing with a CUDA
+    illegal memory access. Until the kernel is fixed for SM12x, select the
+    top ``topk_tokens`` positions of ``logits[i, ks_i:ke_i)`` per row with
+    torch.topk, emit them relative to ``ks_i`` in ascending position order,
+    and pad rows with fewer candidates with ``-1`` (the kernel's short-row
+    contract).
+    """
+    num_cols = logits.shape[1]
+    ks = cu_seqlen_ks.to(torch.long)[:, None]
+    cols = torch.arange(num_cols, device=logits.device)[None, :]
+    valid = (cols >= ks) & (cols < cu_seqlen_ke.to(torch.long)[:, None])
+    # logits is a per-chunk scratch buffer that is dead after top-k; mask it
+    # in place rather than materializing a masked copy.
+    logits.masked_fill_(~valid, float("-inf"))
+    k = min(topk_tokens, num_cols)
+    top_values, top_cols = logits.topk(k, dim=-1)
+    relative = (top_cols - ks).to(torch.int32)
+    # Downstream sparse-attention kernels iterate the selected KV positions
+    # in ascending order; emit position-sorted indices with the ``-1`` pads
+    # at the tail (torch.topk returns score order instead). Non-finite
+    # scores (rows shorter than ``topk_tokens``, or NaN logits) pad.
+    pad_sentinel = torch.iinfo(torch.int32).max
+    relative = torch.where(
+        top_values.isfinite(), relative, relative.new_full((), pad_sentinel)
+    )
+    relative, _ = relative.sort(dim=-1)
+    relative = torch.where(
+        relative == pad_sentinel, relative.new_full((), -1), relative
+    )
+    topk_indices[:, :k] = relative
+    if k < topk_tokens:
+        topk_indices[:, k:] = -1
+
+
 @triton.jit
 def _fused_indexer_q_rope_quant_kernel(
     positions,
@@ -485,16 +532,32 @@ def sparse_attn_indexer(
                         clean_logits=False,
                     )
                 num_rows = logits.shape[0]
-                ops.top_k_per_row_prefill(
-                    logits,
-                    cu_seqlen_ks,
-                    cu_seqlen_ke,
-                    topk_indices,
-                    num_rows,
-                    logits.stride(0),
-                    logits.stride(1),
-                    topk_tokens,
-                )
+                if current_platform.is_cuda() and (
+                    current_platform.is_device_capability_family(120)
+                ):
+                    # top_k_per_row_prefill's histogram path emits
+                    # uninitialized shared memory as indices on SM12x (see
+                    # _top_k_per_row_prefill_torch); the decode path already
+                    # excludes SM12x from cooperative_topk for the same
+                    # kernel family.
+                    _top_k_per_row_prefill_torch(
+                        logits,
+                        cu_seqlen_ks,
+                        cu_seqlen_ke,
+                        topk_indices,
+                        topk_tokens,
+                    )
+                else:
+                    ops.top_k_per_row_prefill(
+                        logits,
+                        cu_seqlen_ks,
+                        cu_seqlen_ke,
+                        topk_indices,
+                        num_rows,
+                        logits.stride(0),
+                        logits.stride(1),
+                        topk_tokens,
+                    )
 
             _merge_dcp_topk_global(
                 logits,