Skip to content

[Spec Decode] Verify in head dtype and drop the FP32 chunk workaround - #53630

Open
jyan-R wants to merge 14 commits into
vllm-project:mainfrom
jyan-R:claude/pr2-remove-fp32-verifier-copy
Open

jyan-R wants to merge 14 commits into
vllm-project:mainfrom
jyan-R:claude/pr2-remove-fp32-verifier-copy

Conversation

@jyan-R

@jyan-R jyan-R commented Aug 24, 2026

Copy link
Copy Markdown
Contributor

Purpose

Speculative decoding expands each request into multiple verification positions. Whenever
any active logits processor is present, the target rejection verifier currently
materializes an FP32 processing canvas of shape [num_verification_tokens, vocab_size].
The production path bounds the transient allocation with a 1 GiB chunk, which in turn
forces a cap on the total logits adaptive verification may schedule per step.

This PR makes the target verifier explicitly opt into storing and processing logits in
the LM head dtype (BF16/FP16), and deletes the FP32 canvas, the chunk iterator, and the
adaptive cap:

  • temperature is applied in FP32 when the rejection kernels, top-k/top-p, and the
    processed-logprob scorer load the logits; it is never written back to the stored
    logits;
  • top-k/top-p computes its boundary in FP32 and writes back only the mask; penalties are
    computed in FP32 and written back to the head dtype with a single rounding; exact mask
    processors keep their original semantics;
  • raw reporting keeps an original head-dtype snapshot only when a processor would
    overwrite the input; processed reporting reads the processed stored logits;
  • the adaptive cost tables, selection formula, grammar mapping, and zero-budget contract
    are unchanged; only the hard cap that existed to serve the FP32 chunk is removed;
  • the PyTorch top-p fallback's softmax and cumulative sum are pinned to FP32, since
    head-dtype target logits make the previously latent BF16 accumulation path reachable;
  • speculative decoding fails closed on any nonzero min_p, and custom
    logprob_token_ids reporting is wired up.

The regular sampler, non-speculative decoding, and zero-draft steps keep the existing
FP32 processing path.

Supported sampling processors and reporting behavior

All supported processors below avoid materializing a full-vocabulary FP32 verifier canvas. This
PR does not change which processors are supported by speculative decoding.

Sampling parameter / processor Property Supported in Spec Decode Effect on stored logits Head-dtype precision
temperature Distribution scaling Yes (unchanged) None; consumers scale on load Computed in FP32
top_k / top_p Selection mask Yes (unchanged) Preserve kept values; mask others with -inf Exact
repetition / frequency / presence penalties Value transformation Yes (unchanged) Rewrite affected values One BF16/FP16 rounding
allowed_token_ids / min_tokens / bad_words Hard token mask Yes (unchanged) Write -inf masks Exact
thinking_token_budget Forced-token override Yes (unchanged) Write a dtype-safe sentinel Semantically exact
structured-output grammar mask Hard token mask applied before verification Yes (unchanged) Write -inf masks Exact
min_p Probability filter No (unchanged) Any nonzero value fails validation
logit_bias Additive value transformation No (unchanged) One BF16/FP16 rounding

Raw and processed logprobs are reporting modes rather than logits processors. Raw reporting reads
a pre-processor head-dtype snapshot when stored logits may be modified. Processed reporting reads
the processed stored logits and applies temperature in FP32 on load. Adaptive-verification output
logprobs were enabled upstream by #52242 (merged 2026-08-25); this PR is rebased on top of it and
preserves its device-tensor boundary reporting through the head-dtype scorer.

Penalties are the only path that writes back to the stored logits lossily: the processor
still computes in FP32, but the result is rounded to BF16/FP16 once. Please pay
particular attention to the penalty-only, stochastic, and top_k=1 results under
Numerical fidelity and model quality below. We did not find an implementation that
preserves FP32-equivalent penalty output without either reintroducing full-vocabulary
FP32 storage or adding significant extra scans/latency; this PR therefore keeps this
numerical trade-off explicit rather than claiming processed penalty logprobs are bitwise
FP32.

Why this is not a repeat of #48641

