[GG] fix(spec-decode): centralize probabilistic draft sampling

[voipmonitor]

Summary

Centralize probabilistic draft Gumbel sampling in DraftModelSpeculator._sample_probabilistic_draft and route the generic, autoregressive, and DSpark paths through it.

The shared method always applies draft_gumbel_pos(), keeping proposal noise disjoint from verifier acceptance/recovery offsets. The autoregressive caller still forwards active_num_reqs, preserving its FULL-graph padded-row safeguard.

Fixes #205.
Supersedes #209.

Why this replaces #209

#209 correctly fixes the live GG autoregressive override, but only changes that call site. GG already contains the broader salted-stream backport corresponding to upstream vLLM vllm-project#47386; a later GG autoregressive override copied the old upstream positions + 1 call and bypassed the corrected base implementation.

Keeping raw gumbel_sample() calls in individual speculators makes that failure mode easy to repeat. This PR makes stream ownership a base-speculator responsibility instead:

• generic probabilistic drafts use the shared method;
• the GG autoregressive override uses it while retaining active_num_reqs;
• DSpark sequential probabilistic sampling uses the same method;
• DFlash continues to inherit the generic path.

Only the base method calls gumbel_sample() in spec_decode; individual speculators pass model positions, not Philox offsets.

Correctness scope

Affected before this fix:

• V2 autoregressive speculators;
• probabilistic draft sampling;
• sampled target requests.

Unaffected:

• greedy drafts and temperature-zero requests;
• the V1 proposer;
• the all-accepted bonus draw.

The first-order defect is rejection recovery reusing the exact vocabulary-sized Gumbel field that selected the rejected draft token. Immediate acceptance probability is unchanged; the fix restores the recovery sampler's independence contract.

Runtime invariants

• No new allocation, kernel launch, synchronization, or public API.
• The same single tensor add is used to derive positions + 1 + 2**30.
• AR FULL-graph active-row protection is unchanged.
• DSpark's existing predicted-position convention is unchanged.
• Probabilistic outputs can change bit-for-bit as expected from correcting the RNG stream.

Verification

Release-container tests:

4 passed

• AR tuple/tensor behavior
• shared salted-offset contract
• AR forwarding of the FULL-graph active-row limit

GPU tests on CUDA:

11 passed

• probabilistic rejection output-distribution regression
• Gumbel sampler tests

Static checks:

Ruff: pass
format: pass
git diff --check: pass

The existing statistical regression inherited from upstream vllm-project#47386 measures the coupled stream outside its target-distribution noise floor and the salted stream within it; this PR preserves that implementation and closes the GG override gap without claiming an unmeasured model-level quality delta.

Summary by CodeRabbit

• Bug Fixes

  • Improved probabilistic draft-token sampling across autoregressive and DSpark decoding.
  • Preserved sampling positions, random seeds, temperature, and active-request state for more consistent results.
  • Added support for target-vocabulary mappings and safer handling of processed logits.

• Tests

  • Added coverage for sampler arguments, disjoint random-sampling positions, and state preservation during autoregressive sampling.

[coderabbitai]

📝 Walkthrough

Walkthrough

The base speculator now owns probabilistic draft sampling. Autoregressive and DSpark speculators delegate to this shared helper, preserving draft position handling, seeds, temperature, step metadata, processed logits, and active-row state. Tests cover sampler arguments and delegation.

Changes

Probabilistic draft sampling

Layer / File(s)	Summary
Centralize probabilistic sampling vllm/v1/worker/gpu/spec_decode/speculator.py	Adds _sample_probabilistic_draft and routes sample_draft through the shared helper, including optional active-row filtering for processed-logit writes.
Adopt shared sampling in speculators vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py, vllm/v1/worker/gpu/spec_decode/dspark/speculator.py, tests/v1/worker/test_gpu_autoregressive_speculator.py	Autoregressive and DSpark sampling delegate through the shared helper, while tests verify Philox position offsets, sampler arguments, and preserved delegation state.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related PRs

• local-inference-lab/vllm#209: Directly related Philox-stream changes for autoregressive probabilistic draft sampling.

Suggested reviewers: lukealonso, mgoin

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: centralizing probabilistic draft sampling for speculative decoding.
Linked Issues check	✅ Passed	The PR centralizes probabilistic sampling, applies salted draft positions, updates the autoregressive path, and adds tests covering disjoint Philox offsets.
Out of Scope Changes check	✅ Passed	The generic, autoregressive, and DSpark updates plus targeted tests directly support the linked issue and stated centralization objective.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/gg-centralize-probabilistic-draft-sampling-20260730

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

vllm/v1/worker/gpu/spec_decode/dspark/speculator.py (1)

253-264: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add a DSpark regression test for the position translation.

Line 258’s -2 adjustment is a DSpark-specific Philox contract; the added tests cover only the base helper and autoregressive delegation. Assert this call forwards sample_pos[:, i] - 2 so a future cleanup cannot shift DSpark’s randomness stream.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@vllm/v1/worker/gpu/spec_decode/dspark/speculator.py` around lines 253 - 264,
Add a DSpark regression test covering the call from the relevant speculator
method to _sample_probabilistic_draft, and assert its positions argument is
exactly sample_pos[:, i] - 2. Keep the assertion focused on preserving this
DSpark-specific position translation and randomness contract.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@vllm/v1/worker/gpu/spec_decode/dspark/speculator.py`:
- Around line 253-264: Add a DSpark regression test covering the call from the
relevant speculator method to _sample_probabilistic_draft, and assert its
positions argument is exactly sample_pos[:, i] - 2. Keep the assertion focused
on preserving this DSpark-specific position translation and randomness contract.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 07142b45-8717-4b54-93dd-25af02a2c48a

