Conversation
Contributor
|
Documentation preview: https://vllm--53630.org.readthedocs.build/en/53630/ |
jyan-R
marked this pull request as ready for review
August 24, 2026 21:22
jyan-R
requested review from
22quinn,
NickLucche,
houseroad and
yewentao256
as code owners
August 24, 2026 21:22
Copilot stopped reviewing on behalf of
jyan-R due to an error
August 24, 2026 21:42
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
jyan-R
force-pushed
the
claude/pr2-remove-fp32-verifier-copy
branch
from
August 25, 2026 14:36
16aa5dd to
55bae3f
Compare
`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
force-pushed
the
claude/pr2-remove-fp32-verifier-copy
branch
from
August 25, 2026 17:37
55bae3f to
4855de2
Compare
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
…2-verifier-copy # Conflicts: # tests/v1/spec_decode/test_rejection_sampler_utils.py Signed-off-by: jyan <r_02213@mail.sjtu.edu.cn>
…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>
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
…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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
processed-logprob scorer load the logits; it is never written back to the stored
logits;
computed in FP32 and written back to the head dtype with a single rounding; exact mask
processors keep their original semantics;
overwrite the input; processed reporting reads the processed stored logits;
are unchanged; only the hard cap that existed to serve the FP32 chunk is removed;
head-dtype target logits make the previously latent BF16 accumulation path reachable;
min_p, and customlogprob_token_idsreporting 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.
-inf-infmasks-infmasksRaw 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=1results underNumerical 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:
FP32 registers at every load site and never touches storage;
-infmasks — keptlogits are bit-identical to what the LM head produced;
guardrailed under Numerical fidelity and model quality below;
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
processed logprobs still fall back to the FP32 canvas, so it cannot delete the
chunking or the adaptive cap. This PR covers every currently legal
processor/reporting combination.
divide by temperature on load" contract on the draft side; this PR aligns the target
verifier to the same contract.
max_total_logitscap exists only to guarantee the FP32 verifier landsin a single safe chunk; this PR removes that premise and does not change the adaptive
verification algorithm.
still falls back to Model Runner V1, so the custom
logprob_token_idse2e requesttemporarily covers Eagle only.
adaptive-verification output-logprobs restriction. This PR is rebased on top of it:
the device-side
cu_num_generated_tokensboundaries it introduced are kept and flowthrough the head-dtype scorer unchanged, and removing the chunked path also removes
the stale-layout hazard its
_verify_in_chunksassert defended against. If[Bugfix][Spec Decode][Structured Output] Drive grammar masks from GPU logit counts #52477 lands first, the grammar-mapping tests will need a rebase update.
materialize avoidable full-vocabulary temporaries; the FP32 canvas this PR deletes is
exactly such a temporary.
Test Plan
thinking budget, top-k/top-p, rejection sampling, adaptive verification, and the
logprob numerical oracles and head-dtype penalty path.
test_spec_decode_logprobsmatrix across 3 speculative model setups and all4
LogprobsModevalues, comparing emitted tokens and raw/processed logprob outputs with thenon-speculative reference.
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.
RejectionSampler.__call__with BF16 logits,K=7,V=129280, andR={8,256,512}across temperature, top-p, penalties, combined processors, and logprobreporting. Compare p50 latency and per-call CUDA peak allocation with the FP32-canvas control.
penalties-only, stochastic top-p with penalties, and
top_k=1; compare accuracy, invalidoutputs, and acceptance length.
Repository test commands
.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
231 passed, 2 skipped(reported from the run; that consolelog was not kept). It covers every supported sampling processor, raw/processed reporting,
adaptive verification, thinking budget and the head-dtype penalty path.
test_spec_decode_logprobsmatrix passed all 3 model setups × 4LogprobsModecombinations (12 passed). Spec and non-spec runs returned identical emittedtoken IDs, logprob token IDs and ranks. Raw scores were bitwise equal; the largest processed-score
difference was
0.08334.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 activesampling wave produced draft tokens, so the test covers both the control path and the new
head-dtype verifier in the same distributed engine.
ruff checkandruff format --checkpassed 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=129280andR=512(4096 verification rows), whichcrosses 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.
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-6and2.13e-4against the FP32-inputreference, 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:
temperature=0.7, top_p=0.9; penalties ontop_k=1temperature=0.7, top_p=0.9; penalties onThe 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
supported_models.mdandexamplesfor 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)