[Bugfix][Spec Decode] Fix prompt_logprobs corruption from CUDA graph output buffer reuse by padded drafter

[hoobnn]

Purpose

FIX #53488

When speculative decoding runs a padded GPU drafter (MTP/EAGLE-style), prompt_logprobs are silently corrupted for requests executed through piecewise CUDA graphs, while the same server without --speculative-config scores every request correctly.

Root cause

In the V1 GPU model runner, the padded GPU drafter runs before _bookkeeping_sync():

target model forward   -> hidden_states backed by a reusable CUDA graph output buffer
sample
padded drafter forward -> reuses/overwrites that graph output buffer
bookkeeping            -> _get_prompt_logprobs_dict() computes prompt logits
                          from the already-overwritten hidden_states

All CUDA graphs share one global memory pool, so the drafter's forward can reclaim the memory backing the target model's graph output. Prompt logprobs are then computed from draft-model output instead of the target hidden states.

This exactly explains the sharp boundary reported in #53488:

• Prompts at or below the CUDA graph capture limit (PIECEWISE) are corrupted.
• Prompts immediately above the limit (cudagraph_mode NONE, eager output) are correct.
• Generation is unaffected, since sampling consumes logits before the drafter runs.

I reproduced this on different hardware from the reporter (RTX 5880 Ada, x86_64, vs. their aarch64 GB10) with one of their exact checkpoints, confirming DEBUG-log cudagraph_mode flips PIECEWISE→NONE precisely at the corruption boundary (512→513 tokens with --max-cudagraph-capture-size 512).

Fix

Add _get_bookkeeping_hidden_states(): only when (1) a GPU model drafter will run before bookkeeping and (2) the current batch contains a request still needing prompt logprobs, snapshot the unpadded scheduled slice of the target hidden states (hidden_states[:total_num_scheduled_tokens].clone()) before the drafter is enqueued, and pass the snapshot to _bookkeeping_sync().

All other paths are zero-cost and unchanged:

• Requests without prompt logprobs incur no clone.
• No speculative decoding: unchanged.
• CPU / post-bookkeeping drafters and GPU n-gram drafting: unchanged.
• The padded drafter fast path itself is untouched.

Not a duplicate

Draft PR #53506 is explicitly "not a fix": it adds scheduler chunk-accounting test coverage for multi-module MTP prefill lookahead and does not touch the GPU model runner, prompt-logprob computation, or CUDA graph buffer lifetime. The reporter's checkpoint has a single MTP layer (use_multi_module_mtp() is False), so it does not take that path. No other open PR fixes #53488.

Test Plan

Unit (new regression tests, red before the fix / green after):

pytest tests/v1/worker/test_gpu_model_runner.py -k bookkeeping_hidden_states -v

Static checks:

pre-commit run --files vllm/v1/worker/gpu_model_runner.py tests/v1/worker/test_gpu_model_runner.py
pre-commit run mypy-3.12 --hook-stage manual --files vllm/v1/worker/gpu_model_runner.py

