[sawsa307]

Purpose

When a speculative decode step stops on an accepted draft (EOS, a stop token, or max_tokens reached before the bonus token), _update_request_with_output trims the step's later tokens, but num_computed_tokens keeps counting them. Rejections are rolled back in update_from_output; trimmed tokens are not.

Ordinary requests are freed right after the stop, so the count is never read again. A resumable (streaming-input) request continues: _update_request_as_session uses num_computed_tokens as the fold boundary, so the next chunk starts "already computed" over positions whose KV belongs to the trimmed drafts.

Prompt [1, 2, 3], max_tokens=2, k=2. The prefill samples 10; a step with drafts [11, 12] returns [11, 12, 13] (both accepted) and stops at 11. Folding in the next chunk [7, 8]:

	all_token_ids	num_computed_tokens	next step computes
main	[1, 2, 3, 10, 11, 7, 8]	6	8 only; 7 reuses draft 12's KV
this PR	[1, 2, 3, 10, 7, 8]	4	7, 8

The fix uncounts the trimmed tokens, so num_computed_tokens again covers all new tokens but the last kept one. The fold then drops that token exactly as it drops the final sampled token of a non-speculative turn, which is what makes a session's output identical with and without speculative decoding (see the GPU check below).

Under async scheduling the trimmed tokens' placeholders were never released either (num_output_placeholders was decremented by the post-trim count), so they are released too. num_computed_tokens - num_output_placeholders is unchanged, so neither the async cache_blocks frontier nor the C - P frontier from #58259 moves.

