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

[mimeding]

Purpose

The same long prompt should not change which compressed-context tokens sparse attention uses, or the generated answer, just because several indexer scores were exactly equal. Today the DeepSeek prefill path calls ops.top_k_per_row_prefill. On equal finite logits that kernel can emit different index sets for the same row when the row batch changes, so later attention sees a different KV subset.

This PR replaces that prefill call with a stable value-desc / index-asc selector: columns in [ks, ke) only, equal scores keep the smaller column index, unused slots stay -1.

Changed behavior

• File: vllm/model_executor/layers/sparse_attn_indexer.py (code owner @zyongye).
• Prefill only. Decode still uses indexer_topk.py (persistent_topk / cooperative_topk / top_k_per_row_decode / …).
• No new CUDA kernel. The selector is torch.argsort(..., descending=True, stable=True) plus a validity mask, matching a production wrap that overwrote the CUDA result.

Related work (not duplicates)

Checked open PRs/issues for top_k_per_row_prefill, prefill topk tie, and #53691.

• [Kernel] Make persistent_topk deterministic #55122 makes persistent_topk deterministic. Different registered operator; its topk.cu / persistent_topk.cuh change does not cover this prefill path.
• [Bugfix][Kernel] Fix top-k selection when the radix threshold bin overflows the smem stash #55314 repairs radix threshold-bin overflow and permits an arbitrary subset after full-key equality, so it does not guarantee this exact-tie membership.
• [Perf] Use 1024 threads per block in the prefill top-k kernel #55363 changes prefill kernel thread count; correctness is still torch.topk set match, not index-stable ties.
• [Kernel] Harden top_k_per_row against NaN and under-filled output #50201 hardens NaN / under-filled output. [Bugfix] Route DSv4 sparse-indexer prefill top-k around NaN-broken kernel path on SM12x #49897 is the SM12x NaN crash fallback.
• [Performance]: DSA indexer is computed redundantly on every TP rank (NVIDIA / CUDA) #53691 (closed; follow-up [Perf][Attention][DSA] Shard prefill indexer rows across TP ranks #54394) includes an offline harness where injected ties make top_k_per_row_prefill selection sets differ under row batching. Their live H100 sample found no boundary ties. This PR does not implement TP sharding.
• [Bugfix] Honor DeepSeek-V4 enable_thinking precedence #56708 is unrelated DeepSeek-V4 enable_thinking precedence. Not stacked with this change.

Validation

CPU, macOS arm64, Python 3.12, PyTorch 2.14.0, no CUDA. Source checkout with PYTHONPATH, not a native vLLM GPU build (vllm._version warning):

PYTHONPATH=. .venv/bin/python -m pytest -q --confcutdir=tests/kernels \
  tests/kernels/test_sparse_prefill_topk.py
# 10 passed in 0.75s

Cases: equal finite scores keep the smaller index; [ks, ke) clipping; empty range; k wider than the valid range; valid -inf vs invalid columns; two-row ties; 244 vs 640 boundary tie; 1-row vs batched-row identity under ties; 20-repeat identity; CPU inverted-bounds error.

Changed-file pre-commit (--hook-stage manual, three paths): ruff check/format, typos, mypy 3.10–3.13, SPDX, root lazy imports, forbidden imports, torch.cuda API, config validation, boolean-with. compileall and git diff --check passed.

Not run on this branch: CUDA tests/kernels/test_top_k_per_row.py, current-main GPU serving, or model evals. Existing kernel tests still allow set-equality under ties for the CUDA op; that op is no longer used on this DeepSeek prefill path.

Supporting deployment evidence (different tree, NVIDIA vLLM 0.27.1 overlay, wrap image b15aa6ff): the same selector contract was 24/24 byte-stable across 10 runs / two serving starts and 48/48 long retrieval including 120k. That is not a general fidelity claim and is not a current-main GPU result.

Limits

• GLM kpool and ROCm aiter prefill still call top_k_per_row_prefill.
• torch.argsort over the logit row is a correctness-first replacement; this PR does not claim a prefill latency win.
• No invented GPU algorithm.

AI assistance and review status

Grok 4.6 assisted with duplicate-work comparison against current origin/main (52dd0d7562adb3c2d556ce6fe7ca7c3226e1976b), implementation, CPU tests, and this description. Per AGENTS.md, a human submitter must review every changed line and run relevant tests; that human line-by-line review and human-run GPU/model tests are not claimed here. This PR is draft until that happens.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