[Bugfix][Model Runner V2][Spec Decode] Decouple draft Gumbel stream from acceptance/recovery noise

[WoosukKwon]

Purpose

With draft_sample_method="probabilistic", the draft token proposed for position P was sampled with Gumbel noise keyed by Philox offset P (sample_draft passed positions + 1, with a comment saying this was "to match the Gumbel noise used for draft and target sampling"). The rejection sampler keys two other random quantities by the same offset:

• the acceptance uniform: u = tl.rand(seed, P), which is bit-for-bit the float conversion of the very Philox draw used as the draft's Gumbel key (tl.randint(seed, P)), and
• the recovery Gumbel noise on rejection (gumbel_block_argmax derives randint(seed, P)), which is therefore the identical noise vector that selected the rejected draft token.

Rejection sampling (Leviathan et al.) requires the draft proposal, acceptance uniform, and recovery draw to be independent; this coupling makes the recovery draw re-favor exactly the tokens the draft noise favored, biasing the output marginal toward draft-preferred tokens. Deliberately intent-matched noise sharing only preserves the target distribution under an accept-iff-equal verification rule, not under the u < p/q rule this sampler uses.

The fix salts the draft-side Philox offsets into a range disjoint from all target-side offsets (positions are bounded by max_model_len ≪ 2³⁰). Acceptance probability is E_q[min(1, p/q)] independent of the coupling, so this costs no acceptance rate. The default draft_sample_method="greedy" consumes no Gumbel noise and is unaffected, as are temperature-0 requests.

Found while auditing the V2 spec-decode sampling path for #47239 (not that issue's cause — the reporter's config uses the default greedy drafts).

Measured effect

262k-sample output-marginal test (V=32k, k=1, temp=1.0, p and q ranked oppositely over 8 tokens), GB200:

draft Gumbel key	TV(output, target p)	acceptance
coupled (old, offset P)	0.01251	0.4994
salted (this PR)	0.00144	0.4987
independent control	0.00211	0.4990

Noise floor ≈ 0.0028; theoretical acceptance Σmin(p,q) = 0.5.

Why this is not duplicating an existing PR

gh pr list --search "gumbel draft probabilistic" and related searches — no open PR touches the draft sampling keys or this bias.

Test Plan

• New deterministic GPU regression test tests/v1/worker/test_gpu_rejection_sampler_dist.py (65k samples, fixed seeds): asserts TV(output marginal, target) < 0.008 and acceptance ≈ 0.5 using the production key derivation (draft_gumbel_pos). With the old coupled key the same test measures TV=0.0140 and fails.
• Existing tests/v1/worker/test_gpu_autoregressive_speculator.py and test_gpu_gumbel_sample.py.

Test Result

$ pytest tests/v1/worker/test_gpu_rejection_sampler_dist.py -q
1 passed
$ pytest tests/v1/worker/test_gpu_autoregressive_speculator.py tests/v1/worker/test_gpu_gumbel_sample.py -q
12 passed
# same dist test with the old coupled key (positions + 1):
AssertionError: output marginal deviates from target: TV=0.01399  (1 failed)

Pre-commit (ruff, mypy) passes on changed files.

Notes

Two related, lower-order couplings are left as-is (pre-existing): re-proposals of the same position after a rejection reuse that position's draft key across engine steps, and Gemma4-style constant-position speculators reuse the key across draft steps within one proposal. Both need a per-proposal counter to fix cleanly and neither shows above the noise floor in the measurement here.

AI assistance disclosure

Developed with AI assistance (Claude Code); I reviewed every changed line and ran the tests above.

🤖 Generated with Claude Code

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review to trigger a review and subscribe this PR to future pushes, or @claude review once for a one-time review.

_{Tip: disable this comment in your organization's Code Review settings.}

[WoosukKwon]

Addressed the review: DSparkSpeculator's sequential Markov sampling was calling gumbel_sample directly with key Q-1 (the verification key), bypassing sample_draft — so its draft_sample_method="probabilistic" path had the same coupling. Now routed through draft_gumbel_pos with the DFlash convention (pass Q-2, salted key = Q-1 + DRAFT_GUMBEL_POS_OFFSET), keeping each per-position key unchanged modulo the salt. Verified this is the last direct gumbel_sample caller in the draft paths (the only other caller is the target-side sampler, which must stay unsalted).

[mergify]

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

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