This count is also what stale block-hash truncation keys on (#52806, and the del session.block_hashes[num_computed_tokens // hash_block_size:] suggested on #58259): with the overshoot, the cut point lands too late and leaves stale hashes behind.

Diffusion schedulers inherit the change; resumable diffusion requests were not exercised.

Not a duplicate

• [Bugfix] Clamp num_computed_tokens after streaming session rebuild #44526 clamps num_computed_tokens to num_prompt_tokens after the fold. That catches only an overshoot past the whole new chunk (num_new_tokens <= 0); in the example above it is a no-op and 7 is still skipped. This PR fixes the count where it goes wrong.
• [Bugfix] Fix resumable request + async scheduling handoff race #58259 (merged) fixed the async handoff race. Its C - P frontier also covers this case under async scheduling, but not under sync scheduling, where num_output_placeholders is always 0. This PR keeps C - P unchanged, so the two compose.
• [Bugfix][V1] Register accepted spec-decode tokens before terminal cleanup #54979 caps its terminal cache_blocks call at request.num_tokens to work around the same overshoot; the cap stays correct.
• Searched open and closed PRs and issues for streaming input + speculative decoding, num_computed_tokens trimming, and _update_request_as_session; no other PR fixes this.

Test Plan

# CPU
VLLM_TARGET_DEVICE=cpu .venv/bin/python -m pytest \
  tests/v1/streaming_input/test_scheduler_streaming.py \
  tests/v1/core/test_scheduler.py tests/v1/core/test_async_scheduler.py -q
VLLM_TARGET_DEVICE=cpu .venv/bin/python -m pytest tests/v1/core tests/v1/streaming_input -q

# GPU (GB200)
.venv/bin/python -m pytest tests/v1/e2e/general/test_streaming_input.py -v
.venv/bin/python -m pytest tests/v1/e2e/general/test_async_scheduling.py::test_with_ngram_gpu_spec_decoding -v

New test: test_next_chunk_fully_computed_after_stop_on_accepted_draft (scheduler, sync and async scheduling): after the trimmed stop, the fold drops the last kept token and the next chunk is scheduled in full.

GPU check, run manually and not kept in CI: a 6-chunk streaming session (opt-125m, fp32 + FLEX_ATTENTION, greedy, max_tokens=2 per chunk) must produce the same tokens and logprobs with and without ngram speculative decoding (k=4), with drafts accepted. The greedy tokens of this repetitive prompt do not change with the context, so logprobs are what catch the difference.

Test Result

• Scheduler test: the sync case fails on main and passes with the fix. The async case passes on main since [Bugfix] Fix resumable request + async scheduling handoff race #58259 and is kept to check that this PR's C/P accounting still composes with it.
• test_scheduler_streaming.py + test_scheduler.py + test_async_scheduler.py: 221 passed.
• tests/v1/core + tests/v1/streaming_input (CPU): 797 passed. The 1 failure and 6 errors (test_reset_prefix_cache_e2e, test_scheduler_e2e) need a GPU and are identical on main.
• GPU on GB200:
  • Streaming session check, main: diverges from the second chunk on, e.g. logprobs [-0.3593, -0.5455] with ngram vs [-0.3427, -0.7063] without.
  • Streaming session check, this PR: tokens match and logprobs agree to ~1e-5; tests/v1/e2e/general/test_streaming_input.py passes.
  • Async spec-decode regression check (the placeholder accounting change applies to every async spec request, not only streaming): tests/v1/e2e/general/test_async_scheduling.py::test_with_ngram_gpu_spec_decoding passes, all 6 configs (sync/async, preemption, chunked prefill, mp/uni executor) match the non-speculative baseline in outputs and logprobs.
  • ngram is not yet supported by Model Runner V2, so the streaming session check uses the V1 model runner. The fix is in the scheduler and runner-independent. Async scheduling is covered by the scheduler test only.

AI assistance (Claude Code) was used to find, reproduce and fix this. I reviewed every changed line and ran the tests above.

---

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.

🤖 Generated with Claude Code

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

[yzong-rh]

Good catch on this bug! LGTM

A good way understand imo is that, when spec decode requests are retired via update_from_outputs, we need to adjust num_computed_tokens and num_output_placeholders (decrementing both by num_rejected).

vllm/vllm/v1/core/sched/scheduler.py

Lines 2054 to 2066 in 7871963

	num_draft_tokens = len(scheduled_spec_token_ids)
	num_sampled = self.num_sampled_tokens_per_step
	num_accepted = max(len(generated_token_ids) - num_sampled, 0)
	num_rejected = num_draft_tokens - num_accepted
	# Rejections roll back num_computed_tokens (and, under async
	# scheduling, num_output_placeholders, which covers the spec
	# tokens). A stale rejection count predates the preemption
	# rollback and must not apply.
	if not output_is_stale:
	if request.num_computed_tokens > 0:
	request.num_computed_tokens -= num_rejected
	if request.num_output_placeholders > 0:
	request.num_output_placeholders -= num_rejected

However, when spec decode tokens are "trimmed" due to a STOP condition, this was not done. Leading to the issues above. Effectively, we need to

trimmed = num_new_tokens - retained

# discard trimmed positions, was missing!
num_computed_tokens -= trimmed
num_output_placeholders -= trimmed

# resolve retained output placeholders
num_output_placeholders -= retained 

As an optimization, to avoid decrementing num_output_placeholders twice (once by retained, once by trimmed, this PR

num_computed_tokens -= trimmed
num_output_placeholders -= num_new_tokens

NIT1: Not sure the logprob comparison GPU test is necessary in CI

NIT2: Perhaps we should assert on num_computed_tokens being equal to len(all_token_ids) when resumable requests stop (after the _update_request_as_session but before appending a new streaming input chunk)

[sawsa307]

Thanks for the review and the clear write-up of the accounting.

NIT1: agreed, the scheduler test already catches this. I'll drop the e2e test and keep its GPU result in the description.

NIT2: The idea makes sense, but it only holds once both this PR and #58259 are in.
With this PR alone, async + spec decode still folds with the in-flight 1 + k positions counted (real AsyncScheduler, k=2: num_computed_tokens=13, num_tokens=11); with #58259 alone, sync + spec decode still has the trim overshoot (6 vs 5), since num_output_placeholders is 0 there.
Either way the assertion would turn silent corruption into an EngineCore crash. Also, a session that stops while preempted folds with num_computed_tokens == 0 while keeping its tokens (#58776, fixed in #58778), so the check should be num_computed_tokens <= num_tokens at fold entry rather than equality after the truncation. It can be added as a small follow-up once both are merged.

[mergify]

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

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

[mergify]

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

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

[njhill]

/ci run

[github-actions]

✅ Triggered Buildkite CI #95177 for commit a035821b4c38.

[github-actions]

✅ @sawsa307, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.