[Bugfix] Make DSA prefill top-k deterministic on equal scores

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,
+        )

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,
                 )
 

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