#48641 removed the same FP32 copy and was reverted in full three days later by #49033
over a never-root-caused DP+EP+MTP hang. Beyond removing the copy, #48641 also changed
the numerics: it applied temperature by writing the scaled values back into the
stored head-dtype logits
, for the regular sampler and the speculative verifier
alike. Every such write-back costs an extra BF16/FP16 rounding, and every downstream
consumer — rejection statistics, top-k/top-p, logprob reporting — then reads the
already-rounded, temperature-scaled values. Whether or not that contributed to the
hang, it is a real precision hazard, and this PR is built specifically to remove it:

  • stored logits always keep their pre-temperature values; temperature is applied in
    FP32 registers at every load site and never touches storage;
  • top-k/top-p computes its boundary in FP32 and writes back only -inf masks — kept
    logits are bit-identical to what the LM head produced;
  • the only remaining lossy write-back is penalties, which is quantified and
    guardrailed under Numerical fidelity and model quality below;
  • the scope is the speculative verifier only; the regular sampler keeps its FP32 path.

On the risk side, a permanent e2e test covers both #49033's pure-greedy control path
and an active path that actually enters the head-dtype verifier, in the same
distributed engine. That test guards the risk quadrant; it does not claim to reproduce
or explain the root cause of the original hang.

Related work

