[Bugfix][Spec Decode] Trim the optimistic spec-decode tokens on every pipeline rank

[Peuqui]

Purpose

Under async scheduling with speculative decoding, _update_states optimistically
assumes every draft token was accepted and appends that many -1 placeholders to
req_state.output_token_ids:

optimistic_num_accepted = req_state.prev_num_draft_len
req_state.output_token_ids.extend([-1] * optimistic_num_accepted)

That extend is not gated on the pipeline rank. The trim that undoes it was:

if not is_last_rank:
    ...
elif num_output_tokens < len(req_state.output_token_ids):
    del req_state.output_token_ids[num_output_tokens:]

so it only ever ran on the final pipeline rank. The deferred correction
(correct_spec_decode_token_counts) adjusts num_computed_tokens and
num_tokens_no_spec, but it never removes the placeholders — the trim is the
only thing that does. On every non-final rank the cached output therefore grew
by the draft length each speculative step and never shrank, so the cached state
drifted away from the scheduler's view of the same request.

The fix turns the elif into an if. Single-rank behaviour is unchanged:
there is_last_rank is True, the preceding branch is not taken, and the trim
ran before this change too.

prev_num_draft_len is assigned in InputBatch.update_req_spec_token_ids,
which runs on every rank, so the growth is not specific to any model or backend.

Test Plan

Environment: checkout at origin/main 4f19ef7 with the compiled extensions of a
1Cat-vLLM 1.5.0 wheel linked in (see Limitations).

python -m pytest tests/v1/worker/test_gpu_model_runner.py -k trims_optimistic -q
python -m pytest tests/v1/worker/test_gpu_model_runner.py -q
python -m pytest tests/v1/worker/test_gpu_model_runner_pp.py -q

# counter-check: gpu_model_runner.py reverted to origin/main, test kept
python -m pytest tests/v1/worker/test_gpu_model_runner.py -q

pre-commit run --files vllm/v1/worker/gpu_model_runner.py \
    tests/v1/worker/test_gpu_model_runner.py
pre-commit run mypy-3.10 --hook-stage manual --files <same two files>

Test Result

New test alone: 1 passed.

tests/v1/worker/test_gpu_model_runner.py: 39 passed, 2 failed, then the
process aborts during cleanup_dist_env_and_memory. With
vllm/v1/worker/gpu_model_runner.py reverted to origin/main and the new test
kept, the same command gives 38 passed, 3 failed and the identical abort —
the one extra failure is the new test:

assert len(non_last_state.output_token_ids) == len(last_state.output_token_ids)
AssertionError: assert 4 == 1
 +  where 4 = len([111, -1, -1, -1])
 +  and   1 = len([111])

The two failures present on both trees are
test_non_last_pp_rank_profiles_with_speculative_config[ngram] and
[draft_model]. They are pre-existing and unrelated: both die inside
FlashInfer with

RuntimeError: Check failed: (status == cudaSuccess) is false:
    BatchPrefillWithPagedKVCache failed with error unspecified launch failure
torch.AcceleratorError: CUDA error: unspecified launch failure

i.e. the FlashInfer paged-prefill kernel does not launch on this bench's GPUs
(Turing/Volta). The corrupted CUDA context is also what aborts the interpreter
at teardown, which is why pytest never prints its own failure report for them.

tests/v1/worker/test_gpu_model_runner_pp.py: 11 passed.

The existing PP consistency tests
test_update_states_pp_non_async_multi_request_keeps_token_buffers_consistent
and test_update_states_pp_async_multi_request_keeps_rank_state_consistent
pass both before and after, so the newly reachable trim does not disturb the
non-async hand-off path.

pre-commit over both files: ruff check, ruff format, typos, mypy-local, SPDX
headers, root lazy imports, forbidden imports, the torch.cuda-call check,
config-docstring check, attention-backend docs and the boolean-ops check all
Passed. pre-commit run mypy-3.10 --hook-stage manual: Passed.

Not a duplicate

Checked on 2026-09-08 against every open PR. Three open PRs touch
vllm/v1/worker/gpu_model_runner.py — #241, #239 and #235 — and none of their
hunks in that file contain output_token_ids, is_last_rank,
num_output_tokens or _update_states. No open issue mentions the optimistic
spec-decode extend or the output trim.

Why it matters in practice

On our own deployment (2x RTX 8000 + 2x V100, TP2 x PP2, MTP) the unbounded
overshoot is not cosmetic. The inflated length feeds the discard/chunked-prefill
bookkeeping; after enough speculative rounds the discard mask flips on the
non-final rank only, the PP broadcast guard then disagrees between ranks — the
final rank keeps sending sampled-token broadcasts while the first rank skips
them — and the run wedges in NCCL. That is what sent us looking for this in the
first place; the unit test above is the reduced, deterministic form of it.

Already reported upstream

The same elif is present in upstream vLLM (vllm/v1/worker/gpu_model_runner.py
at 9e905f7, 2026-09-01), so this is not a fork-specific regression. I reported it
there on 2026-08-28 as vllm-project/vllm#54260, including the failure chain
above; that issue is still open and unanswered. I am opening the fix here
because this is the tree I run and can test against.

Limitations

Unit tests and linting run against this tree. I have not booted current main
from source on this hardware — that is a multi-hour CUDA build — so there is no
end-to-end throughput number attached to this change. The defect and the fix are
demonstrated by the unit test above, which fails on the unmodified tree and
passes with the change.

AI assistance: this change was developed with Claude (Anthropic) as a coding
assistant. Every changed line was reviewed by me and the test runs above were
executed on my hardware; I can defend the change end to end.

🤖 Generated with Claude Code