diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index 3a1ad0f0d23c..6e75506eac57 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/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" + ) diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 5b8e2bf008e9..ab840b4ff8dd 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/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,