Test Plan

  1. Run the focused CUDA test matrix for sampling-parameter validation, sampler state,
    thinking budget, top-k/top-p, rejection sampling, adaptive verification, and the
    logprob numerical oracles and head-dtype penalty path.
  2. Run the complete test_spec_decode_logprobs matrix across 3 speculative model setups and all
    4 LogprobsMode values, comparing emitted tokens and raw/processed logprob outputs with the
    non-speculative reference.
  3. Run the DP2+EP+MTP regression twice. In the same distributed engine, exercise both the
    pure-greedy Revert "[Sampler] Stop upcasting logits to fp32 in apply_sampling_params" (#48641) #49033 control signature and an active-sampling request that enters the head-dtype
    verifier; require each wave to draft and match the non-speculative output.
  4. Benchmark RejectionSampler.__call__ with BF16 logits, K=7, V=129280, and
    R={8,256,512} across temperature, top-p, penalties, combined processors, and logprob
    reporting. Compare p50 latency and per-call CUDA peak allocation with the FP32-canvas control.
  5. Run a paired 1319-question, 5-shot GSM8K A/B evaluation covering the default path,
    penalties-only, stochastic top-p with penalties, and top_k=1; compare accuracy, invalid
    outputs, and acceptance length.
Repository test commands
.venv/bin/python -m pytest -q \
  tests/test_sampling_params.py \
  tests/v1/worker/test_gpu_sampler_flags.py \
  tests/v1/worker/test_gpu_thinking_budget.py \
  tests/v1/sample/test_topk_topp_sampler.py \
  tests/v1/spec_decode/test_rejection_sampler_utils.py \
  tests/v1/spec_decode/test_adaptive_verification.py \
  tests/v1/sample/test_logprobs.py::test_token_logprobs_apply_virtual_temperature
.venv/bin/python -m pytest -vv -s \
  tests/v1/sample/test_logprobs.py::test_spec_decode_logprobs
.venv/bin/python -m pytest -vv -s \
  'tests/v1/e2e/spec_decode/test_mtp_parallel_load.py::test_deepseek_mtp_load_dp[dp2]'

Test Result

Correctness and path coverage

  • The targeted GPU suite passed with 231 passed, 2 skipped (reported from the run; that console
    log was not kept). It covers every supported sampling processor, raw/processed reporting,
    adaptive verification, thinking budget and the head-dtype penalty path.
  • The full public-model test_spec_decode_logprobs matrix passed all 3 model setups × 4
    LogprobsMode combinations (12 passed). Spec and non-spec runs returned identical emitted
    token IDs, logprob token IDs and ranks. Raw scores were bitwise equal; the largest processed-score
    difference was 0.08334.
  • The DP2+EP+MTP surrogate passed twice consecutively with a minimum spec/non-spec token match ratio
    of 1.000. Both the pure-greedy wave matching the Revert "[Sampler] Stop upcasting logits to fp32 in apply_sampling_params" (#48641) #49033 incident signature and the active
    sampling wave produced draft tokens, so the test covers both the control path and the new
    head-dtype verifier in the same distributed engine.
  • ruff check and ruff format --check passed for all 19 changed Python files.

These results verify that the optimization is reached in production paths while preserving token
selection and logprob reporting contracts; they do not rely only on isolated kernel tests.

Verifier memory and latency

The benchmark uses BF16 logits with K=7, V=129280 and R=512 (4096 verification rows), which
crosses the old approximately 2076-row chunk boundary. The control is the same candidate with the
head-dtype verifier disabled, so the comparison isolates this PR's verifier mode while keeping the
PyTorch top-p correctness fix on both sides.

Processor FP32 canvas p50 (ms) Head dtype p50 (ms) Per-call peak allocated
temperature 5.054–5.103 0.979–1.093 1023.20 → 2.33 MiB
top-p 31.178–31.219 25.846–26.057 1023.20 → 2.33 MiB
penalties 11.960–12.025 5.316–8.148 1023.20 → 2.33 MiB
penalties + top-p 39.058–39.361 30.424–34.095 1023.20 → 2.33 MiB

The measured processor cases reduce the per-call allocation from 1023.20 MiB to 2.33 MiB and
improve p50 latency in every case; temperature-only is approximately five times faster. Processed
logprob reporting similarly falls from 1032.44 MiB to 20.68 MiB. This demonstrates that the FP32
canvas is removed rather than moved behind a different chunk boundary.

Raw reporting with a mutating processor is the intentional exception: preserving pre-processor
values still requires an approximately 1010 MiB BF16 snapshot, for a total around 1030.66 MiB. The
claim is therefore the removal of the verifier FP32 canvas and its local memory/latency cost, not a
general end-to-end throughput improvement.

Numerical fidelity and model quality

Mask-only processors remain exact because kept logits are never rewritten. Temperature and top-p
processed scores have maximum finite errors of 1.91e-6 and 2.13e-4 against the FP32-input
reference, with no rank differences. Raw-reporting digests are bitwise identical between base and
head-dtype runs.

Penalties are the only lossy stored-processor path. Their FP32 calculation is rounded once on
writeback to BF16/FP16; the penalty + top-p oracle measures a maximum finite score error of
0.0203. We therefore evaluated all penalty-enabled combinations on the full 1319-question,
5-shot GSM8K set, in addition to the default control:

Case Sampling Base Head dtype(PR) Accuracy Δ Acceptance length (base/head)
standard greedy; neutral penalties(off) 82.942% 82.942% 0.000 pp 4.2015 / 4.1509
penalties-only greedy; penalties on 82.259% 81.577% −0.682 pp 4.0394 / 4.0341
stochastic temperature=0.7, top_p=0.9; penalties on 79.075% 79.606% +0.531 pp 3.9586 / 3.9658
top_k=1 temperature=0.7, top_p=0.9; penalties on 82.259% 81.653% −0.607 pp 3.9686 / 3.9878

The standard control produced identical output token sequences and correctness for all 1319
questions. Invalid counts were unchanged in every paired run. The three penalty-enabled,
fixed-seed accuracy deltas stay within ±0.682 percentage points, while their acceptance lengths
remain close to base. These are quality guardrails rather than a claim of statistical equivalence;
together with the FP32 oracle they bound the observed effect of the intentional head-dtype penalty
writeback.

Duplicate check

Checked RFC #42259 and open PRs by issue number and relevant area keywords on 2026-08-24, and
re-checked on 2026-08-25. The only directly overlapping open PR is #53090; the scope difference is
described above.

AI assistance

AI assistance was used for implementation, test design, GPU experiment orchestration and drafting
this description. The human submitter reviewed lines, resolved rebase conflicts and
ran the tests above, and can explain and defend the change end to end.


Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan, such as providing test command.
  • The test results, such as pasting the results comparison before and after, or e2e results
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing (anything written below this line will be removed by GitHub Actions)

@mergify

mergify Bot commented Aug 24, 2026

Copy link
Copy Markdown
Contributor

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

@mergify mergify Bot added documentation Improvements or additions to documentation speculative-decoding mrv2 Model Runner V2 specific labels Aug 24, 2026
@jyan-R
jyan-R marked this pull request as ready for review August 24, 2026 21:22
Copilot AI lite review requested due to automatic review settings August 24, 2026 21:22

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

mergify Bot commented Aug 25, 2026

Copy link
Copy Markdown
Contributor

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

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

@mergify mergify Bot added the needs-rebase label Aug 25, 2026
@jyan-R
jyan-R force-pushed the claude/pr2-remove-fp32-verifier-copy branch from 16aa5dd to 55bae3f Compare August 25, 2026 14:36
@mergify mergify Bot removed the needs-rebase label Aug 25, 2026
jyan-R and others added 10 commits August 26, 2026 01:33
`apply_top_k_top_p_pytorch` accumulated the nucleus in the logits dtype.
With a BF16 input the running sum crosses `1 - top_p` too late and extra
low-probability tokens stay in the nucleus: for `linspace(-4, 4)` at
`top_p=0.9` it keeps 2 extra tokens at `V=8192` and 23 at `V=129280`.

`apply_top_k_top_p_triton` asserts float32 input, and the two adjacent
softmax calls in this file already pass `dtype=torch.float32`. This makes
the third one consistent.

Every in-tree caller upcasts to float32 before reaching this helper, so
the defect is latent on main today. It becomes reachable once the
speculative verifier hands head-dtype logits to the sampler.

Co-authored-by: OpenAI Codex <codex@openai.com>
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: jyan <r_02213@mail.sjtu.edu.cn>
`apply_top_k_top_p_triton` accepts an optional per-row `temperatures`
tensor and a non-FP32 logits dtype. When temperatures are supplied every
pivot, probability and tie comparison reads `fp32(stored) / t_eff`, while
the final write keeps the original stored value for kept positions and
`-inf` for dropped ones. A zero temperature maps to an effective 1.0.

This keeps the mask decision in the scaled space without ever writing a
scaled value back to memory, so a caller can hold pre-temperature logits
in the model head dtype.

The scratch buffer and the two lookup tables are pinned to FP32 instead
of following the logits dtype, so a mixed-dtype workload no longer keeps
one full-vocab buffer per dtype in the cache.

Existing callers pass no temperatures and are unaffected; the dispatch in
`apply_top_k_top_p` only diverts small batches to Triton when
temperatures are present, because the PyTorch path cannot apply them.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: jyan <r_02213@mail.sjtu.edu.cn>
`compute_token_logprobs` and `compute_topk_scores` accept an optional
per-row `temperatures` tensor. When supplied, the max, sum-exp and the
selected-token gather all read `fp32(stored) / t_eff`; a zero temperature
maps to an effective 1.0. Top-k ids and ranks only depend on ordering,
which division by a positive temperature preserves, so they are computed
from the stored values unchanged.

The two identical `logits_mode` score blocks are hoisted out of the
custom-token-ids branch so the scaling lives in one place. The
`masked_fill` stays gated on the branch that produces `valid_mask`.

Existing callers pass no temperatures and are unaffected.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: jyan <r_02213@mail.sjtu.edu.cn>
`Sampler.apply_sampling_params` gains `use_head_dtype`. With it set the
caller owns the incoming tensor: no FP32 `[N,V]` copy is allocated, the
exact stored processors (logit bias, penalties, bad words, thinking
budget) write the head dtype in place, temperature is left to the
consumers, and min_p is skipped because speculative decoding rejects it.

Because temperature is no longer materialized, the existing
`needs_logits_processing` cache is too wide to decide whether the stored
logits change. `needs_stored_logits_processing` is the narrower contract:
it drops temperature-only and min_p and gates the head-dtype path.

The thinking-budget sentinel is clamped to the target dtype maximum. At
FP16 the previous `1.0e9` became `+inf`, which turns the row logsumexp
into `inf` and every logprob into NaN.

No production caller passes `use_head_dtype=True` yet.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: jyan <r_02213@mail.sjtu.edu.cn>
`RejectionSampler._verify` is now the only caller that passes
`use_head_dtype=True`, so a legal speculative batch never materializes an
FP32 `[N,V]` processing canvas. `__call__` runs one `_verify` over the
whole batch. Raw reporting keeps one head-dtype snapshot, and only when a
stored processor would actually overwrite the values; processed reporting
reads the stored logits and divides by the request temperature on load.

The rejection kernels now interpret target logits the same way they
already interpreted draft logits: as pre-temperature values. Local stats,
the global token logprob and logsumexp, the residual mass and the
resample all divide by the request temperature in FP32 registers. Greedy
requests keep the unscaled argmax, and the gumbel draw after the residual
is already in temperature space and is not scaled again. The existing NaN
guards on the local-max buffers are unaffected, since NaN survives the
division.

With no FP32 buffer to bound there is nothing left to chunk, so
`MAX_CHUNK_BYTES`, `_FP32_BYTES`, `get_max_chunk_logits`,
`_iter_request_chunks`, `_verify_in_chunks` and the chunk logprob
concatenation are removed. `AdaptiveVerificationManager` loses the
`max_total_logits` constructor argument and the hard truncation of
scheduled drafts that only existed to keep a compacted batch inside one
chunk; its cost curve, grammar mapping and budget logic are unchanged.
`max_draft_budget` is now simply the number of scheduled draft scores.

The head-dtype path never applies `min_p`, so speculative validation now
rejects any `min_p > 0` instead of letting `0 < min_p <= 1e-5` slip through
and be silently ignored. The penalty kernels still compute in FP32 but
write the result back in the head dtype; the spec-decode docs state that
rounding and that `--generation-config auto` can pick penalties up from
the model's `generation_config.json`.

`tests/v1/worker/test_gpu_rejection_sampler_chunking.py` only covered the
deleted wrapper via a fake `_verify`, so it is deleted with it. Kernel
partition equivariance stays covered by the rejection-utils tests.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: jyan <r_02213@mail.sjtu.edu.cn>
The non-speculative sampler already honours `SamplingParams.logprob_token_ids`,
but `RejectionSampler` only consulted `max_num_logprobs`, so a speculative
request that asked for custom token logprobs silently received none.

`__call__` now also asks `LogprobTokenIdsState` for the batch's largest
custom list and returns logprobs when either source wants them, and
`_get_logprobs_tensors` forwards the state and the expanded request
mapping to `compute_topk_scores`, which already knows how to gather both.

`test_spec_decode_logprobs` gains a custom `logprob_token_ids` request so
the speculative output is compared against the non-speculative reference.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: jyan <r_02213@mail.sjtu.edu.cn>
`test_deepseek_mtp_load_dp` is extended so it keeps the vllm-project#49033 incident
signature and also reaches the new path:

- DP2 now runs with expert parallelism, async scheduling and CUDA graphs
  on SM90+, matching the parallelism layout of the nightly that failed
  while vllm-project#48641 was in tree;
- the first wave stays pure greedy, which is exactly the request shape
  that failed there and in which every processor early-exits;
- a second wave uses `temperature=0.7, top_k=1` so the head-dtype
  verifier is actually exercised in the same engine;
- `vllm:spec_decode_num_drafts` is summed across DP ranks and its delta
  asserted non-zero per wave, so neither wave can pass on the other's
  drafts;
- the engine shutdown and the test itself are bounded, so a hang reports
  as a failure instead of stalling the job.

`test_spec_decode_logprobs` gains an active-sampling request alongside
the existing default, penalty and custom-token-ids ones, so all four
logprobs modes are compared against the non-speculative reference on a
request that reaches the new path.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: jyan <r_02213@mail.sjtu.edu.cn>
Co-authored-by: OpenAI Codex <noreply@openai.com>

Signed-off-by: jyan <r_02213@mail.sjtu.edu.cn>
Co-authored-by: OpenAI Codex <noreply@openai.com>

Signed-off-by: jyan <r_02213@mail.sjtu.edu.cn>
Co-authored-by: OpenAI Codex <noreply@openai.com>

Signed-off-by: jyan <r_02213@mail.sjtu.edu.cn>
@jyan-R
jyan-R force-pushed the claude/pr2-remove-fp32-verifier-copy branch from 55bae3f to 4855de2 Compare August 25, 2026 17:37
@mergify

mergify Bot commented Aug 29, 2026

Copy link
Copy Markdown
Contributor

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

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

…2-verifier-copy

# Conflicts:
#	tests/v1/spec_decode/test_rejection_sampler_utils.py

Signed-off-by: jyan <r_02213@mail.sjtu.edu.cn>
@mergify mergify Bot removed the needs-rebase label Aug 31, 2026
…2-verifier-copy

# Conflicts:
#	vllm/v1/sample/ops/topk_topp_sampler.py
#	vllm/v1/sample/ops/topk_topp_triton.py

Signed-off-by: jyan <r_02213@mail.sjtu.edu.cn>
…2-verifier-copy

# Conflicts:
#	tests/test_sampling_params.py

Signed-off-by: jyan <r_02213@mail.sjtu.edu.cn>
@mergify

mergify Bot commented Sep 12, 2026

Copy link
Copy Markdown
Contributor

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

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

@mergify mergify Bot added the needs-rebase label Sep 12, 2026
…2-verifier-copy

Resolves conflicts with vllm-project#56323 (sampling kernels launched through the JIT
warmup dispatcher) and vllm-project#56122 (speculative watermarking):

- Thread `temperatures` through the `_topk_topp` dispatcher so the
  verifier keeps scaling stored head-dtype logits on load; the
  dispatcher's SPLIT_COVERS_PONLY now mirrors use_split (rows owing a
  temperature division never take the split pipeline).
- Warmup inputs pass temperatures=None: only Model Runner V2 uses the
  scaling specialization, and it registers no sampling warmups (vllm-project#56654).
- Drop the Iterable/Iterator imports only the deleted chunking used.
- The new watermark tests store target logits pre-temperature and scale
  both rows in the reference, matching rejection_sample's divide-on-load
  contract (all callers use temperature 1.0, so outcomes are unchanged).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: jyan <r_02213@mail.sjtu.edu.cn>
@mergify mergify Bot removed the needs-rebase label Sep 14, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation mrv2 Model Runner V2 specific speculative-decoding

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant