diff --git a/tests/kernels/test_sparse_prefill_topk.py b/tests/kernels/test_sparse_prefill_topk.py new file mode 100644 index 000000000000..889f1bd8ba10 --- /dev/null +++ b/tests/kernels/test_sparse_prefill_topk.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU contract for deterministic DSA prefill top-k. + +Guards equal-score membership: the same valid range must keep the same +compressed-context indices, including when a row is batched with others. +The CUDA top_k_per_row_prefill kernel is not exercised here. +""" + +import pytest +import torch + +from vllm.model_executor.layers.sparse_prefill_topk import ( + stable_prefill_topk_from_valid_range, +) + + +def _run( + logits: torch.Tensor, + ks: list[int], + ke: list[int], + topk_tokens: int, +) -> list[list[int]]: + idx = torch.full( + (logits.shape[0], topk_tokens), + -1, + dtype=torch.int32, + ) + stable_prefill_topk_from_valid_range( + logits, + torch.tensor(ks, dtype=torch.int32), + torch.tensor(ke, dtype=torch.int32), + idx, + topk_tokens, + ) + return idx.tolist() + + +def test_equal_finite_scores_keep_smaller_index() -> None: + logits = torch.tensor([[1.0, 5.0, 5.0, 0.0, 9.0]]) + assert _run(logits, [0], [5], 3) == [[4, 1, 2]] + + +def test_valid_range_excludes_columns_outside_ks_ke() -> None: + logits = torch.tensor([[9.0, 1.0, 3.0, 3.0, 0.0]]) + assert _run(logits, [1], [5], 2) == [[2, 3]] + + +def test_empty_range_writes_padding() -> None: + logits = torch.tensor([[1.0, 2.0, 3.0]]) + assert _run(logits, [2], [2], 2) == [[-1, -1]] + + +def test_k_wider_than_valid_range_pads() -> None: + logits = torch.tensor([[4.0, 1.0, 2.0]]) + assert _run(logits, [0], [2], 4) == [[0, 1, -1, -1]] + + +def test_valid_neg_inf_is_not_replaced_by_invalid_column() -> None: + logits = torch.tensor([[100.0, float("-inf"), 1.0, 50.0]]) + assert _run(logits, [1], [3], 2) == [[2, 1]] + + +def test_two_row_ties_keep_lower_index_per_row() -> None: + logits = torch.tensor([[1.0, 1.0, 0.0], [0.0, 5.0, 5.0]]) + assert _run(logits, [0, 1], [3, 3], 2) == [[0, 1], [1, 2]] + + +def test_boundary_tie_prefers_lower_compressed_index() -> None: + """Equal finite scores at 244 and 640; k=1 must keep 244.""" + cols = 641 + logits = torch.full((1, cols), -2.0) + logits[0, 244] = 1.0 + logits[0, 640] = 1.0 + assert _run(logits, [0], [cols], 1) == [[244]] + assert _run(logits, [0], [cols], 2) == [[244, 640]] + + +def test_batched_rows_match_single_row_under_ties() -> None: + """Same logits must not change selection when the row batch changes.""" + cols = 32 + k = 4 + row0 = torch.zeros(cols) + # Four-way tie at the cut, plus two unique higher scores. + row0[3] = 9.0 + row0[11] = 8.0 + row0[5] = 1.0 + row0[7] = 1.0 + row0[13] = 1.0 + row0[19] = 1.0 + row1 = torch.arange(cols, dtype=torch.float32) + solo = _run(row0.unsqueeze(0), [0], [cols], k) + batched = _run(torch.stack([row0, row1]), [0, 0], [cols, cols], k) + assert solo[0] == batched[0] + assert solo[0] == [3, 11, 5, 7] + + +def test_selection_is_repeatable() -> None: + logits = torch.tensor([[1.0, 5.0, 5.0, 0.0, 9.0], [0.0, 5.0, 5.0, 5.0, 1.0]]) + first = _run(logits, [0, 1], [5, 5], 3) + for _ in range(19): + assert _run(logits, [0, 1], [5, 5], 3) == first + assert first == [[4, 1, 2], [1, 2, 3]] + + +def test_cpu_rejects_inverted_bounds() -> None: + logits = torch.tensor([[1.0, 2.0, 3.0]]) + idx = torch.full((1, 2), -1, dtype=torch.int32) + with pytest.raises(ValueError, match="ks/ke"): + stable_prefill_topk_from_valid_range( + logits, + torch.tensor([2], dtype=torch.int32), + torch.tensor([1], dtype=torch.int32), + idx, + 2, + ) diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 32902901c72b..eaffa7e02d25 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -27,6 +27,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( get_fp8_min_max, ) +from vllm.model_executor.layers.sparse_prefill_topk import ( + stable_prefill_topk_from_valid_range, +) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.deep_gemm import ( @@ -589,7 +592,6 @@ def sparse_attn_indexer( cu_seqlen_ke, clean_logits=False, ) - num_rows = logits.shape[0] if candidate_blocks is not None: # Two-level selection (v4.1): the candidate source # publishes its top blocks; later indexers mask their @@ -614,14 +616,13 @@ def sparse_attn_indexer( chunk_candidates, candidate_block_size, ) - ops.top_k_per_row_prefill( + # Value-desc / index-asc. CUDA top_k_per_row_prefill is not + # membership-stable when equal finite scores sit at the cut. + stable_prefill_topk_from_valid_range( logits, cu_seqlen_ks, cu_seqlen_ke, topk_indices, - num_rows, - logits.stride(0), - logits.stride(1), topk_tokens, ) diff --git a/vllm/model_executor/layers/sparse_prefill_topk.py b/vllm/model_executor/layers/sparse_prefill_topk.py new file mode 100644 index 000000000000..80d10dd4a8e7 --- /dev/null +++ b/vllm/model_executor/layers/sparse_prefill_topk.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Deterministic prefill top-k for the DSA sparse attention indexer.""" + +import torch + + +def stable_prefill_topk_from_valid_range( + logits: torch.Tensor, + cu_seqlen_ks: torch.Tensor, + cu_seqlen_ke: torch.Tensor, + topk_indices: torch.Tensor, + topk_tokens: int, +) -> None: + """Select per-row top-k indices over columns in ``[ks, ke)``. + + Invalid columns are dropped after a stable descending argsort, so a + valid ``-inf`` score cannot lose to a masked column. Equal finite + scores keep the smaller column index. Writes into ``topk_indices``; + unused slots are ``-1``. + + Args: + logits: Prefill indexer scores, shape ``(num_rows, cols)``. + cu_seqlen_ks: Inclusive start column per row. + cu_seqlen_ke: Exclusive end column per row. + topk_indices: Output index buffer, shape + ``(num_rows, >= topk_tokens)``. + topk_tokens: Number of indices to write per row. + + Raises: + ValueError: If shapes are inconsistent, or CPU ``ks``/``ke`` + fall outside the logit columns. + """ + if logits.ndim != 2 or topk_indices.ndim != 2: + raise ValueError("stable prefill topk expects 2D logits and indices") + num_rows, cols = logits.shape + if topk_indices.shape[0] != num_rows: + raise ValueError("topk rows differ from logits") + if topk_tokens <= 0 or topk_indices.shape[1] < topk_tokens: + raise ValueError("topk_tokens wider than index buffer") + ks = cu_seqlen_ks.to(device=logits.device, dtype=torch.int64).reshape(-1) + ke = cu_seqlen_ke.to(device=logits.device, dtype=torch.int64).reshape(-1) + if ks.numel() != num_rows or ke.numel() != num_rows: + raise ValueError("ks/ke rows differ from logits") + # CPU contract checks stay strict. On CUDA, skip .any()/.item() host + # syncs; production ks/ke come from scheduler metadata and the mask + # below only keeps columns in [ks, ke). + if ks.device.type == "cpu" and bool( + (ks < 0).any() or (ke < ks).any() or (ke > cols).any() + ): + raise ValueError("ks/ke outside logits columns") + k = min(int(topk_tokens), cols) + if cols >= (1 << 30): + raise ValueError("logits wider than packed column field") + if num_rows == 0 or cols == 0 or k == 0: + if k: + topk_indices[:, :k] = -1 + return + col = torch.arange(cols, device=logits.device, dtype=torch.int64) + ok = (col[None, :] >= ks[:, None]) & (col[None, :] < ke[:, None]) + # Invalid columns are filled with -inf for the argsort, then dropped + # so a valid -inf cannot lose to an invalid column. + masked = logits.masked_fill(~ok, float("-inf")) + order = torch.argsort(masked, dim=1, descending=True, stable=True) + is_valid = ok.gather(1, order) + rank = is_valid.to(torch.int64).cumsum(dim=1) + slot = torch.where(is_valid, rank - 1, torch.full_like(rank, k)) + chosen = torch.full( + (num_rows, k), + -1, + dtype=topk_indices.dtype, + device=logits.device, + ) + rows_i = torch.arange(num_rows, device=logits.device)[:, None].expand_as(order) + sel = slot < k + chosen[rows_i[sel], slot[sel]] = order[sel].to(topk_indices.dtype) + width = (ke - ks).clamp(min=0) + keep = torch.arange(k, device=logits.device)[None, :] < width[:, None] + chosen = torch.where(keep, chosen, torch.full_like(chosen, -1)) + topk_indices[:, :k] = chosen