[Perf][DSA] Use DeepSelect for eligible sparse-indexer prefill chunks

[WoosukKwon]

Purpose

Extend the existing opt-in deep_select sparse-indexer backend to supported prefill chunks. Today the option accelerates decode, while prefill still calls top_k_per_row_prefill. For eligible chunks, reuse DeepSelect with the causal per-row end bounds and the existing preallocated output buffer.

The fast path requires a single request in the indexer chunk, SM100-family hardware, no PCP or DCP, and supported input/output alignment. A single request's gathered KV starts at column zero, so DeepSelect's end argument represents the original causal interval. Multi-request chunks, context parallelism, and unsupported layouts retain the existing kernel. A multi-request batch can still benefit when the chunk planner places each request in its own chunk.

Enable with --kernel-config '{"sparse_indexer_topk_backend":"deep_select"}'. Defaults and decode dispatch are unchanged. No new kernels or dependency changes.

Duplicate-work check

Follow-up to merged #56464, which integrated DeepSelect for decode. Checked #56217 and its comments, open PRs referring to that issue, and DeepSelect/prefill-top-k searches.

• [Perf] Use 1024 threads per block in the prefill top-k kernel #55363 changes the existing prefill kernel's thread count; this selects DeepSelect through the existing explicit backend option.
• [Bugfix] Make DSA prefill top-k deterministic on equal scores #56749 introduces stable tie-breaking with a different selector; this does not promise index-stable ties.
• [Perf][PCP] Enable TP-sharded Indexer prefill and direct TopK publication #56258 changes context-parallel prefill sharding/publication; this fast path excludes context parallelism.
• [Perf][DSA] Add FlashInfer gvr_2 as a sparse-indexer decode top-k backend #57109 adds a decode GVR backend; this extends prefill only.

None implements this DeepSelect prefill path.

Test Plan

.venv/bin/python -m pytest tests/kernels/test_top_k_per_row.py -k deep_select_topk -q
pre-commit run --files vllm/config/kernel.py vllm/model_executor/layers/sparse_attn_indexer.py tests/kernels/test_top_k_per_row.py
pre-commit run mypy-3.12 --files vllm/config/kernel.py vllm/model_executor/layers/sparse_attn_indexer.py tests/kernels/test_top_k_per_row.py --hook-stage manual

Extend the existing preallocated-output test to cover bounded selection and CUDA graph replay after updating row ends. Includes empty/short rows, lengths around top-k, full-length rows, strided output, untouched padding, and -1 tails. Compare selected scores against PyTorch.

Test Result

Standalone commit 3e850c278d, based on 42919b49c5: 26 GPU tests passed on GB200 (250 deselected). Applicable pre-commit hooks, including mypy 3.10, and explicit mypy 3.12 passed. The three changed files exactly match the previously validated patch.

Earlier validation of the same patch included three full-model parity probes: 4K prefill/~100K KV, eight 1K prefills/~100K KV, and a mixed-chunk case exercising both DeepSelect and the nonzero-start fallback. All observed output errors were zero.

Performance evidence and limitations

The earlier approximately 6.3x top-k speedup is a dummy-model benchmark result, not a demonstrated real-weight serving speedup. Those profiles used v41_benchmark_dummy weights and generated KV history. The 4K-query/~100K-KV profile measured 7.52 ms versus 1.19 ms summed over ten prefill top-k launches per forward; both arms had matching launch counts and row counts. It was not a 6.3x full-model speedup. The previously quoted exploratory 8.0%/10.8% full-forward gains also predate the benchmark metadata-reset correction and should not be treated as validated production gains.

Captured real-model logits

We repeated kernel timing on logits captured from actual deepseek-ai/DeepSeek-V4.1-Flash weights (revision dba1be0a40aa45a94ad051997016db3960a90277), TP4/EP4, MegaMoE, FlashInfer FP8 KV and the MXFP4 NEW sparse indexer. Prompts contained repository documentation, with actual context lengths of 8,192, 51,200 and 102,400 tokens. We captured the first eligible matrix at each column width; these samples do not cover every layer or prompt distribution.