End-to-end: cyankiwi/Qwen3.8-27B-AWQ-BF16-INT4 (one of the issue's exact checkpoints) on RTX 5880 Ada. Plain vs. MTP ({"method":"mtp","num_speculative_tokens":3}) servers with identical flags (--max-model-len 32768 --no-enable-prefix-caching --enable-chunked-prefill --max-num-batched-tokens 2048 --max-num-scheduled-tokens 2048 --max-cudagraph-capture-size 512). 18 deterministic prompts spanning 448–1714 tokens scored via /v1/completions with "max_tokens": 1, "temperature": 0, "prompt_logprobs": 0, strictly verifying identical prompt token IDs and greedy output token IDs between servers.

Test Result

Unit: 4 passed, 44 deselected. Pre-commit and mypy-3.12: all applicable hooks passed.

E2E, MTP/plain mean-NLL ratio per prompt length:

Prompt tokens	cudagraph_mode	Before fix	After fix
448	PIECEWISE	6.9254x	1.00098x
457	PIECEWISE	55.5964x	0.99734x
480	PIECEWISE	5.8075x	0.99874x
496	PIECEWISE	6.9857x	1.00104x
511	PIECEWISE	6.9448x	1.00192x
512	PIECEWISE	6.9607x	1.00194x
513	NONE	0.99910x	0.99910x
520	NONE	1.00026x	1.00026x
525	NONE	0.99902x	0.99902x
544	NONE	1.00051x	1.00051x
600	NONE	1.00055x	1.00055x
1714	NONE	0.99892x	0.99892x

Across all 18 prompts the ratio range goes from 0.99892x – 55.59636x (before) to 0.99734x – 1.00194x (after). Plain-server numbers are bit-identical before and after the fix (no regression on the non-speculative path), and greedy generation token IDs remain identical between Plain and MTP servers. The worst case (457 tokens) drops from PPL ≈ 2.2e8 back to 1.41, matching plain.

Reproduction script, server logs, and all 72 raw API responses are available on request.

AI assistance was used for reproduction, root-cause analysis, implementation, and testing. I have reviewed every changed line and take responsibility for the change end-to-end.

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

[hoobnn]

cc @benchislett @MatthewBonanni — would appreciate a review when you have a chance.

This fixes a prompt_logprobs corruption caused by the padded drafter reusing the target model's CUDA graph output buffer: when drafting runs before bookkeeping (introduced with efficient padded speculation, #24539), the drafter can overwrite the hidden states that prompt logprobs bookkeeping still needs to read. The fix clones the relevant slice only when a request in the batch actually has prompt_logprobs enabled, so the common path stays copy-free.

Tagging you since this touches the padded-drafter path from #24539, and it's closely related to the cudagraph-padding buffer issue fixed in #50065.

[nyhhome]

Independent validation of this PR on a second platform. I filed #53488; this fix resolves it here.

Setup

Hardware	NVIDIA Jetson AGX Thor (sm110a), aarch64, 122.8 GiB unified memory, single GPU
Image	vllm/vllm-openai:nightly-aarch64, vLLM 0.26.1rc1.dev1102+ge9d1398d9 — the same build used in the #53488 report
Model	Qwen3.8-27B NVFP4 (RadixArk checkpoint, ships its own MTP head)
Serve flags	--max-model-len 32768 --gpu-memory-utilization 0.58 --kv-cache-dtype fp8 --kv-cache-memory-bytes 8589934592 --max-num-seqs 128 --no-enable-prefix-caching --trust-remote-code
Speculation	--speculative-config '{"method":"mtp","num_speculative_tokens":3}'
CUDA graphs	on (enforce_eager=False, CUDAGraphMode.FULL_AND_PIECEWISE, max_cudagraph_capture_size = 512 on the speculative rows, 256 on the non-speculative row)

Since vllm/v1/worker/gpu_model_runner.py is pure Python, I did not rebuild: the patch applies to the image's own file with zero fuzz (patch -F 0, all four hunks, uniform −2 line offset), and I bind-mounted the result read-only over the in-image path. The control and baseline legs bind-mount the image's own unmodified file at the same path, so the only difference between legs is file content, not the presence of a mount. Verified inside each container:

• control: sha256 c4a375f1…, hasattr(GPUModelRunner, "_get_bookkeeping_hidden_states") → False
• patched: sha256 6692240f…, hasattr(...) → True, and inspect.getsource prints your function verbatim

Three legs, each the first and only vLLM row of its own boot (this box does not fully return unified memory when a container exits), all launched at MemAvailable 105.1–105.2 GiB.

1. The corruption reproduces here, and the 512/513 boundary reproduces exactly

Scoring a fixed 8-chunk corpus with prompt_logprobs=0 (/v1/completions, max_tokens=1, temperature=0). The only chunks whose own input_tokens ≤ 512 are chunk 1 (458 tok) and chunk 2 (481 tok) — and they are exactly the corrupted ones:

chunk	input_tokens	baseline (no spec) mean_nll	unpatched + MTP	patched + MTP
0	1715	1.87396	1.88783	1.86770
1	458	1.28882	11.75770	1.28305
2	481	2.06185	20.66963	2.07117
3	526	1.79109	1.81191	1.78667
4	667	0.94024	0.93227	0.91779
5	600	0.31440	0.31795	0.31967
6	823	0.40277	0.39407	0.40023
7	976	0.42162	0.41300	0.42211
whole corpus	6246	1.16827 (ppl 3.2164)	3.36961 (ppl 29.07)	1.16432 (ppl 3.2037)

Length scan over prefixes of the same corpus, each request standalone, HTTP 200 unless noted:

input_tokens	unpatched	patched
400	2.01776 clean	2.02249 clean
450	11.5084 corrupt	1.95101 clean
480	12.67211 corrupt	1.98547 clean
500	HTTP 400 — see below	1.98641 clean
508	15.87925 corrupt	1.95169 clean
512	15.87800 corrupt	1.94532 clean
513	1.93850 clean	1.94516 clean
516 / 520 / 525 / 560	clean	clean

The 512 corrupt | 513 clean adjacent pair is exactly the same adjacent pair (512 corrupt / 513 clean) I bisected on a different machine with a different checkpoint, so the boundary now has an independent second-platform reproduction. Note the lower edge does not close: 400 tokens is clean on both platforms, so this is not "every short prompt".

2. One symptom worth adding to the issue: it can escalate to HTTP 400

At the 500-token point the unpatched server does not return garbage — it fails outright:

HTTP 400 {"error":{"message":"Out of range float values are not JSON compliant: nan",
                   "type":"BadRequestError","param":null,"code":400}}

The overwritten hidden states can produce NaN logprobs that vLLM's own JSON encoder refuses to serialize. So besides silently wrong numbers, this bug can present as a hard request failure. I hit this twice from independent runs against the same unpatched server — my first scan script died on it before I hardened the scan to record it — so it is repeatable, not a one-off. After the patch, that point returns a normal 1.98641.

3. After the patch the boundary is gone, not moved

polluted_tokens: [] across all 11 scan points. The 512→513 step changes from 15.878 → 1.939 (an 8× discontinuity) to 1.94532 → 1.94516 (0.008%) — the discontinuity is flattened rather than relocated.

To rule out the obvious alternative explanation — that the patched leg simply stopped capturing graphs, or stopped speculating — both legs logged the identical max_cudagraph_capture_size: 512 and the identical cudagraph_capture_sizes list (… 464, 480, 496, 512), so capture still covers the whole corrupted range; and the patched leg ran at 36.07 tok/s against 13.01 tok/s for the non-speculative baseline, so speculation was still active. The fix removes the corruption without removing either mechanism.

4. Speculation is fully preserved; no generation-side regression

	unpatched	patched
spec_decode_num_drafts_total	4533	4531
spec_decode_num_draft_tokens_total	13599	13593
spec_decode_num_accepted_tokens_total	9124	9170
draft acceptance rate	67.09%	67.46%
mean acceptance length	3.013	3.024
single-stream tok/s (code / math / prose, 512 tok, temp 0)	34.75 / 32.17 / 19.90	36.07 / 32.68 / 21.37

The non-speculative baseline leg runs at 13.01 / 13.03 / 13.01 tok/s, so speculation is genuinely active on both speculative legs. The tok/s improvement crosses a reboot and a different thermal state (48.6 W vs 35.9 W mean GPU power), so I would not claim the patch makes anything faster — only that nothing regressed. That is the expected shape: the added hidden_states[:total_num_scheduled_tokens].clone() only fires when a request in the step actually asked for prompt_logprobs, and none of the benchmark requests do.

Caveats, stated plainly

• This validates the fix. It does not identify which knob sets the 512 boundary — in my configuration max_num_scheduled_tokens = 2048/(K+1), max_num_seqs × (K+1) = 128 × 4, and the speculative row's max_cudagraph_capture_size are all exactly 512, and I did not run the discriminating single-variable sweep.
• However, CUDA graphs do look necessary: two earlier rows on this same box with MTP K=3 and --enforce-eager scored completely clean, and --enforce-eager removes CUDA graphs while leaving 2048/(K+1) = 512 intact. That is consistent with the mechanism described in this PR.
• VLLM_GDN_DECODE_KERNEL=triton was set on all three legs. It is an unrelated workaround for [Bug]: GDN MTP fused decode kernel (fused_gdn_decode_post_conv_mtp) crashes with "no kernel image is available" on SM110a (Jetson Thor) — capability guard checks symbol presence, not cubin arch #53462 (no-kernel-image for the fused GDN decode kernel on sm110a); without it the server does not start on this hardware. It is constant across legs and therefore not a confounder here.
• Single GPU, single checkpoint, MTP only — I did not test EAGLE/EAGLE3 or tensor parallelism.
• Cross-boot greedy token-level determinism does not hold on this box (even a plain-vs-plain self-comparison diverges), so I deliberately did not use token-identity as an acceptance criterion; the generation-side evidence above is throughput, acceptance counters, and a numeric canary.

LGTM from my side — this fixes the issue I reported, on hardware and a checkpoint different from the ones the PR was developed against.

Disclosure: this validation was run and written up with AI assistance; the hardware, the runs, and the numbers are real and reproducible, and the account owner reviewed the result before posting.

[nyhhome]

Correction to my validation comment above, plus the discriminating run I said I had not done.

1. One caveat item was wrong. I wrote that max_num_scheduled_tokens = 2048/(K+1). It is 2048. The engine logs it explicitly —

WARNING [vllm.py:1931] max_num_scheduled_tokens is set to 2048 based on the speculative decoding settings.

— and _set_max_num_scheduled_tokens contains no division. So the 512 confluence in my configuration was two-fold, not three-fold: max_num_seqs × (K+1) = 128 × 4, and the speculative row's max_cudagraph_capture_size. A consequence worth flagging for anyone who tried to follow my caveat: --max-num-batched-tokens is not a discriminating knob (512 / 2048 / 8192 all produce the same capture list).

2. The discriminating single-variable run has now been done, and the boundary follows the CUDA graph capture limit.

Same box (Jetson AGX Thor, sm110a), same image (0.26.1rc1.dev1102+ge9d1398d9), same unpatched gpu_model_runner.py, same --max-num-seqs 128 and num_speculative_tokens=3, CUDA graphs on. The only changed knob is an explicit --compilation-config '{"max_cudagraph_capture_size":256}':

prompt tokens	mean NLL	verdict
200	13.544	polluted
256	13.577	polluted
257	1.990	clean
300	2.089	clean
480	1.985	clean
512	1.945	clean
513	1.936	clean

Boundary: 256 polluted / 257 clean, against 512 / 513 on the otherwise identical row with the default capture size of 512.

max_num_seqs × (K+1) was unchanged at 128 × 4 = 512 across both rows, and 512 tokens is now clean. That excludes it and leaves the capture limit — the mechanism this PR describes. Startup assertions on this row: 'max_cudagraph_capture_size': 256, capture list tail …240, 248, 256, enforce_eager=False, V1 model runner confirmed in the logs.

Two honest notes on this run: the corpus and scoring request are bit-identical to my earlier rows, but I had to lower --gpu-memory-utilization from 0.58 to 0.40 because another process on the box was holding memory; the KV pool is pinned by --kv-cache-memory-bytes and both capture size and max_num_seqs are explicit, so this should not bear on the boundary — and the boundary landing exactly on the new capture size is itself the strongest evidence. Second, these are max_tokens=1 prompt-scoring requests, so the spec-decode acceptance counters stay at zero on all rows, including the baseline.

So this correction strengthens the conclusion rather than weakening it: my --enforce-eager rows were clean while max_num_seqs × (K+1) = 512 was unchanged, and now moving only the capture limit moves the boundary with it.

3. On the force-push. I re-checked head 78e87980be8d against the revision I tested: the patch is byte-for-byte identical (hunk line offsets only), so the validation result carries over unchanged to the current head.

Separately, and from reading the config code only — I have not measured this: #53183 ("Use MRV2 for all models by default") landed between this PR's old and new base and removed the hybrid/architecture gate in _is_default_v2_model_runner_model. On the new base a hybrid + MTP model like the one I tested defaults to the V2 model runner, where vllm/v1/worker/gpu_model_runner.py is not the runner serving the request. Worth noting that this does not narrow what the PR fixes: V1 remains the runner wherever MRV2 cannot serve, and #53183 landed on 2026-08-27, after v0.28.0 was published on 2026-08-26 — so on every released build a hybrid + MTP model still defaults to V1, and the bug this PR fixes is live there regardless of what main now defaults to. Practically: anyone re-testing this PR on the current base probably wants VLLM_USE_V2_MODEL_RUNNER=0 pinned on both legs, or they will silently measure a path this patch does not touch.

Disclosure: this analysis and the runs behind it were done with AI assistance; the hardware, the runs and the numbers are real and reproducible, and the account owner reviewed the content before posting. Raw artifacts available on request.

[nyhhome]

One more data point for the reviewers, from the first nightly after #55375 landed — vllm/vllm-openai:nightly-aarch64 built 2026-09-06 (VLLM_BUILD_COMMIT=1970f3ed, 0.28.1rc1.dev451+g1970f3ed4), this PR not applied. Same Jetson AGX Thor, same corpus, scoring request and flags as my 08-31 row above (--max-num-seqs 128, MTP K=3, CUDA graphs on, default max_cudagraph_capture_size = 512 with the same 51-entry capture list, VLLM_GDN_DECODE_KERNEL=triton, GMU 0.40, KV pinned). The only knob varied is the model runner:

model runner	15 prompt lengths, 200 → 1100 tokens
VLLM_USE_V2_MODEL_RUNNER=0 (V1, gpu_model_runner.py — the runner this PR patches)	polluted at 200 / 256 / 400 / 480 / 508 / 512, clean from 513 up — the 512 | 513 boundary is unchanged
default on this nightly (Using V2 Model Runner)	clean at every length, mean NLL 1.90–2.03
fully stock: default runner + the engine-chosen cuda GDN decode kernel (#53835)	clean at every length, mean NLL 1.89–2.03

Two things follow. The corruption is still there on the V1 path on current main, so this fix is still needed for anyone who pins VLLM_USE_V2_MODEL_RUNNER=0 (or lands on V1 through one of the fallbacks). And the V2 runner does not reproduce it with the same drafter and capture list, which is consistent with the mechanism this PR describes living in the V1 runner's padded-drafter / output-buffer path rather than in the kernels or the scheduler. For the record, #55375 (the Mamba state-index stride fix discussed on #53488) changed nothing here, as expected for this model.

Full context, including how this relates to the Mamba state-index stride class, is in my comment on #53488: #53488 (comment)

Disclosure: the runs and the write-up were done with AI assistance; the hardware, the runs and the numbers are real and reproducible; the account owner reviewed the content before posting.