[MrGeva]

Overview

With tensor parallelism, DSpark speculative decoding repeats full-vocabulary work on every TP rank, twice per decode step: the drafter's sequential Markov sampling loop, and the target's verification (an all-gather of the target logits, then temperature, top-p and the rejection sampler over the whole vocabulary). This PR splits both by vocabulary across the TP ranks. It has two parts, each behind its own opt-in flag (both off by default):

1. Vocab-parallel DSpark Markov sampling (speculative_config.dspark_tp_markov_sampling, experimental): each rank runs the Markov loop on its share of the vocabulary and the ranks exchange per-block maxima.
2. Vocab-sharded verification (speculative_config.vocab_sharded_verification): each rank verifies on its own vocabulary blocks, without gathering the target logits, and with part 1 on, the drafter also works only on those blocks.

They ship together because they share the exchange machinery: symmetric-memory tagged words with never-zero tags that the reader clears, a bounded spin with a timeout flag in pinned host memory, and a setup fallback that all ranks agree on. Also, the draft side of part 2 is part 1's Markov loop restricted to the ranks' owned columns.

Part 1: vocab-parallel DSpark Markov sampling

DSpark drafts a block of num_speculative_tokens tokens with a parallel backbone pass, then runs a sequential Markov loop. Each of the loop's steps adds a Markov bias, from the previously sampled token, to the base draft logits and Gumbel-samples the next token. With tensor parallelism the loop runs replicated on every TP rank. Each step:

• reads the whole markov_w2 projection ([draft_vocab, markov_rank] bf16, 102 MB for a 200k vocabulary) to compute a few rows of bias;
• Gumbel-samples the full vocabulary.

Both are memory bound, and every rank does the full work. For MiniMax-M3.1 Flash (K=7) the loop is 150–570 µs per decode step at concurrency 1–32 on GB300, on the draft's critical path.

With dspark_tp_markov_sampling, each TP rank owns a contiguous range of the sampler's 1024-column blocks and, per step:

1. computes the bias of its columns only, from a row slice of markov_w2 (1/TP of the weight traffic);
2. Gumbel-samples its blocks, with the noise keys the replicated loop uses (gumbel_noised_argmax, spec_step_seed);
3. publishes one 8-byte word per block, [block max (fp32) | 22-bit tag | 10-bit index in the block], to every rank through torch symmetric memory. Each word carries its own tag, so no fence or flag is needed (NCCL's LL protocol). Tags are never 0, and the reader clears each word it consumes, so a slot holds either nothing or the current step's word and a stale word can never be taken for a fresh one. A sys-scope release/acquire flag variant and an NCCL all-gather were both slower than the replicated loop.
4. Every rank picks the winner as the replicated loop does: the first NaN block, otherwise the first maximum.

The block maxima are the values the replicated loop computes, so the drafted tokens are the same.

The rejection sampler reads the draft-logits cache at full vocabulary on every rank (unless part 2 is on, see below). Each rank writes its own columns while sampling, and the rest is completed in one of two ways:

• ≤ 4 requests: while sampling, ranks also push their biased logits to their peers' staging buffers, two bf16 per tagged 8-byte word, and one kernel unpacks them after the loop.
• More requests: every rank recomputes the peers' columns for all steps at once, from the sampled tokens, projecting only the peers' rows of markov_w2 (left and right of its own). This reads them once per draft step instead of once per Markov step. Pushing a whole cache would cost more NVLink time than it saves.

Exactness guard. The sliced and batched projections must equal the replicated loop's projection bitwise. cuBLAS gives that for 64-aligned vocabularies, but not for odd ones (e.g. 50017), where it picks different kernels. So at load time the speculator checks it on the real weight against the real model.markov_bias (logits processor and GEMM dispatch included), for every batch size up to 128, and all TP ranks agree on the result (0.42 s and a 367 MiB transient peak per rank for M3.1 Flash at TP4). If any size differs, or the configuration is unsupported, the replicated loop runs. Unsupported configurations are:

• greedy drafting, top-k drafting, a reduced draft vocabulary, batch-sharded draft logits;
• watermarking, adaptive verification, fp64 Gumbel, a model-specific sample_sequential;
• a scaled or soft-capped logits processor, a quantized or skinny-GEMM markov_w2, TP = 1, non-CUDA.

Batches above 128 requests also run the replicated loop.

Part 2: vocab-sharded verification

With a vocab-parallel LM head, verification all-gathers the target logits, so every TP rank holds [num_logits, vocab] and applies temperature, top-p and the rejection sampler over the whole vocabulary: the same work on every rank. With probabilistic drafting the drafter does the same with its own logits: an all-gather per draft step, and a full-vocabulary draft-logits cache on every rank.

With vocab_sharded_verification, each TP rank owns whole 8192-column blocks of the vocabulary (the rejection sampler's statistics blocks, split as evenly as possible) and works on those only:

• Target logits: the rank's own lm_head shard, plus the few columns of its owned range that sit in the right neighbour's shard, which that neighbour pushes. Nothing is all-gathered.
• Temperature and top-p: on the owned columns. Top-p is an exact select run jointly by all ranks: the smallest prefix (by logit, then token id) whose probability mass reaches p, plus the row max. Masses are summed in 64-bit fixed point, so every rank, and every run, gets the same kept set. Greedy rows do not trigger it.
• Rejection sampling: the rejection sampler's per-block kernels run on the owned blocks only (they now take a block offset and a column offset); the per-block partials are exchanged and the per-request kernels combine all blocks in the stock order. The drafted tokens' target and draft logits come from their owners. Given the same processed logits, the result is bitwise the stock one.
• Drafting (DSpark with dspark_tp_markov_sampling): the draft logits are computed the same way (own shard + right neighbour's columns), part 1's Markov loop runs on the same owned ranges, and only the owned columns of the draft-logits cache are written, since verification reads only those.

All exchanges are small (per-block partials, per-row histograms, the boundary columns) and use part 1's tagged-word scheme. Each verification step uses its own tag epoch and alternates between two buffer halves. The sharded verification kernels are compiled at load, collectively on every rank, so that no rank compiles while its peers wait on an exchange. The warmup runs one step per kernel variant that serving can launch: for each top-p tie-split count, a row count that is a multiple of 16 and one that is not (warmup_num_logits, commit 946ffe5). Verified: no vocab-shard Triton compiles after load in a c32 serving smoke, and the unit tests pass.

Steps with requests that need the full vocabulary (logprobs, top-k, min_p, logits processors or a thinking budget, structured output) take the stock path. Before such a step, the cached draft logits of its requests are completed with an all-gather, bounded to 64 MiB per chunk.

Opt-in, flags and fallbacks

Both options are off by default and change nothing unless set.

flag	falls back (logged on every rank alike) to	when
dspark_tp_markov_sampling	the replicated Markov loop	the exactness guard fails; a configuration listed under part 1; batches above 128 requests; symmetric memory unavailable
vocab_sharded_verification	gathered verification, as today	no probabilistic drafting; watermarking, FLy, adaptive verification, batch-sharded sampling, prefill context parallelism, a custom rejection sampler, VLLM_COMPUTE_NANS_IN_LOGITS, sampling masks or trace replay; a target model without compute_logits_local; a draft-logits cache narrower than the target vocabulary; symmetric memory unavailable
vocab_sharded_verification, per step	gathered verification for that step	a request in the batch needs the full vocabulary (above)
both, draft side	part 1's loop on full-vocabulary columns	the draft model or vocabulary does not support owned columns

Failure handling (both parts):

• If symmetric memory cannot be allocated or rendezvoused on any rank (no P2P, multi-node TP, a torch build without it), every rank logs the reason and falls back. Each setup step is agreed across ranks first, so no rank waits in a collective for a peer that failed.
• An exchange that waits more than about 2 s for a peer stops spinning and sets a flag in pinned host memory. The next step reads the flag without a sync and raises, instead of producing tokens from incomplete data.

Claims

Part 1 (vocab-parallel Markov sampling)

• Same tokens and the same draft-logits cache, bitwise. On MiniMax-M3.1 Flash's real draft model (TP4, GB300), the replicated loop and this loop were run from identical draft metadata. The metadata includes a padding request, a greedy request and a temperature-0.7 request. The drafted tokens and the entire draft-logits cache match on all 4 ranks at B = 1, 2, 4, 8, 16, 32, 64, 96 and 128, in both completion modes, with 0 exchange timeouts.

• Loop time (CUDA-graph timed on the real draft model, TP4, µs, best of two runs on one node; the shared draft lm_head + logits all-gather is included in both):

  requests	replicated	vocab-parallel
  1	291	233
  4	371	315
  8	440	389
  16	563	509
  32	705	600
  64	1156	995
  96	1575	1372
  128	2010	1759

  The review changes do not cost loop time. Against the first revision on the same node, run ABAB: B1 232 → 233, B32 601 → 600, B64 998 → 995, B128 1759 → 1759.

• Draft CUDA graph per decode step in serving: −61 µs at concurrency 1, −35 to −42 µs at 8, −128 to −133 µs at 32 (Validation).

• Memory: about 43 MiB of symmetric memory per rank, allocated and zeroed once at load: push staging for 4 requests, 2 · TP · 4 · K · ⌈blocks/TP⌉ · 512 · 8 B (42.9 MiB at TP2/TP4, 43.8 MiB at TP8 for V = 200064, K = 7), plus 0.4 MiB of block words for 128 requests.

Part 2 (vocab-sharded verification)

• Same tokens: for the same processed logits, the sampled tokens and counts equal the full-vocabulary rejection sampler's (block, standard and synthetic modes). Drafting on owned columns draws the same tokens as part 1's loop.

• Exact, deterministic top-p: the kept set is the float64-exact one except within fp32 rounding of p at the boundary (on real logits: 4 of 3611 rows, one tied token each); never a superset.

• Faster from 8 requests up (MiniMax-M3.1 Flash, TP4, DSpark K=7, GB300; base = the same build with part 1 on and gathered verification). On the final revision, step period −6.3% at concurrency 16 (15.62 → 14.63 ms) and −6.7% at concurrency 32 (18.82 → 17.56 ms). Time per decode step outside the target model's CUDA graph (step period − target graph span; the target graph itself is unchanged):

  concurrency	base	this PR	change
  1	964 µs	985 µs	+21 µs
  8	1523 µs	1243 µs	−280 µs
  16	2297 µs	1398 µs	−900 µs
  32	2955 µs	1695 µs	−1260 µs

  c16 / c32: final revision, two runs per arm. c1 / c8: first revision (the review changes kept the c16 / c32 gain within noise of the first revision's −882 / −1210 µs).

• At concurrency 1 the exchanges' latency outweighs the work saved (verification 325 → 377 µs, drafting 591 → 565 µs).

• Memory: one symmetric buffer per rank, allocated at load (before the KV cache is sized), proportional to max_num_seqs: 253 MiB at TP4, K = 7, V = 200064, max_num_seqs 256, plus 50 MiB on the draft side (sized for the Markov sampler's 128-request limit).

Accuracy: GSM8K unchanged within noise with both options on (Validation).

Validation

GB300 (SM103), driver 580.178.04, image vllm/vllm-openai:minimax-m31-flash-arm64-cu134@sha256:5d763476… (vLLM 36ae880) with this PR's files, adapted to that image's API where main differs (the image predates DraftLogitsCache). Serving uses the InferenceX GB300 recipe: TP4, DSpark K=7 probabilistic drafting, block verification, NVFP4 KV, FULL_AND_PIECEWISE.

Both options, final head

Stock image against the image with only this PR's files at the head commit, both options on, on one node, arms back to back. GSM8K 8-shot through the chat API, 320 questions at concurrency 32, block verification:

	greedy	temperature 1, top-p 0.95	mean acceptance length
stock image	0.9812	0.9656	4.65
this PR, both options on	0.9750	0.9719	4.72

Both differences are 2 questions out of 320, in opposite directions, within one standard error (about 0.008).

Serving smoke (same node, both options on, 64 greedy + 64 sampled GSM8K requests at concurrency 32, with temporary per-step counters): all 4 ranks logged Vocab-sharded verification: this rank owns vocabulary columns [...] and DSpark drafting is vocab-sharded.. Every serving step with drafts was verified vocab-sharded, and every draft step ran the owned-columns Markov loop, on all 4 ranks. There were no fallback warnings, full-vocabulary verification steps or exchange timeouts. The only full-vocabulary steps were three warm-up dummy runs at load, before CUDA graph capture.

Tests

• tests/v1/spec_decode/test_dspark_tp_markov.py, run in CI by the new 2×B200 step :nvidia: (B200) V1 Spec Decode DSpark TP Markov:
  • Objects. A real DSparkMarkovHead (ParallelLMHead), LogitsProcessor, DraftLogitsCache and InputBuffers. Both arms run DSparkSpeculator._sample_sequential: the replicated _sample_markov_chain as reference, and the vocab-parallel early return.
  • Equality. Tokens and the whole cache are checked bitwise at 1, 2, 4, 5, 33 and 128 requests, in both completion modes and with CUDA-graph replays. Rows include a padding, a greedy and a T=0.6 request; TP2 and TP4; vocabularies 200064 and 151936.
  • Owned columns (part 2's draft side). Drafted tokens and the owned columns of the cache bitwise equal to the replicated loop; after complete_cache, the whole cache is equal, at 1, 5 and 128 requests (several chunks).
  • Fallback and failure paths. Every unsupported_reason; symmetric memory failing on one rank (every rank falls back); a late peer at epochs where a 14-bit tag would wrap; an exchange timeout, which check_exchange raises.
  • The test skips with the reason if the projections are inexact on the GPU, since the feature then falls back by design.
• tests/v1/spec_decode/test_vocab_sharded_verify.py, run in CI by the new 2×B200 step :nvidia: (B200) V1 Spec Decode Vocab-Sharded Verification:
  • Owned columns (TP 2 and 4, vocab 200064 and 151936, bf16 and fp32): equal to the gathered logits' columns.
  • Top-p (rows 1/8/64/200; LM-like bf16, quarter-step ties, 30% −inf, peaked and flat rows; p in [0.01, 0.99], p = 0.9999, 1 − 1e-7 and 1; rows without a finite logit): the kept set equals a float64 reference, except where the reference's cumulative mass at the boundary is within 1e-6 of p. Two runs give identical masks.
  • Rejection sampling (block, standard and synthetic modes; greedy and sampled requests; −1 placeholder drafts; 1/3/32 requests): sampled tokens and counts bitwise equal to the full-vocabulary rejection sampler.
  • Exchange protocol: a late peer at tag-wrap epochs, an odd shard size, and a timeout.
• tests/v1/worker/test_gpu_sampler_flags.py: one case per request feature for Sampler.needs_full_vocab (min_p, top-k, logprobs, thinking budget, logit bias, penalties and bad words set it; greedy-with-top-p, temperature and top-p alone do not).

On GB300 (image port of the tests, TP2 and TP4 × both vocabularies): 8 passed for the two multi-GPU files, 26 passed for the sampler flags. The TP4 cases need 4 GPUs and do not run in CI.

Not covered by CI: an LLM-level run needs a DSpark checkpoint; the smallest are MiniMax-M3.1 Flash and DSV4. The real-model checks and serving runs here are on M3.1 Flash on GB300.

Part 1 serving

Agentic-like code prompts of about 22k tokens; nsys, same node, arms back to back; draft CUDA graph span p50 per step, µs:

concurrency (requests in the draft)	draft graph, replicated	draft graph, vocab-parallel (two runs)	saved per step	share of the step
1	646	586 / 585	61	0.9%
8	917	875 / 882	35–42	0.3%
32	1368	1235 / 1240	128–133	0.7%

Arms on one node per concurrency (replicated, vocab-parallel, first revision, vocab-parallel). The first revision measured 582, 869 and 1245 µs, so the review changes keep the gain within run-to-run noise. An earlier c16 run of the first revision measured 1079 → 1003 µs. Every vocab-parallel run logged DSpark Markov sampling is vocab-parallel over TP. on all 4 ranks, with no fallback, timeout or error. The target graph does not change; whole-step periods move by up to ±4% between server restarts on this setup, so the draft graph span is the reliable comparison.

Part 2 serving

Agentic-like code prompts of about 22k tokens, temperature 1, top-p 0.95, part 1 on in both arms; nsys on one node, arms alternated base / PR / base / PR.

• Final revision: time outside the target graph c16 2300 / 2294 → 1403 / 1392 µs (draft graph 972 → 838 µs; period 15.62 → 14.63 ms, −6.3%); c32 2964 / 2945 → 1704 / 1685 µs (draft graph 1180 → 956 µs; period 18.82 → 17.56 ms, −6.7%). Engagement lines on all 4 ranks.

• First revision, per part (averages of two runs):

  concurrency	verification (target graph end → draft graph start)	draft graph	eager logits path (NCCL gathers + copies)
  1	325 → 377 µs	591 → 565 µs	25 → 0–2 µs
  8	626 → 438 µs	846 → 743 µs	120–140 → 0 µs
  16	1274 → 522 µs	953 → 831 µs	–
  32	1674 (one run; other 2218) → 668 µs	1187 → 964 µs	–

• Against batch-sharded sampling (enable_batch_sharded_sampling, same build without part 1, which it excludes), time outside the target graph: 1422 → 1314 µs at concurrency 8, 1631 → 1416 µs at 16, 2370 → 1695 µs at 32.

Top-p on real logits: 3611 verify rows dumped from MiniMax-M3.1 Flash DSpark serving (64-, 105-, 186- and 256-row batches, p = 0.95). 4 rows differ from the float64 reference, each by one tied token at the boundary (reference cumulative mass within 1.5e-8 of p); no row keeps extra tokens. CUDA-graph timed, per call on 4 ranks: 104 µs at 64 rows, 211 µs at 256 rows.

Earlier accuracy runs (GSM8K chat, 320 questions, concurrency 32, greedy / temperature 1 top-p 0.95): part 1 with MiniMax-M3 model changes in both arms: 0.9719 → 0.9812 / 0.9719 → 0.9750 and 0.9781 → 0.9750 / 0.9750 → 0.9719. Part 2 against part 1 alone: 0.9750 → 0.9781 / 0.9688 → 0.9750 (first revision) and 0.9781 → 0.9812 / 0.9750 → 0.9688 (review revision). Mean acceptance length unchanged within run-to-run noise.

Details

• Integration with main (8530a77):
  • dspark_tp_markov_sampling is a method-specific option: it is listed in _METHOD_SPECIFIC_OPTIONS and in DSparkSpeculatorConfig.supported_options. Its type is bool | None (unset means off), so the generic option check rejects it for other methods only when it is set.
  • vocab_sharded_verification is a plain SpeculativeConfig field; it is a property of verification, not of a drafting method.
  • DSparkSpeculator.load_model calls _try_enable_tp_markov_sampling when only part 1 is on. With part 2 also on, the runner decides verification first and then calls enable_vocab_sharded_drafting(shard_or_None), which builds exactly one Markov sampler (owned-columns or plain).
• Commits:
  • f8d461a: part 1;
  • eddf09a, db489e9, f2ffea6: part 1 review changes (agreed symmetric-memory fallback, surfaced timeouts, never-zero 22-bit tags with consume-and-clear, the guard against model.markov_bias, the skinny-GEMM exclusion, peers-only recompute; tests on the real sampling path and the CI step; the timeout flag in pinned host memory);
  • 0ae150c: part 2;
  • 41a2302, 54030bb, 9741123: part 2 review changes (agreed fallback via VocabShard.create, load-time eligibility and collective warmup, do_not_specialize on per-step scalars, never-zero tags with consume-and-clear, bounded completion chunks, tighter buffer sizing, CI step and tests; the timeout flag in pinned host memory).
  • 946ffe5: the load-time warmup covers every kernel variant serving can launch (the top-p tie kernels' split count and the exchange strides' divisibility by 16), with a test that every step size is covered.
• Code map (part 2):
  • vllm/v1/worker/gpu/spec_decode/vocab_shard.py (new): ownership, the exchanges, the owned-logits assembly and the sharded top-p.
  • rejection_sampler_utils.py: rejection_sample(..., vocab_shard=...); the stats, residual-mass and resample kernels take block_start / target_col_start, and the logprob helper optionally reads the drafted token's logits from a small array. Stock calls pass 0 / None.
  • RejectionSampler.supports_vocab_sharding / vocab_sharding_unsupported_reason / warmup_vocab_sharded / verify_vocab_sharded; Sampler.needs_full_vocab (per request); GPUModelRunner.sample takes the sharded path when eligible.
  • DSpark: TPMarkovSampler.create(speculator, owned_ranges), DSparkSpeculator.enable_vocab_sharded_drafting / complete_draft_logits.
• Validation scope: the measurements ran on the MiniMax-M3.1 Flash day-0 image with this PR's files, adapted to that image's API where main differs. The A/B serving runs of each part had the same unrelated MiniMax-M3 model changes in both arms; the final GSM8K run above compares the stock image with only this PR's files. The option wiring for main's per-method speculator config has not run end to end on main, which has no MiniMax-M3.1 NVIDIA model path.
• Related: batch-sharded sampling (enable_batch_sharded_sampling) splits the batch across ranks instead of the vocabulary. It is exclusive with both options here: part 1 is ignored when batch-sharded draft logits are in use, part 2 when batch-sharded sampling is on. Part 2's verification side works with any probabilistic drafter whose draft-logits cache is replicated; only its draft side is DSpark-specific.

---

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.

AI assistance: the kernels, tests and benchmarks were developed with Claude Code (Claude Opus 5.5). Every line was reviewed by the human submitter, who also confirmed the validation above.

🤖 Generated with Claude Code

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

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

🚀

[MrGeva]

No fallback when symmetric memory is unavailable. TPMarkovSampler(self) is constructed unguarded from DSparkSpeculator._try_enable_tp_markov_sampling, but the private import here and symm_mem.empty / symm_mem.rendezvous (L348-360) can raise: RuntimeError on PCIe-only, no-P2P or multi-node TP, ImportError on a torch build without it. The test itself catches RuntimeError ("symmetric memory is unavailable"), and SymmMemCommunicator wraps the same calls in try/except RuntimeError and disables itself. Here the exception escapes load_model, so setting dspark_tp_markov_sampling on such a node kills engine start-up. That contradicts the config docstring ("otherwise the replicated loop runs") and the "ignored because ..." warning path. rendezvous is collective, so if only some ranks fail the others stay blocked in it. unsupported_reason also has no check that the TP group is single-node or peer-reachable.

Suggest wrapping the import and construction in try/except, logging the warning and leaving _tp_markov_sampler = None, with the decision agreed across ranks (same cpu_group all_reduce as the projection check).

[MrGeva]

Agreed, thanks. Fixed in eddf09a (test in db489e9).

An exception from the import, symm_mem.empty or rendezvous escaped load_model, and since rendezvous is collective, a failure on one rank could leave the others waiting.

TPMarkovSampler.create() now sets up the exchange in two steps, and every TP rank agrees on each one over the CPU group, as the projection check already did:

1. Import and symm_mem.empty run in a try. An all_reduce(MAX) of the failure flag follows, so no rank enters rendezvous unless every rank has allocated.
2. rendezvous runs in a try, followed by a second agreement.

On any failure every rank returns (None, reason). The speculator then logs "dspark_tp_markov_sampling is ignored because symmetric memory is unavailable (...)" and keeps the replicated loop. The test makes symm_mem.empty raise on rank 1 only, then checks that both ranks fall back without hanging. A multi-node or non-P2P group falls back through the same path (failure in empty/rendezvous), so I did not add a separate node check.

One case is still not covered: a rank that fails inside rendezvous before it publishes its handle to the store. The other ranks would then wait in torch's store exchange, which we can't interrupt from here. A peer's failure while importing handles happens after the exchange, so that case is caught.

[MrGeva]

Exchange timeouts are silently swallowed. After SPIN_LIMIT iterations the kernel only does atomic_add(err_ptr, 1) and carries on with whatever is in w, i.e. stale or zero words. The sampled token is then wrong, and in push mode (L275) stale or zero bf16 values are written into the draft-logits cache. errors is never read on the host outside the test. In production a late or missing peer therefore produces drafts whose cached logits no longer match the drawn tokens, which biases the rejection-sampling probability ratio, with no log line.

Also, _SPIN_LIMIT = 1 << 24 is described as "about a second" (L72). Each iteration is a volatile global load plus a block reduction (roughly 0.5-1 us), so it is closer to 10 s per kernel, times 7 combine kernels per draft step if a peer is gone. That is effectively a hang with a delayed, silent corruption instead of an error.

Suggest checking errors on the host (e.g. once per propose or every N steps) and raising or disabling the feature, and fixing the comment.

[MrGeva]

Agreed, the timeout was silent: it was only counted in errors and never read. Fixed in eddf09a and f2ffea6.

On the spin limit: it was neither about 1 s nor about 10 s. I measured it on GB300: 0.24 µs per iteration (one 196-block row), so 1 << 24 is about 4.0 s per kernel.

• _SPIN_LIMIT = 1 << 23, about 2 s per kernel, and the comment now gives the measured cost.
• On a timeout the kernels set a flag directly in pinned host memory (a UvaBuffer), so the normal path has no extra work.
• DSparkSpeculator.propose() calls sampler.check_exchange(). It reads that flag without synchronizing and raises RuntimeError("DSpark vocab-parallel Markov sampling timed out waiting for a TP peer; ..."). The error therefore surfaces at most one step later, instead of producing silent bad drafts.
• An intermediate version copied a device counter to the host after every draft step instead. That memcpy node cost about 5 µs per step in the draft graph, so I dropped it.
• The test checks this: only rank 0 samples (spin limit patched down), and check_exchange() raises.

I chose raising over disabling the feature at runtime. The sampler's kernels are captured in the draft CUDA graphs, so disabling it would mean recapturing them. Also, a peer that misses an exchange by seconds in a lockstep TP group means the group itself is unhealthy: the target's NCCL collectives would stall on it too.

[MrGeva]

14-bit tag aliasing. Freshness is tag == e & 0x3FFF. A (row, block) word is refreshed only while that row is active, so a slot not written for a multiple of 16384 epochs (about 2340 draft steps at K=7, i.e. seconds at serving rates) and then read again has the expected tag and is accepted without waiting.

• The zero-initialised buffer: at epoch 16384 the tag is 0, so any never-written word passes immediately ("Tags never match a zeroed word: the first epoch is 1" holds only for the first 16383 epochs).
• Concurrency drops below row r, stays there for exactly 16384*k epochs (same parity slot), then rises: row r's old words (tag from epoch e-16384k) are read as fresh before the peers' writes land, so a token and max from a past draw are used.

The combine kernel (L192) has the same check. The push path avoids this with the full 32-bit epoch.

The word is full (32 value + 14 tag + 18 token), but the block index is implied by the word's position, so only the 10-bit in-block index needs storing. That frees 8 bits for a 22-bit tag, and removes the 2^18 vocabulary limit.

[MrGeva]

Confirmed, and reproduced on GB300 against the previous head. Fixed in eddf09a (regression test in db489e9).

The reproducer used a fresh, zeroed words buffer with the epoch set so the first tag is 0 (epoch 16383), and delayed the peer by about 1 s with torch.cuda._sleep before its sample. The rank that is on time then accepts the zero words, and the drafted tokens differ from the replicated loop's, with errors == 0:

case	TP2 rank 0 token mismatches	TP4 mismatches per rank
epoch 100, B=4 (control)	0	0
epoch 16383 (tag 0 on zeroed words), B=4 push	2	6–10
epoch 16383, B=8 recompute	10	11–17
epoch 32767 (stale words written 16384 epochs earlier, same tag), B=4	2	6–10

Fixed as you suggested, plus one more step:

• Word layout. The word is now [fp32 block max | 22-bit tag | 10-bit index in the block]. The token is rebuilt from the word's block position, which also removes the 2^18 vocabulary limit.
• Tags are never 0. The tag is epoch % (2^22 - 1) + 1, so a zeroed or cleared word never matches. The epoch is int64.
• Readers clear consumed words. The combine kernel (and the push-unpack kernel) clear each word right after reading it. The next write to a slot comes two epochs later, after the peer has seen this rank's next word, so a slot only ever holds 0 or the word of the epoch being read. A stale word can no longer alias, whatever the tag width.
• Regression test. A late peer at epochs 0, 2^14 − 1 and 2^22 − 2 is waited for, and tokens and cache match the replicated loop.

[MrGeva]

The exactness guard does not test the stock path. inexact_projection_reason compares F.linear(x, w2[c0:c1]) and the batched F.linear with the full F.linear. The replicated loop's bias is model.markov_bias -> LogitsProcessor._apply_head -> markov_w2.quant_method.apply -> dispatch_unquantized_gemm(). On CUDA that can route to torch.ops.vllm.cuda_skinny_unquantized_gemm (layer.use_skinny_gemm, set for ParallelLMHead by maybe_use_skinny_gemm in UnquantizedEmbeddingMethod.process_weights_after_loading) or, with VLLM_BATCH_INVARIANT, to linear_batch_invariant. The latter is excluded in unsupported_reason; the former is not.

sample() calls F.linear directly on cached weight views (here for the recompute path and per step). If the stock projection ever uses a different kernel (a skinny plan for a (vocab, markov_rank) shape, a dispatch change), the guard still reports "exact" while tokens and cache silently diverge from the replicated loop.

Suggest adding getattr(head.markov_w2, "use_skinny_gemm", False) to unsupported_reason, and using the real model.markov_bias(...) output as the guard's reference.

[MrGeva]

Agreed, as a latent risk. There is no divergence today: skinny-GEMM plans exist only for SM121, and none for a (vocab, markov_rank) shape; VLLM_BATCH_INVARIANT was already excluded. Fixed both ways in eddf09a:

• unsupported_reason now rejects markov_w2.use_skinny_gemm ("markov_w2 uses the skinny GEMM").
• The guard's reference is now the real model.markov_bias(...) (DSparkMarkovHead → LogitsProcessor → quant_method.apply → dispatch_unquantized_gemm), at every batch size. A different kernel or dispatch on the stock path now fails the guard, and the replicated loop runs. The guard checks this rank's row slice at M = num_reqs, and the peers' row slices that the recompute path uses at M = 7·num_reqs, row by row against the per-step stock projection.

[MrGeva]

This test never runs in CI. The only step that selects tests/v1/spec_decode as a directory is the single-device "H200 MIG 35GB V1 Spec Decode" step in .buildkite/test_areas/spec_decode.yaml. There device_count() < world_size, so every case hits pytest.skip("Not enough GPUs to run the test."). Multi-GPU spec-decode tests (e.g. v1/e2e/spec_decode/test_sharded_sampling.py) are listed explicitly in distributed.yaml. This PR adds no step and the test has no multi_gpu_test / distributed marker.

tests/AGENTS.md: "A test no step in .buildkite/test_areas/*.yaml selects never runs, which looks like passing." The PR description's "CI runs the test on TP >= 2 CUDA runners only" is not backed by any CI config.

Suggest adding it to a >=2-GPU step and using multi_gpu_test(num_gpus=...) from tests/utils.py instead of the hand-rolled mp.spawn plus skip.

[MrGeva]

Right, the H200 MIG step only skipped it. Fixed in db489e9.

I added a :nvidia: (B200) V1 Spec Decode DSpark TP Markov step (b200-k8s, 2 GPUs) to .buildkite/test_areas/spec_decode.yaml. It has source_file_dependencies on the DSpark speculator directory, gumbel.py, the logits processor, the vocab-parallel embedding, qwen3_dspark.py and the test. The tests use multi_gpu_test(num_gpus=2) / multi_gpu_test(num_gpus=4) instead of the hand-rolled skip.

TP4 does not run in CI (the step has 2 GPUs). I ran it on GB300 (TP2 and TP4, both vocabularies: 4 passed). The PR description is corrected.

[MrGeva]

Stand-ins instead of real objects, and the reference is not the stock loop. tests/AGENTS.md: "Construct vLLM objects with their real constructors; do not stand them in with SimpleNamespace or MagicMock". _fake_speculator builds the speculator, draft model, markov_head, logits_processor and draft_logits from SimpleNamespace. _replicated_loop re-implements the loop as F.linear(markov_embed(prev), w2) plus gumbel_sample.

So the test proves equality with the author's model of the stock loop, not with DSparkSpeculator._sample_markov_chain / model.markov_bias (LogitsProcessor + quant_method.apply). Any difference on the real path (GEMM dispatch, dtype handling, row_mapping) is invisible.

Suggest building a real DSparkMarkovHead / DSparkSpeculator (test_dspark_topk.py already builds a real DSparkMarkovHead) and calling the real _sample_markov_chain as the reference.

[MrGeva]

Agreed. Fixed where feasible in db489e9. The test now uses real objects:

• a real DSparkMarkovHead built with its constructor (ParallelLMHead + UnquantizedEmbeddingMethod.process_weights_after_loading);
• a real LogitsProcessor, DraftLogitsCache and InputBuffers.

Both arms run the real DSparkSpeculator._sample_sequential. With the sampler off, that is the stock _sample_markov_chain → model.markov_bias → _sample_logits → gumbel_sample. With it on, it takes the new early return. The test model wires markov_embed / markov_bias exactly as Qwen3 DSpark does.

What is still not built with its constructor is the DSparkSpeculator itself. Its constructor and load_model need a DSpark checkpoint's config and draft weights, so the test sets the state its sampling reads via __new__, as test_dspark_topk.py does. An LLM-level TP2 test would need a DSpark checkpoint, and the smallest are M3.1 Flash / DSV4, which CI cannot load.

The real-model end-to-end check is out of tree. It used M3.1 Flash's draft model at TP4 on GB300, running the replicated _sample_sequential and the vocab-parallel one from identical draft metadata. Tokens and the whole draft-logits cache were bitwise equal at B = 1…128 on all ranks, in both modes.

[MrGeva]

The enable/fallback glue has no test. TPMarkovSampler.unsupported_reason, _try_enable_tp_markov_sampling, the load_model hook and the _sample_sequential early return are not exercised anywhere. The new test calls only inexact_projection_reason and the sampler on a SimpleNamespace that has none of _draft_topk, draft_watermarker, acceptance_estimator, use_fp64_gumbel, use_confidence_head, _d2t_scatter_index. An attribute typo in unsupported_reason would only surface as an AttributeError at engine load. The PR description itself says the new glue "has not run end to end on main".

At minimum, a unit test over unsupported_reason for each fallback reason, plus one LLM-level run with dspark_tp_markov_sampling on TP=2 comparing against the replicated loop.

[MrGeva]

Agreed. Added in db489e9:

• unsupported_reason. One case per reason: greedy, batch-sharded cache, non-bf16 cache, top-k, reduced vocab, confidence head, fp64 Gumbel, scale, soft cap, head dtype, skinny GEMM. Each sets one attribute and checks the reason, so a typo there fails the test instead of the engine load.
• _try_enable_tp_markov_sampling. Exercised with symmetric memory failing on one rank.
• The _sample_sequential early return. Exercised because both arms go through _sample_sequential, eagerly and under CUDA-graph capture/replay.
• propose → check_exchange. Covered by the timeout case.

An LLM-level TP2 run is not possible in CI without a DSpark checkpoint (see the previous thread). On GB300 the M3.1 Flash serving smoke (TP4, concurrency 1/8/32, two runs each) logs "DSpark Markov sampling is vocab-parallel over TP." on all 4 ranks in every run, with no fallback warning, no timeout and no error, and mean acceptance length in the same range as the base arm.

[MrGeva]

Hard assert on a hardware-dependent property. inexact_projection_reason depends on which cuBLAS kernels get picked, which the PR says varies by shape and GPU (validated on SM103 only). Asserting is None turns the by-design outcome "feature falls back to the replicated loop" on another GPU or cuBLAS version into a test failure, once the test actually runs in CI on H100/B200. The comment directly above says the fallback is expected where results are inexact.

Either skip with the returned reason (equality is only meaningful when the sampler is eligible), or assert on the reason string per known GPU.

[MrGeva]

Agreed. Fixed in db489e9: the worker now calls TPMarkovSampler.create(). If it falls back (inexact projection on this GPU or cuBLAS, or no symmetric memory), the test skips with the returned reason instead of failing. Equality is only asserted when the sampler is active.

[MrGeva]

Coverage stops at 32 requests, but the feature is enabled up to 128. MAX_REQS = 32 and the batch list is (1, 2, 4, 5, 16, 32). The recompute path at 33-128 requests (where _MAX_REQS, the load-time sweep and the n_spec * num_reqs batched GEMM matter most, with a ~360 MB bias tensor) is covered only by the out-of-tree harness mentioned in the PR body (B=64).

Suggest parametrising MAX_REQS to 128 (or at least 64) with at least one case above 32.

[MrGeva]

Fixed in db489e9. MAX_REQS = 128 now, with cases 1, 2, 4, 5, 33 and 128 requests in both completion modes, plus CUDA-graph replays at 3 and 8. The load-time guard runs over 1…128 in every test. On the real M3.1 Flash draft model (TP4, GB300) I checked B = 1, 2, 4, 8, 16, 32, 64, 96 and 128: all bitwise equal.

[MrGeva]

Memory claim is off by roughly 10x. The staging buffer is 2 * world * push_max_reqs * n_spec * per * (BLOCK_SIZE // 2) int64, i.e. 2447495128 B ~ 43 MiB at TP4, V=200064, K=7 (similar at TP8). The PR says "a few MB (words, plus push staging for 4 requests)". The words buffer adds only ~0.4 MB. It is allocated and zeroed on every rank up front, even for a workload that never drops below push_max_reqs.

Worth correcting the description, and considering whether push staging should be sized for fewer slots.

[MrGeva]

Correct. The push staging buffer is 2 · world · 4 · 7 · per · 512 · 8 B, which is 42.9 MiB at TP2 and TP4 and 43.8 MiB at TP8 for V = 200064, K = 7. The word buffer adds 0.38 MiB at 128 requests (measured: stage 42.9 MiB, words 0.38 MiB per rank). I fixed the description: about 43 MiB of symmetric memory per rank, allocated and zeroed once at load, and only when the option is on.

I left the sizing as is (about 0.015% of a GB300's memory). It could be halved by dropping the parity double buffer, which the consume-and-clear protocol no longer strictly needs, but I'd rather not change the layout in this PR.

[MrGeva]

Recompute mode projects columns it then throws away. Each rank computes the full [num_reqs * n_spec, vocab] bias (reads all of w2, up to ~360 MB of bf16 output at 128 requests), and _tp_markov_fill_kernel then skips this rank's own [own_start, own_end) blocks. Projecting only the peers' row ranges (w2[:col_start], w2[col_end:]) would cut GEMM traffic and the temporary by 1/TP. If this is deliberately avoided because the bitwise guard only covers the full-weight shape, that is worth a comment, since it is the dominant cost of the path used for 5-128 requests.

[MrGeva]

Agreed. Fixed in eddf09a: the recompute path now projects only the peers' rows, w2[:col_start] and w2[col_end:], as two GEMMs (one on the first and last rank). The fill kernel reads the left or right part by block, and the guard checks exactly those two shapes at M = 7·num_reqs, so the cache stays bitwise equal; this was verified at B = 8…128 on the real model.

Measured on GB300 at TP4: the recompute GEMM alone gets 4–7 µs faster on the slowest rank (B8: 23.1 → 16.5 µs, B64: 51.2 → 44.1, B128: 88.2 → 83.1). The middle ranks need two GEMMs, which eats part of the 1/TP saving. At the level of the whole loop, the difference against projecting the full weight is within run-to-run noise (±5 µs at B ≥ 8, ABAB on one node).

[MrGeva]

Load-time sweep is heavy. On every TP rank, for each of up to 128 batch sizes, this runs n_spec + 2 GEMMs against the whole ~100 MB w2, materialises two ~360 MB intermediates at the top end, and blocks on torch.equal (host sync) several times per size. cuBLAS heuristics are piecewise in M, so checking the CUDA-graph capture sizes (what actually runs) plus a sparse set of other sizes would give the same practical coverage at a fraction of the startup cost and peak memory.

[MrGeva]

Thanks. Partly valid: I measured it before changing anything. On GB300 the old guard took 0.26 s per rank at TP4 and 1.24 s at TP2 (first call, including cuBLAS warm-up), with about 1 GiB of transient allocations from the stacked per-step outputs. That is high for what it is, but not a startup problem.

I kept checking every batch size from 1 to max_reqs rather than only the capture sizes plus a sparse set. The draft also runs eagerly or piecewise, for example above the largest capture size and on non-FULL paths, at any batch size, and the guard is about bitwise equality, where interpolating between sizes isn't safe. Instead I made it cheaper (eddf09a):

• no host sync per size: mismatches accumulate in a device flag, read once;
• comparisons per step against model.markov_bias, without materializing all steps;
• empty_cache() afterwards.

On the real M3.1 Flash draft model at TP4, the whole create() (guard plus symmetric-memory setup) now takes 0.42 s, with a 367 MiB peak. The guard is now also stricter: real reference, peers' slices.

[MrGeva]

Design fit / complexity vs. measured gain. This adds ~590 lines: a custom spin-wait and tagging protocol over private torch symmetric-memory APIs, a second cache-write path (own columns, push/unpack, recompute/fill), and a load-time numerical-equality sweep. The measured saving is 33-110 us per draft step, i.e. 0.25-0.9% of the step in the PR's own serving table (loop-only: 59-159 us). The failure modes it adds sit in the shared decode hot path: a GPU spin on a missing peer, tag aliasing, and silent drift if GEMM dispatch changes (see the other comments). The PR notes that an NCCL all-gather and a flag variant were slower, so this is a judgment call, but it seems worth the maintainers explicitly accepting that trade-off, or marking the option experimental in the config docstring.

[MrGeva]

Fair question. Here are the numbers, so maintainers can weigh it.

Scope. The option is opt-in, default off, and now marked Experimental in its docstring. With it off, nothing changes on the shared decode path: _sample_sequential gains one is not None check, and propose one more. With it on, unsupported configurations and inexact GEMMs fall back to the replicated loop on every rank alike, a failed symmetric-memory setup falls back, and a peer timeout raises instead of drifting. The other comments' failure modes (aliasing, silent timeouts, dispatch drift) are fixed and tested in this revision (eddf09a, db489e9, f2ffea6).

What it saves (MiniMax-M3.1 Flash, TP4, K=7, GB300).

• Markov loop, graph-timed on the real draft model (replicated → vocab-parallel, µs): B1 291 → 233, B4 371 → 315, B8 440 → 389, B16 563 → 509, B32 705 → 600, B64 1156 → 995, B128 2010 → 1759. That saves 51–251 µs, or 10–20% of the loop. The absolute saving grows with batch size, because the replicated loop's cost scales with B at full vocabulary on every rank.
• Serving (seeded, same node, nsys; draft CUDA graph p50 per step, µs): concurrency 1: 646 → 586, concurrency 8: 917 → 875–882, concurrency 32: 1368 → 1235–1240. That saves 61, 35–42 and 128–133 µs per step. The whole step is 6.8, 13.4 and 19.7 ms, so this is 0.3–0.9% of the step, as in the original table.

Why this design. A plain NCCL all-gather of the per-block maxima, and a sys-scope flag variant, were both slower than the replicated loop at these sizes (the exchange is about 1.5 KB per row per step, so latency is everything). That is why the exchange uses tagged LL words over symmetric memory. The exchange and sampler are also meant as the base of a planned follow-up (vocab-sharded verification), which would reuse TPMarkovSampler to draft on each rank's owned columns.

If the maintainers prefer not to carry an LL protocol in spec_decode/dspark/ for a sub-1% step gain on its own, I'm fine keeping it out of tree until the follow-up shows the larger combined effect. Let me know which you prefer.

[MrGeva]

_TAG_BITS is defined but never used; only _TOKEN_BITS is (in unsupported_reason). The kernels hard-code (1 << 14) - 1, << 18, >> 18 and (1 << 18) - 1 in several places (sample and combine kernels). Changing either width means finding every literal. Define them as tl.constexpr module constants (as gumbel.py does with _DRAFT_NOISE_SALT = tl.constexpr(1 << 30)) and derive masks and shifts from them.

[MrGeva]

Fixed in eddf09a. _INDEX_BITS = tl.constexpr(10) and _TAG_BITS = tl.constexpr(22) are module constants (as in gumbel.py), and every shift and mask is derived from them. A _tag(e, BITS) helper gives both exchanges their tag (22 bits for the sampling words, 32 bits for push staging), and tl.static_assert(BLOCK_SIZE == 1 << _INDEX_BITS) ties the block size to the layout. _TOKEN_BITS and the vocabulary limit are gone.

[MrGeva]

The shard geometry (vocab = min(w2.shape[0], org_vocab_size), num_blocks, per, col_start / col_end) is computed twice: in __init__ (L334-341) and again here. If one drifts, the guard validates a different slice than the sampler uses. Factor one helper returning (vocab, num_blocks, per, col_start, col_end). Also vllm.utils.math_utils.cdiv (already imported by speculator.py) could replace triton.cdiv.

[MrGeva]

Fixed in eddf09a. The geometry (vocab, blocks, per, col_start/col_end, the row slices) is computed once in TPMarkovSampler.__init__, which no longer allocates or communicates. create() builds the sampler first and then runs the guard (_inexact_projection_reason, now an instance method) on that same object, so the guard checks exactly the slices the sampler uses. triton.cdiv was replaced with vllm.utils.math_utils.cdiv.

[MrGeva]

TPMarkovSampler captures part of the speculator in __init__ (w2 views, n_spec, max_reqs) but sample() takes the speculator again on every call to reach model, draft_logits.values, draft_tokens, input_buffers, _anchor_idx (a private attribute), seeds and temperature. That leaves two sources of truth: self.n_spec vs speculator.num_speculative_steps, and vocab_size = base_logits.shape[-1] for the kernels vs self.vocab_size used to derive num_blocks / per. A mismatch would silently mis-index the words buffer. Keep the reference from construction and assert base_logits.shape[-1] == self.vocab_size.

[MrGeva]

Fixed in eddf09a. The sampler keeps the speculator reference from construction, and sample() no longer takes it. The loop uses self.n_spec throughout. The kernels use self.vocab_size, with assert base_logits.shape[-1] == self.vocab_size, so the words-buffer geometry and the logits cannot disagree.

[MrGeva]

I folded vocab-sharded speculative verification (the opt-in vocab_sharded_verification) into this PR, at 9741123. It was reviewed separately as a stacked PR. I merged it here because it uses the same exchange machinery as the vocab-parallel Markov loop: tagged symmetric-memory words, the bounded spin with a host-memory timeout flag, and the setup fallback that all ranks agree on. Its draft side is also this PR's Markov loop running on the owned columns, so it could not land on its own. Both options stay off by default. The description now covers both parts.

[mergify]

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

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