📥 Commits

Reviewing files that changed from the base of the PR and between 443e1d6 and 7bf5ddd.

📒 Files selected for processing (4)

• tests/v1/worker/test_gpu_autoregressive_speculator.py
• vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
• vllm/v1/worker/gpu/spec_decode/dspark/speculator.py
• vllm/v1/worker/gpu/spec_decode/speculator.py

[malaiwah]

Rental-machine validation — PASS (focused and composed-source gates)

Tested merged commit: 30038602b71395f481ef4a6edfe4fcf8551d9c15
Original PR head: 7bf5ddd98ffc1d503ad6017e9a0e8edd4bc9bf0d
Tested branch: dev/gilded-gnosis
Composed stack: #211 + #210 at b135a9b79377386ee805c22dee950e985f07ed28

Host: Vast.ai 46335896, Ubuntu 24.04.4 / kernel 7.0.0-28, 2x EPYC 9455, 1.5 TiB RAM.
GPU/topology: 4x RTX PRO 6000 Blackwell; all NODE on NUMA 0; all 12 directed P2P copies passed.
Driver/CUDA/PyTorch: 610.43.02 / CUDA UMD 13.3 / PyTorch 2.12.0+cu132 (CUDA 13.2).
Test environment: retained turnkey container; isolated cold/warm caches.

Commands and results

• pytest tests/v1/worker/test_gpu_autoregressive_speculator.py — 4 passed; Ruff clean.
• GPU Gumbel/rejection-sampler suites — 11 passed cold (19.75 s), 11 passed warm (16.84 s).
• Focused DSpark spy — PASS; caller supplied exactly sample_pos[:, i] - 2.
• With [GG] Bound EXL3 Trellis prefill arena capacity #210 — 16 combined CPU tests passed; [GG] fix(spec-decode): centralize probabilistic draft sampling #211 GPU suites passed 11 cold and 11 warm.

Greedy/temperature-zero coverage passed and no capture-time compile, Xid, hang, or request error was observed. The later full-model live-probabilistic gate was not run because independent SparkInfer and LMCache source-composition gates failed first.

Report and durable logs: field-review result

No unreported source modifications were present; the tested worktrees were clean.

---

Exact r14-base composition — PASS

The immutable r14 vLLM lock was replayed from base f978d009fab996617f9a3cadef36ce727bcd83cd and independently matched published tree 749050edab1b6664937c52fa1b0be360be632c1e. Original #211 head 7bf5ddd98ffc1d503ad6017e9a0e8edd4bc9bf0d plus exact #210 then applied without conflict.

Resulting candidate 5a7ac1481e1fc24b0bbc35efe160b2ff34797bee passed 16/16 CPU/static tests and the #211 probabilistic GPU union 11/11 cold and 11/11 warm on physical GPU 0. Ruff and diff --check were clean; unchanged r14 native binaries used for worktree import were hash-recorded and removed after the gate.

This extends the earlier current-development-branch result to the actual r14 image source boundary. Exact evidence: r14 ledger at 4d2aa18.

[malaiwah]

Live full-model probabilistic follow-up — PASS

Tested merged change: 30038602b71395f481ef4a6edfe4fcf8551d9c15 (original PR head 7bf5ddd98ffc1d503ad6017e9a0e8edd4bc9bf0d)
Exact r14 composed vLLM candidate: 5a7ac1481e1fc24b0bbc35efe160b2ff34797bee
Host: Vast.ai 46335896, 4x RTX PRO 6000 Blackwell, driver 610.43.02 / CUDA UMD 13.3 / PyTorch 2.12.0+cu132.

The previously missing live gate is now complete. The full GLM-5.2 3.25-bpw EXL3 server ran the V2 runner with MTP3, draft_sample_method=probabilistic, temperature 1.0, and a fixed seed. This does execute #211: method=mtp selects MTPSpeculator, which inherits AutoRegressiveSpeculator; non-null probabilistic draft logits delegate to the shared salted sampler.

Observed across 32 live requests / 4,096 generated tokens:

• 1,526 speculative steps;
• 4,578 drafted tokens;
• 2,563 accepted tokens;
• pooled mean acceptance length: 2.680;
• pooled draft-token acceptance: 55.99%;
• 0 request errors, 0 preemptions;
• all short/32K/65K/126K correctness and retrieval checks passed without degeneration.

The exact-r14 focused distribution suite already passed 11/11 cold and 11/11 warm, including greedy and temperature-zero controls. A separate plain-checkpoint download is unnecessary: this EXL3 MTP path reaches the same changed sampler.

Log scope: cold priming compiled 28 previously unseen SparkInfer request shapes, which the test plan permits. The subsequent warm audit contains zero compile, ERROR, CUDA/IPC, or process-failure findings. Four XGrammar FSM errors occurred earlier in structured-output checks because this exact-source gate deliberately disabled the appliance compatibility patch to preserve SHA attribution; that separate built-image requirement is not attributable to #211.

Evidence: campaign ledger, raw benchmark JSON, and clean warm-window audit.