[ischencheng]

Overview

DSpark walks its draft block left to right: each step adds W2 @ W1[prev] to the base logits (W2 is V x 256) and then takes the argmax or a sample. Today every draft token is a full-vocab GEMV, an add and a sampler, and W2 (78 MB in bf16 at V = 151936) is read from HBM again on every step.

This adds dspark_markov_quantization="int8" to the speculative config. With it, sm_90 runs the whole walk in one cooperative kernel. W2 is quantized to int8 with one fp32 scale per row, most of it stays in shared memory across all SMs and the rest streams from L2, and the per-step argmax across CTAs goes through one atomic {key, count} pair. Off by default.

The kernel is adapted from the MIT-licensed int8 Markov walk kernels in https://github.com/tie-pilot-qxw/dspark-markov-walk (notice kept in csrc/dspark_markov_walk.cu).

Claims

• Markov walk on 1x H100, eager loop vs this kernel: 298 -> 60 us at V = 151936 with 7 draft tokens (bs 1, greedy), 431 -> 68 us at bs 16, 563 -> 112 us at bs 64. 3.4-6.3x across the table below.
• Qwen3-4B on 1x H100, output tok/s over the eager loop: +4.3% at c1 and +4.6% at c16 with the 7-token DSpark drafter (greedy), +9.3% and +8.8% with the 16-token one. Probabilistic at c1: +5.5% and +11.7%.
• Greedy outputs at c1 match the eager walk (80/80). Acceptance length moves by at most 0.06 tokens and GSM8K by at most 0.025 on these prompts.

Validation

All on 1x H100 at 5af18a7. The one later commit only changes how the launch sets its shared memory attribute, and the kernel tests pass on it too.

Kernel (python benchmarks/kernels/benchmark_dspark_markov_walk.py, CUPTI, CUDA graph, cold L2, median us per call for all draft steps). The eager arm is _sample_sequential's loop, with greedy_argmax and gumbel_sample:

V	draft tokens	bs	greedy eager -> int8	probabilistic eager -> int8
151936	7	1	298.2 -> 60.5 (4.93x)	356.2 -> 80.9 (4.40x)
151936	7	16	431.0 -> 68.0 (6.34x)	563.1 -> 91.4 (6.16x)
151936	7	64	563.2 -> 112.3 (5.01x)	878.6 -> 179.8 (4.89x)
151936	16	1	682.7 -> 124.9 (5.47x)	819.0 -> 172.4 (4.75x)
129280	5	1	198.6 -> 41.8 (4.75x)	237.5 -> 54.8 (4.33x)
129280	5	64	341.6 -> 77.1 (4.43x)	535.4 -> 114.3 (4.69x)
248320	16	1	959.2 -> 270.3 (3.55x)	1099.9 -> 324.6 (3.39x)

e2e, Qwen3-4B, GSM8K prompts (80 at c1, 320 at c16, max 512 tokens), eager and int8 arms interleaved, 2 runs each in a fresh process (runs of an arm within 0.7%), mean output tok/s:

drafter	setting	eager	int8	accept len eager / int8	GSM8K eager / int8
deepseek-ai/dspark_qwen3_4b_block7 (7)	greedy c1	1031	1076 (+4.3%)	6.255 / 6.253	0.850 / 0.850
	greedy c16	13047	13651 (+4.6%)	6.351 / 6.369	0.834 / 0.816
	probabilistic c1	1012	1069 (+5.5%)	6.214 / 6.246	0.813 / 0.838
shreybirmiwal/Qwen3-4B-DSpark-b16 (16)	greedy c1	1057	1155 (+9.3%)	7.046 / 7.041	0.850 / 0.850
	greedy c16	12046	13111 (+8.8%)	7.226 / 7.239	0.822 / 0.828
	probabilistic c1	1015	1133 (+11.7%)	6.918 / 6.975	0.800 / 0.813

Tests:

• pytest tests/kernels/test_dspark_markov_walk.py: 66 passed.
• pytest "tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py::test_dspark_correctness_and_acceptance_rate[qwen3-deepspec-int8-walk]" with VLLM_GPU_SYNC_CHECK=error: passed (GSM8K 0.791, acceptance length 4.009; the stock qwen3-deepspec case in the same run: 0.778, 4.011).

Details

• Only acceptance can change. int8 W2 can move near-ties, and the target still verifies every draft token. For probabilistic drafting the kernel writes the bf16 logits it sampled from into draft_logits, so rejection sampling uses exactly that distribution.
• Sampling noise is Philox keyed by the request seed and the draft position, so seeded requests stay reproducible. A seed draws different tokens than it does with the eager loop.
• TP: every rank runs the whole walk, so with this option the draft logits cache stays replicated and the sequential sampling is not batch-sharded. I only ran it at TP1.
• Supported: sm_90, vanilla Markov head with rank 256 in bf16, fp16 or fp32, full-vocab drafts, up to 16 draft tokens, V up to 270336 on 132 SMs, bs <= 64 (larger batches keep the eager loop). Loading fails with a clear error for draft watermarking, the online acceptance estimator, a logit soft cap or output logit transform, a draft model with its own sample_sequential, dspark_draft_topk, or a quantized W2.
• Memory: 114 MiB of int8 layouts at V = 151936.
• CI: the kernel test runs on H100. The new qwen3-deepspec-int8-walk case skips on GPUs with too few SMs to hold W2 (e.g. a MIG slice), so I ran it on a full H100 (above).
• Not a duplicate: [Perf][DSv4.1] DSpark draft on the specialized decode #61869 fuses the Markov loop for DeepSeek-V4.1's specialized decode on GB300. [Perf][DSpark][DSv4.1] Fuse the Markov sampling loop on NVIDIA by inverse CDF #64442 fuses it for DeepSeek-V4.1 on NVIDIA by inverse CDF, with a bias GEMM per step. [Spec Decode][DSpark] Markov sampling in the Gumbel kernel, wider DFlash padding, top-p fix #64476 adds the bias inside the Gumbel kernel, also with a GEMV per step, so if it lands first the probabilistic eager column needs a re-measure. [Perf][Spec Decode] Vocab-parallel DSpark Markov sampling and vocab-sharded verification over TP #65007 shards the walk over TP, [Perf][Spec Decode] DSpark: run the top-k Markov walk in one kernel #64033 is the top-k walk and [Perf][DSpark] lossless improve DSpark throughput by W2 matrix topk pruned #56694 prunes W2 with top-k. This PR is the generic DSpark path on Hopper with W2 kept on chip.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[mergify]

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

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.