Logits columns	Query rows in captured chunk	Existing prefill top-k	DeepSelect	Native / DeepSelect
4,096	8,192	110.4 us	193.6 us	0.57x
8,192	8,192	135.6 us	227.8 us	0.60x
25,600	2,048	111.5 us	151.8 us	0.73x
51,200	2,048	176.7 us	195.4 us	0.90x
102,400	1,310	233.7 us	180.9 us	1.29x

Here “native” means vLLM's existing CUDA top_k_per_row_prefill, not PyTorch or FlashInfer. Both selectors used identical captured FP32 logits, causal bounds, top-k 512 and preallocated outputs. Measurements used a GB200 on gb200-rack1-04, CUPTI CUDA-graph timing with cold L2, 25 warmup calls and four alternating timing rounds. Selected-score equality, valid/unique indices and short-row tails passed for every case before timing. Matched normal-random controls produced similar trends.

The original dummy-input result overstated the benefit for these sampled real-model workloads. DeepSelect regressed on the captured 4K–50K-column matrices and improved the 100K-column matrix by about 1.29x, not 6–7x. These are isolated selector timings, not end-to-end serving gains. In particular, 100K actual context also includes compressed 50K-column selections where DeepSelect was slower.

Synthetic controls

Fresh single-GB200 kernel measurements show that both shape and score distribution matter. With normal-random FP32 logits, top-k 512, 4K query rows, causal row bounds, uncompressed KV, preallocated outputs, CUPTI CUDA-graph timing and cold L2:

KV / logits columns	Existing prefill top-k	DeepSelect	Native / DeepSelect
8,192	80.9 us	135.0 us	0.60x
51,200	321.4 us	352.4 us	0.91x
102,400	688.3 us	499.9 us	1.38x

A controlled distribution ablation at the same 8K-query/8K-KV shape demonstrates why random logits cannot stand in for every indexer workload:

FP32 score distribution	Existing prefill top-k	DeepSelect	Native / DeepSelect
Normal random	122.4 us	225.1 us	0.54x
Same random values + 100	1168.0 us	241.8 us	4.83x
Same random values * 1e-8	1153.6 us	224.3 us	5.14x

All cases passed selected-score correctness checks. The native selector initially bins scores through FP16 conversion; concentrated values increase histogram contention and can require further selection work. DeepSelect is much less affected in this ablation. This demonstrates input-dependent performance, but does not by itself establish the cause of the exact 6.3x historical model-profile result.

For the tested DSV4.1 checkpoint, 102,400 actual context tokens produce 51,200-entry KV in layers 2–19 and 102,400-entry KV in layers 20–39. With the NEW sparse indexer, dense prefill top-k runs at index sources 2/8/14 over the compressed context and source 20 over the uncompressed context; later indexers use selected candidate blocks. A single nominal KV-length threshold does not describe all of these selections.

No universal 6–7x speedup or real-weight end-to-end gain is claimed by this PR. The captured real-model measurements argue against using DeepSelect unconditionally for prefill. Any automatic dispatch should depend on the actual logits width and be validated with corrected model-level measurements; the random-logit sweep alone is insufficient to establish a production cutoff.

Real-weight evaluation previously completed with this patch in the combined optimized model configuration: deepseek-ai/DeepSeek-V4.1-Flash, revision dba1be0a40aa45a94ad051997016db3960a90277, TP4/EP4, native MegaMoE/shared fusion, FlashInfer sparse attention. Full GSM8K, lm-eval 0.4.12, five-shot multi-turn chat, thinking disabled, greedy generation, 1024 output tokens, seed 123:

Configuration	Correct / 1319	Accuracy
Auto-MoE/FlashMLA FP8-indexer baseline, three runs	1270–1274	96.2851–96.5883%
Combined optimized configuration, FP8 indexer	1269	96.2092%
Combined optimized configuration, MXFP4 sparse indexer	1266	95.9818%

These combined evaluations include separate attention and MoE changes; they do not isolate DeepSelect's quality effect. Selection tests check score correctness, not identical membership among tied scores.

AI assistance

Codex assisted with implementation, validation, and this description. Woosuk Kwon requested upstreaming this change after reviewing its scope and reported results.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[WoosukKwon]

/ci run

[github-actions]

❌ This PR is 5 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

[WoosukKwon]

/ci run

[github-actions]

✅ Triggered Buildkite CI #89392 for commit a0f8ffee5fed.