[Spec Decode] Add top-k DSpark Markov projection

[askliar]

Purpose

DSpark computes the base logits for all draft positions in parallel, then applies the Markov bias sequentially over the full draft vocabulary. The repeated full-vocabulary Markov projection is on the serial drafting path.

This PR adds dspark_draft_topk for Qwen3 DSpark. It selects the top-k candidates from the base logits once, gathers the corresponding markov_w2 rows, and evaluates the sequential Markov bias only for those candidates.

The current implementation deliberately reuses the existing dense sampling path:

• Compute dense base logits and select the top-k candidates once for all draft positions.
• Initialize the reusable dense logits buffer to -inf.
• At each sequential step, apply the Markov bias only to the selected candidates and scatter the corrected values into that dense row.
• Use the existing greedy or probabilistic Gumbel/rejection sampling path unchanged.

The base LM-head output and probabilistic draft-logits storage remain dense. This optimization specifically reduces the repeated sequential Markov projection while avoiding custom proposal and rejection-sampling plumbing.

Current scope:

• Qwen3 DSpark

Test Plan

• Compare dense DSpark against dspark_draft_topk across lower- and higher-concurrency workloads.
• Check GSM8K accuracy and acceptance length for greedy and probabilistic drafting.
• Profile the decode window to confirm that savings occur in the drafter projection.
• Run focused configuration/top-k tests and repository pre-commit hooks.

Results

End-to-end comparison

Tested with draft_sample_method: probabilistic, temperature 1, and 7 speculative tokens on GB10.

Implementation	Top-k	GSM8K strict / flex	AIPerf tok/s/user¹	AIPerf total tok/s¹
Initial PR implementation	512	0.82 / 0.83	65.3	44.6
Initial PR implementation	Off	0.80 / 0.83	53.0	40.6
Current simplified implementation	512	0.84 / 0.84	78.1	52.1
Current simplified implementation	Off	0.84 / 0.85	54.0	41.3

¹ Concurrency 64, ISL 32768 / OSL 1024.

• At concurrency 64, top-k 512 improves per-user throughput by 23% for the initial implementation (53.0 → 65.3) and 45% for the current simplified implementation (54.0 → 78.1).
• With top-k 512, the current implementation is 20% faster per user than the initial implementation (78.1 versus 65.3). Their no-top-k baselines are equivalent.
• These results confirm that top-k benefits higher-concurrency, GPU-bound decode.
• GSM8K scores remain within the observed run-to-run range; there is no measured accuracy regression.

Decode GPU profile

GPU time by selected kernel category, in milliseconds, over a concurrency-1 / approximately 8K-context decode window:

Kernel category	Initial / Off	Initial / 512	Current / Off	Current / 512
Drafter projection (cutlass / GEMV)	155	105	154	107
Elementwise / norm	28.7	34.4	28.3	28.9
Sort / top-k / select	5.5	6.2	5.5	6.7
Total GPU time	878.5	822.9	870.8	831.4

The profiles confirm that the saving occurs where expected:

• Top-k 512 reduces the drafter projection by approximately 47–50 ms. The top-k selection itself adds only approximately 1 ms.
• Total decode GPU time falls by 6.3% for the initial implementation and 4.5% for the current implementation.
• The initial path adds 5.7 ms of elementwise/index work, while the current implementation adds only 0.6 ms. The profiler reports 35 additional distinct kernels for the initial path versus 14 for the current path.
• Major target-model categories such as MoE routing/GEMMs, dense matmuls, and attention are unchanged within run-to-run variation.

The concurrency-1 profile establishes the per-step compute saving. At concurrency 64 with a 32K-token input, decode becomes GPU-bound and that saving converts into substantial throughput. The current path's advantage over the initial path at high concurrency is consistent with its lower auxiliary kernel/index overhead, although that contribution was not isolated independently.

Acceptance-length sweep

Top-k	t=0 (greedy)	t=1 (probabilistic)
512	3.9509	3.6582
1024	3.9448	3.6573
2048	3.9445	3.6472
4096	3.9403	3.6455
8192	3.9403	3.6556
16384	3.9648	3.6473
32768	3.9253	3.6568
65536	3.9403	3.6429
Off	3.9253	3.6413

Focused checks

• pre-commit run --files <changed files>: passed, including Ruff and mypy.
• pytest -q tests/v1/spec_decode/test_dspark_topk.py: 2 passed.
• pytest -q tests/test_config.py -k 'draft_sample_method': 2 passed.
• pytest -q tests/v1/worker/test_gpu_rejection_sampler_chunking.py -k 'preserves_request_boundaries': 1 passed.

Benchmark caveats

• AIPerf used 5 measured requests plus 2 warmup requests, so exact throughput values have some run-to-run variance. The top-k 512 differences are substantially larger than the observed noise.
• AIPerf and GSM8K runs used gpu-memory-utilization=0.8.
• Nsight Systems runs used gpu-memory-utilization=0.7 because CUPTI trace buffers caused model-load OOM at 0.8 on this node. The decode kernels were unchanged, so the profiles remain representative of the code paths.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR is described.
• The test plan is included.
• Test results are included.
• Documentation changes are not required for the current model-specific scope.

[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.

[mergify]

Documentation preview: https://vllm--49969.org.readthedocs.build/en/49969/

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @askliar.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[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.

[benchislett]

High level: I feel like this approach may be overly optimized.

I'm not convinced that sampling from a smaller distribution will bring efficiency here. If the savings are just coming from the sparse markov projections, then we could probably gain a lot just by having the bias module directly overwrite the logits of the backbone (with within-topk being written as calculated and outside-topk set to -inf or 0). That would eliminate the need for custom sampling code

[benchislett]

Is it possible to use addmm here?
See #50737

If it requires a layout change or something heavy, don't bother. But if it happens to be a drop-in, seems like a nice potential improvement

[askliar]

Replaced with baddbmm - that should be the easiest drop-in replacement.

[benchislett]

This constraint is no longer required. Please rebase to latest main, which includes a change to replicate the markov_w2 on all gpus (TP disabled)

[askliar]

Fixed, thank you.

[benchislett]

LGTM