[Bugfix][Core] Uncount stop-trimmed spec tokens from num_computed_tokens (#58668)

Signed-off-by: David Cheung <d7cheung@gmail.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>

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

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.

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