From d486219fc2f27a981d6eda53eb2b95bfde1a1d63 Mon Sep 17 00:00:00 2001 From: David Cheung Date: Thu, 15 Oct 2026 12:50:06 -0700 Subject: [PATCH] [Bugfix][Core] Uncount stop-trimmed spec tokens from num_computed_tokens (#58668) Signed-off-by: David Cheung Signed-off-by: Nick Hill Co-authored-by: Claude Opus 5.5 Co-authored-by: Nick Hill --- .../test_scheduler_streaming.py | 42 ++++++++++++++++++- vllm/v1/core/sched/async_scheduler.py | 6 ++- vllm/v1/core/sched/scheduler.py | 5 ++- 3 files changed, 49 insertions(+), 4 deletions(-) diff --git a/tests/v1/streaming_input/test_scheduler_streaming.py b/tests/v1/streaming_input/test_scheduler_streaming.py index 1f97d8a1c0..c6df585a58 100644 --- a/tests/v1/streaming_input/test_scheduler_streaming.py +++ b/tests/v1/streaming_input/test_scheduler_streaming.py @@ -5,9 +5,11 @@ import unittest from unittest.mock import MagicMock import numpy as np +import pytest import torch import vllm.v1.core.kv_cache_utils as kv_cache_utils +from tests.v1.core.utils import create_scheduler as create_spec_scheduler from vllm.config import DeviceConfig, VllmConfig from vllm.multimodal.inputs import ( MultiModalFeatureSpec, @@ -28,7 +30,7 @@ from vllm.v1.kv_cache_interface import ( KVCacheConfig, KVCacheGroupSpec, ) -from vllm.v1.outputs import LogprobsLists, ModelRunnerOutput +from vllm.v1.outputs import DraftTokenIds, LogprobsLists, ModelRunnerOutput from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm.v1.structured_output import StructuredOutputManager @@ -676,3 +678,41 @@ class TestStreamingScheduler(unittest.TestCase): cached_state_cycle1["prompt_token_ids"] is not cached_state_cycle3["prompt_token_ids"] ), "Cached states from different cycles should be independent objects." + + +def _model_output(request_id: str, token_ids: list[int]) -> ModelRunnerOutput: + return ModelRunnerOutput( + req_ids=[request_id], + req_id_to_index={request_id: 0}, + sampled_token_ids=[token_ids], + logprobs=None, + prompt_logprobs_dict={request_id: None}, + pooler_output=[], + ) + + +@pytest.mark.parametrize("async_scheduling", [False, True]) +def test_next_chunk_fully_computed_after_stop_on_accepted_draft(async_scheduling): + """A stop on an accepted draft trims the step's later tokens. The next + chunk must still be computed in full rather than landing on positions + whose KV belongs to the trimmed drafts.""" + scheduler = create_spec_scheduler( + num_speculative_tokens=2, + async_scheduling=async_scheduling, + speculative_method="ngram_gpu" if async_scheduling else None, + ) + session = DummyRequest("session", prompt_token_ids=[1, 2, 3], max_tokens=2) + scheduler.add_request(session) + scheduler.update_from_output(scheduler.schedule(), _model_output("session", [10])) + scheduler.update_draft_token_ids(DraftTokenIds(["session"], [[11, 12]])) + # Both drafts are accepted, but max_tokens stops at 11. + scheduler.update_from_output( + scheduler.schedule(), _model_output("session", [11, 12, 13]) + ) + assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ + + scheduler.add_request(DummyRequest("session", prompt_token_ids=[7, 8])) + output = scheduler.schedule() + + assert list(session.prompt_token_ids) == [1, 2, 3, 10, 7, 8] + assert output.num_scheduled_tokens["session"] == 2 diff --git a/vllm/v1/core/sched/async_scheduler.py b/vllm/v1/core/sched/async_scheduler.py index 290204d58c..0b3ab67b0f 100644 --- a/vllm/v1/core/sched/async_scheduler.py +++ b/vllm/v1/core/sched/async_scheduler.py @@ -60,14 +60,16 @@ class AsyncScheduler(Scheduler): self, request: Request, new_token_ids: list[int], is_stale: bool = False ) -> tuple[list[int], bool]: status_before_update = request.status + # Tokens trimmed at a stop resolve their placeholders too. + num_new_tokens = len(new_token_ids) new_token_ids, stopped = super()._update_request_with_output( - request, new_token_ids + request, new_token_ids, is_stale ) # Placeholders were zeroed at preemption; a stale delivery must not # decrement them (it would underflow). if not is_stale: - request.num_output_placeholders -= len(new_token_ids) + request.num_output_placeholders -= num_new_tokens assert request.num_output_placeholders >= 0 # Cache the new tokens. Preempted requests should be skipped. diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index c7ece2a1f0..c7619a2f66 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -2574,7 +2574,6 @@ class Scheduler(SchedulerInterface): def _update_request_with_output( self, request: Request, new_token_ids: list[int], is_stale: bool = False ) -> tuple[list[int], bool]: - # is_stale is only used by the AsyncScheduler override. # Append generated tokens and check for stop. Note that if # a request is still being prefilled, we expect the model runner # to return empty token ids for the request. @@ -2586,6 +2585,10 @@ class Scheduler(SchedulerInterface): # This must be called before we make the EngineCoreOutput. stopped = check_stop(request, self.max_model_len) if stopped: + # num_computed_tokens covers all new tokens but the last one. + # Keep that true after trimming: resumable requests resume from it. + if not is_stale: + request.num_computed_tokens -= len(new_token_ids) - num_new del new_token_ids[num_new:] # Trim new tokens if needed. break return new_token_ids, stopped