[BugFix] Fix num_output_placeholders preemption underflow (#48245)

Signed-off-by: Chris Eastwood <chris.eastwood@pwn4g3.dev>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Chris Eastwood <chris.eastwood@pwn4g3.dev>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>

tests/v1/core/test_async_scheduler.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from collections import deque
+from collections import defaultdict, deque
+from collections.abc import Callable
 from unittest.mock import Mock
 
 import pytest
@@ -329,46 +330,311 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
     assert not scheduler.running
 
 
-def test_no_placeholder_underflow_on_discarded_spec_frame():
-    num_spec = 5
+class PipelinedEngine:
+    """Drive a real AsyncScheduler like EngineCore.step_with_batch_queue:
+    schedule until the batch queue is full, then process the oldest step's
+    output. Async PP runs pp_size+1 concurrent batches, so up to pp_size
+    steps are in flight at each schedule() call -- the window in which
+    preemption must handle output that has not yet returned. (Single-GPU e2e
+    tests can never create this window: at PP=1, exactly one step is in
+    flight and it is processed before a preempted request can resume.)
+
+    The model runner is emulated with the V2 runner's own bookkeeping, from
+    only what the scheduler serializes to it: slots flushed on
+    preempted_req_ids, resumed requests re-added from the NewRequestData
+    snapshot, sampling when a step reaches the end of the runner's own view
+    of the sequence. This makes preemption races observable: a stale token
+    delivered after a resume is scheduled extends the scheduler's sequence
+    but not the runner's.
+
+    Every sample emits a globally unique token tagged with its sampled
+    position, so tests can assert exact delivery.
+    """
+
+    def __init__(
+        self,
+        scheduler: AsyncScheduler,
+        queue_size: int,
+        accept_drafts: Callable[[int, str, int], int] | None = None,
+    ):
+        self.scheduler = scheduler
+        self.queue_size = queue_size
+        self.accept_drafts = accept_drafts
+        # In-flight steps: (scheduler_output, new_reqs snapshot) in FIFO order.
+        self.queue: deque[tuple[SchedulerOutput, list[tuple[str, int, int]]]] = deque()
+        # Runner-side request state: req_id -> [seq_len, num_computed] as the
+        # runner sees them (its own sampled tokens, not the scheduler's).
+        self.runner_view: dict[str, list[int]] = {}
+        # All tokens the fake runner ever sampled, per request, in order.
+        self.emitted: dict[str, list[int]] = defaultdict(list)
+        # Sequence position each (globally unique) token was sampled for.
+        self.emitted_position: dict[int, int] = {}
+        self.step_idx = 0
+        self._next_token = 1000
+
+    def _schedule(self) -> bool:
+        scheduler_output = self.scheduler.schedule()
+        self.step_idx += 1
+        # Snapshot what NewRequestData serializes at schedule time (both new
+        # and resumed requests for the V2 runner).
+        new_reqs = [
+            (r.req_id, len(r.prefill_token_ids), r.num_computed_tokens)
+            for r in scheduler_output.scheduled_new_reqs
+        ]
+        # Enqueue empty steps too (the engine executes them), so the runner
+        # still observes their preempted/finished request ids in step order.
+        self.queue.appendleft((scheduler_output, new_reqs))
+        return True
+
+    def _process_oldest_step(self) -> None:
+        scheduler_output, new_reqs = self.queue.pop()
+        # Worker-side state updates, in step order: flush preempted/finished
+        # slots, then (re-)add new/resumed requests.
+        for req_id in scheduler_output.preempted_req_ids or ():
+            self.runner_view.pop(req_id, None)
+        for req_id in scheduler_output.finished_req_ids or ():
+            self.runner_view.pop(req_id, None)
+        for req_id, seq_len, num_computed in new_reqs:
+            self.runner_view[req_id] = [seq_len, num_computed]
+
+        req_ids = list(scheduler_output.num_scheduled_tokens.keys())
+        sampled_token_ids: list[list[int]] = []
+        for req_id in req_ids:
+            num_scheduled = scheduler_output.num_scheduled_tokens[req_id]
+            view = self.runner_view.get(req_id)
+            if view is None:
+                # Slot already flushed (request finished/aborted mid-flight).
+                sampled_token_ids.append([])
+                continue
+            seq_len, num_computed = view
+            end = num_computed + num_scheduled
+            if end < seq_len:
+                # Partial prefill by the runner's own bookkeeping: no sample.
+                view[1] = end
+                sampled_token_ids.append([])
+                continue
+            drafts = scheduler_output.scheduled_spec_decode_tokens.get(req_id, ())
+            num_accepted = (
+                min(self.accept_drafts(self.step_idx, req_id, len(drafts)), len(drafts))
+                if drafts and self.accept_drafts
+                else 0
+            )
+            num_rejected = len(drafts) - num_accepted
+            tokens = list(range(self._next_token, self._next_token + 1 + num_accepted))
+            self._next_token += 1 + num_accepted
+            self.emitted[req_id].extend(tokens)
+            sampled_token_ids.append(tokens)
+            # Rejected drafts roll back computed; the sampled tokens extend
+            # the runner's sequence.
+            view[1] = end - num_rejected
+            view[0] = view[1] + 1
+            for offset, token in enumerate(tokens):
+                self.emitted_position[token] = view[0] - len(tokens) + offset
+        model_runner_output = ModelRunnerOutput(
+            req_ids=req_ids,
+            req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)},
+            sampled_token_ids=sampled_token_ids,
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=[],
+        )
+        self.scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    def run(
+        self,
+        max_steps: int = 2000,
+        before_step: Callable[[int, "PipelinedEngine"], None] | None = None,
+    ) -> None:
+        for i in range(max_steps):
+            if not self.scheduler.has_requests() and not self.queue:
+                return
+            if before_step is not None:
+                before_step(i, self)
+            scheduled = (
+                self.scheduler.has_requests()
+                and len(self.queue) < self.queue_size
+                and self._schedule()
+            )
+            if scheduled and len(self.queue) < self.queue_size:
+                # Queue not yet full: the engine returns without blocking.
+                continue
+            if self.queue:
+                self._process_oldest_step()
+        raise AssertionError("engine loop did not converge")
+
+
+def _create_async_pp_scheduler(
+    num_spec: int, pp_size: int = 3, num_blocks: int = 5
+) -> AsyncScheduler:
     scheduler = create_scheduler(
         async_scheduling=True,
-        num_speculative_tokens=num_spec,
-        speculative_method="ngram_gpu",
+        num_speculative_tokens=num_spec or None,
+        speculative_method="ngram_gpu" if num_spec else None,
+        use_v2_model_runner=True,
+        num_blocks=num_blocks,
+        block_size=16,
+        max_num_batched_tokens=512,
     )
-    req = create_requests(num_requests=1, max_tokens=20)[0]
-    req.num_computed_tokens = req.num_tokens
-    scheduler.requests[req.request_id] = req
-    scheduler.running.append(req)
-    req.status = RequestStatus.RUNNING
+    # Emulate PP at the scheduler level; constructing with
+    # pipeline_parallel_size>1 requires that many visible GPUs. Drive with
+    # queue_size=pp_size+1 (V2 async PP runs pp_size+1 concurrent batches).
+    scheduler.pp_size = pp_size
+    scheduler.use_pp = pp_size > 1
+    return scheduler
 
-    req.num_output_placeholders = 1
-    req.async_tokens_to_discard = num_spec
-    computed_before = req.num_computed_tokens
 
-    scheduler_output = SchedulerOutput(
-        scheduled_new_reqs=[],
-        scheduled_cached_reqs=CachedRequestData.make_empty(),
-        num_scheduled_tokens={req.request_id: num_spec + 1},
-        total_num_scheduled_tokens=num_spec + 1,
-        scheduled_encoder_inputs={},
-        scheduled_spec_decode_tokens={req.request_id: [10] * num_spec},
-        num_common_prefix_blocks=[],
-        finished_req_ids=set(),
-        free_encoder_mm_hashes=[],
+def _assert_ordered_subset(delivered: list[int], emitted: list[int]) -> None:
+    """Delivered tokens must be an order-preserving subset of the emitted
+    tokens with no duplicates (tokens are globally unique)."""
+    it = iter(emitted)
+    for token in delivered:
+        assert token in it, f"token {token} delivered out of order or twice"
+
+
+def _assert_positions_consistent(req, engine: PipelinedEngine) -> None:
+    """The i-th delivered output token must be one the runner sampled for
+    exactly sequence position prompt_len + i: catches a preempted request's
+    stale output landing on a position the resumed request resampled (or
+    vice versa), which token-stream equality alone cannot see."""
+    for i, token in enumerate(req.output_token_ids):
+        expected = req.num_prompt_tokens + i
+        actual = engine.emitted_position[token]
+        assert actual == expected, (
+            f"output {i} of {req.request_id}: token sampled for position "
+            f"{actual}, delivered as position {expected}"
+        )
+
+
+@pytest.mark.parametrize("num_spec", [0, 3])
+def test_kv_pressure_preemption_with_inflight_output(num_spec: int):
+    """KV-pressure preemption of requests with in-flight async output.
+
+    PP=3 + async scheduling (batch queue of 4), a block pool small enough
+    that decodes contend and preempt mid-flight, and staggered arrivals so
+    the batch queue actually pipelines. A preempted request's in-flight steps
+    still return: their tokens must be delivered exactly once, their stale
+    spec-rejection counts must not corrupt the rolled-back counters, and the
+    resume must not resample a position that output later delivers.
+
+    Regression for the num_output_placeholders underflow EngineCore crash:
+    with the fix reverted, the num_spec=3 variant fails with exactly
+    ``assert request.num_output_placeholders >= 0`` when a stale spec output
+    returns after the preempted request was resumed and sampled.
+    """
+    max_tokens = 24
+    scheduler = _create_async_pp_scheduler(num_spec)
+    requests = create_requests(
+        num_requests=8, num_tokens=8, max_tokens=max_tokens, ignore_eos=True
     )
-    model_runner_output = ModelRunnerOutput(
-        req_ids=[req.request_id],
-        req_id_to_index={req.request_id: 0},
-        sampled_token_ids=[[999]],
-        logprobs=None,
-        prompt_logprobs_dict={},
-        pooler_output=[],
+    pending = list(requests)
+    for _ in range(2):
+        scheduler.add_request(pending.pop(0))
+
+    # Observe that the scenario under test actually occurs.
+    preempts_with_inflight_output = 0
+    orig_preempt = scheduler._preempt_request
+
+    def counting_preempt(request, timestamp, **kwargs):
+        nonlocal preempts_with_inflight_output
+        if request.num_in_flight_tokens > 0:
+            preempts_with_inflight_output += 1
+        return orig_preempt(request, timestamp, **kwargs)
+
+    scheduler._preempt_request = counting_preempt
+
+    def add_requests(step: int, engine: PipelinedEngine):
+        if pending:
+            scheduler.add_request(pending.pop(0))
+
+    engine = PipelinedEngine(
+        scheduler,
+        queue_size=4,
+        # Deterministically vary spec acceptance so stale outputs carry
+        # nonzero rejection counts.
+        accept_drafts=lambda step, req_id, n: (step + int(req_id)) % (n + 1),
     )
+    engine.run(before_step=add_requests)
 
-    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert preempts_with_inflight_output > 0, (
+        "test did not exercise preemption with in-flight output"
+    )
+    for req in requests:
+        assert req.is_finished()
+        assert req.num_output_tokens == max_tokens
+        # Lossless: delivered tokens are exactly the sampled tokens, in order
+        # (the excluded tail was emitted after the request finished).
+        emitted = engine.emitted[req.request_id]
+        assert list(req.output_token_ids) == emitted[:max_tokens]
+        _assert_positions_consistent(req, engine)
 
-    assert req.num_output_placeholders == 1
-    assert req.num_computed_tokens == computed_before
-    assert req.async_tokens_to_discard == num_spec - 1
-    assert req.status == RequestStatus.RUNNING
+
+@pytest.mark.parametrize("pp_size", [1, 3])
+def test_reset_prefix_cache_with_inflight_output_under_kv_pressure(pp_size: int):
+    """reset_prefix_cache(reset_running_requests=True) resumes requests in
+    the same step it preempts them, so in-flight output must be dropped (the
+    resume resamples those positions).
+
+    pp_size=1: regression for the frame-based discard this fix replaces,
+    which with spec decode drained one *token* count per output frame and
+    over-discarded, corrupting the fresh frames after the resume.
+    pp_size=3: back-to-back resets, so the second re-preempts requests whose
+    dropped stale share is still in flight -- it must be recorded once (not
+    accumulated) and stay dropped.
+    """
+    max_tokens = 24
+    scheduler = _create_async_pp_scheduler(num_spec=3, pp_size=pp_size)
+    requests = create_requests(
+        num_requests=8, num_tokens=8, max_tokens=max_tokens, ignore_eos=True
+    )
+    pending = list(requests)
+    for _ in range(2):
+        scheduler.add_request(pending.pop(0))
+
+    # Observe re-preemptions with an undrained stale share (the
+    # double-count hazard).
+    repreempts_with_stale = 0
+    orig_preempt = scheduler._preempt_request
+
+    def counting_preempt(request, timestamp, **kwargs):
+        nonlocal repreempts_with_stale
+        if getattr(request, "num_stale_output_tokens", 0) > 0:
+            repreempts_with_stale += 1
+        return orig_preempt(request, timestamp, **kwargs)
+
+    scheduler._preempt_request = counting_preempt
+
+    resets = 0
+    reset_steps = {6, 14} if pp_size == 1 else {6, 7, 18, 19}
+
+    def before_step(step: int, engine: PipelinedEngine):
+        nonlocal resets
+        if pending:
+            scheduler.add_request(pending.pop(0))
+        if step in reset_steps and (engine.queue or scheduler.running):
+            scheduler.reset_prefix_cache(reset_running_requests=True)
+            resets += 1
+
+    engine = PipelinedEngine(
+        scheduler,
+        queue_size=pp_size + 1,
+        accept_drafts=lambda step, req_id, n: (step + int(req_id)) % (n + 1),
+    )
+    engine.run(before_step=before_step)
+
+    assert resets > 0, "test did not exercise reset_prefix_cache"
+    if pp_size > 1:
+        # The re-preempt-while-stale-pending window needs pipeline depth.
+        assert repreempts_with_stale > 0, (
+            "test did not exercise re-preemption with an undrained stale share"
+        )
+    for req in requests:
+        assert req.is_finished()
+        assert req.num_output_tokens == max_tokens
+        # Dropped tokens are never delivered; order must be preserved with
+        # no duplicates.
+        _assert_ordered_subset(
+            list(req.output_token_ids), engine.emitted[req.request_id]
+        )
+        _assert_positions_consistent(req, engine)
+        # All stale shares fully drained by the end.
+        assert getattr(req, "num_stale_output_tokens", 0) == 0

tests/v1/core/test_scheduler.py

@@ -1107,7 +1107,7 @@ def test_preemption_re_records_prefix_cache_query():
     request = create_requests(num_requests=1)[0]
     scheduler.add_request(request)
 
-    scheduler.schedule()
+    scheduler_output = scheduler.schedule()
     stats = scheduler.kv_cache_manager.prefix_cache_stats
     assert stats is not None
     assert (stats.requests, stats.preempted_requests) == (1, 0)
@@ -1116,6 +1116,21 @@ def test_preemption_re_records_prefix_cache_query():
     scheduler._preempt_request(request, 0.0)
     assert request.status == RequestStatus.PREEMPTED
 
+    scheduler.update_from_output(
+        scheduler_output,
+        ModelRunnerOutput(
+            req_ids=[request.request_id],
+            req_id_to_index={request.request_id: 0},
+            sampled_token_ids=[[1000]],
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=[],
+        ),
+    )
+    assert request.num_stale_output_tokens == 0
+    stats = scheduler.kv_cache_manager.prefix_cache_stats
+    assert stats is not None
+
     scheduler.schedule()
     assert request.status == RequestStatus.RUNNING
     assert stats.preempted_requests == 1

vllm/v1/core/sched/async_scheduler.py

@@ -49,23 +49,18 @@ class AsyncScheduler(Scheduler):
                 request.next_decode_eligible_step = self.current_step + self.pp_size
 
     def _update_request_with_output(
-        self, request: Request, new_token_ids: list[int]
+        self, request: Request, new_token_ids: list[int], is_stale: bool = False
     ) -> tuple[list[int], bool]:
-        if request.async_tokens_to_discard > 0:
-            # The request was force-preempted in reset_prefix_cache; drop one
-            # stale in-flight async output frame per call until the counter
-            # is drained.
-            request.async_tokens_to_discard -= 1
-            return [], False
-
         status_before_update = request.status
         new_token_ids, stopped = super()._update_request_with_output(
             request, new_token_ids
         )
 
-        # Update the number of output placeholders.
-        request.num_output_placeholders -= len(new_token_ids)
-        assert request.num_output_placeholders >= 0
+        # 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)
+            assert request.num_output_placeholders >= 0
 
         # Cache the new tokens. Preempted requests should be skipped.
         if status_before_update == RequestStatus.RUNNING:

vllm/v1/core/sched/scheduler.py

@@ -694,6 +694,17 @@ class Scheduler(SchedulerInterface):
                     step_skipped_waiting.prepend_request(request)
                     continue
 
+                if (
+                    request.num_stale_output_tokens > 0
+                    and not request.drop_stale_output
+                ):
+                    # Deliverable stale output still in flight: resuming now
+                    # could resample a position that output later delivers.
+                    # It drains within the pipeline depth.
+                    request_queue.pop_request()
+                    step_skipped_waiting.prepend_request(request)
+                    continue
+
                 # Check that adding the request still respects the max_loras
                 # constraint.
                 if (
@@ -1244,11 +1255,17 @@ class Scheduler(SchedulerInterface):
 
         return new_block_ids_to_zero or None
 
-    def _preempt_request(self, request: Request, timestamp: float) -> None:
+    def _preempt_request(
+        self, request: Request, timestamp: float, drop_stale_output: bool = False
+    ) -> None:
         """Preempt a request and put it back to the waiting queue.
 
         NOTE: The request should be popped from the running queue outside of this
         method.
+
+        drop_stale_output: drop (rather than deliver) any in-flight output; used
+        by reset_prefix_cache, whose same-step resume would otherwise deliver
+        tokens out of order.
         """
         assert request.status == RequestStatus.RUNNING, (
             "Only running requests can be preempted"
@@ -1260,6 +1277,18 @@ class Scheduler(SchedulerInterface):
         request.num_computed_tokens = 0
         if request.spec_token_ids:
             request.spec_token_ids = []
+        # Async scheduling: mark all in-flight output as stale. Its tokens are
+        # still delivered on return (dropping them would perturb spec-decode
+        # acceptance) but must not mutate the reset counters; each step drains
+        # its share in update_from_output. num_in_flight_tokens already
+        # includes any undrained stale share, so assign rather than accumulate.
+        # An undrained drop-mode share stays dropped: its positions have
+        # already been resampled.
+        request.drop_stale_output = drop_stale_output or (
+            request.drop_stale_output and request.num_stale_output_tokens > 0
+        )
+        request.num_stale_output_tokens = request.num_in_flight_tokens
+        request.num_output_placeholders = 0
         request.num_preemptions += 1
         if self.log_stats:
             request.record_event(EngineCoreEventType.PREEMPTED, timestamp)
@@ -1687,8 +1716,14 @@ class Scheduler(SchedulerInterface):
         for req_id, num_tokens_scheduled in num_scheduled_tokens.items():
             assert num_tokens_scheduled > 0
             request = self.requests.get(req_id)
+            output_is_stale = False
             if request is not None:
                 request.num_in_flight_tokens -= num_tokens_scheduled
+                # Drain any stale share (see _preempt_request) in lockstep.
+                if request.num_stale_output_tokens > 0:
+                    output_is_stale = True
+                    request.num_stale_output_tokens -= num_tokens_scheduled
+                    assert request.num_stale_output_tokens >= 0
             if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids:
                 # skip failed or rescheduled requests from KV load failure
                 continue
@@ -1702,6 +1737,10 @@ class Scheduler(SchedulerInterface):
                 # In this case, we use is_finished() to check.
                 continue
 
+            # Drop-mode stale output (same-step resume) is discarded entirely.
+            if output_is_stale and request.drop_stale_output:
+                continue
+
             req_index = model_runner_output.req_id_to_index[req_id]
             generated_token_ids = (
                 sampled_token_ids[req_index] if sampled_token_ids else []
@@ -1710,28 +1749,22 @@ class Scheduler(SchedulerInterface):
             scheduled_spec_token_ids = (
                 scheduler_output.scheduled_spec_decode_tokens.get(req_id)
             )
-            # Skip a stale frame still pending discard (async_tokens_to_discard
-            # > 0): its pre-reset rejection count would underflow the counters.
-            if (
-                scheduled_spec_token_ids
-                and (generated_token_ids or self.num_sampled_tokens_per_step == 0)
-                and request.async_tokens_to_discard == 0
+            if scheduled_spec_token_ids and (
+                generated_token_ids or self.num_sampled_tokens_per_step == 0
             ):
                 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
-                # num_computed_tokens represents the number of tokens
-                # processed in the current step, considering scheduled
-                # tokens and rejections. If some tokens are rejected,
-                # num_computed_tokens is decreased by the number of rejected
-                # tokens.
-                if request.num_computed_tokens > 0:
-                    request.num_computed_tokens -= num_rejected
-                # If async scheduling, num_output_placeholders also includes
-                # the scheduled spec tokens count and so is similarly adjusted.
-                if request.num_output_placeholders > 0:
-                    request.num_output_placeholders -= num_rejected
+                # 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
                 spec_decoding_stats = self.make_spec_decoding_stats(
                     spec_decoding_stats,
                     num_draft_tokens=num_draft_tokens,
@@ -1757,7 +1790,7 @@ class Scheduler(SchedulerInterface):
             # Check for stop and update request status.
             if new_token_ids:
                 new_token_ids, stopped = self._update_request_with_output(
-                    request, new_token_ids
+                    request, new_token_ids, is_stale=output_is_stale
                 )
             elif request.pooling_params and pooler_output is not None:
                 # Pooling stops as soon as there is output.
@@ -1899,6 +1932,7 @@ class Scheduler(SchedulerInterface):
         if stopped_preempted_reqs:
             # This is a rare case and unlikely to impact performance.
             self.waiting.remove_requests(stopped_preempted_reqs)
+            self.skipped_waiting.remove_requests(stopped_preempted_reqs)
 
         error_req_ids = set(self.grammar_compile_error_reqs)
         self.grammar_compile_error_reqs.clear()
@@ -2041,8 +2075,9 @@ class Scheduler(SchedulerInterface):
         return False
 
     def _update_request_with_output(
-        self, request: Request, new_token_ids: list[int]
+        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.
@@ -2388,15 +2423,7 @@ class Scheduler(SchedulerInterface):
             # running queue in FIFO order.
             while self.running:
                 request = self.running.pop()
-                self._preempt_request(request, timestamp)
-                # For async scheduling, any output frames already in flight at
-                # preemption time are now stale and must be discarded when they
-                # return. num_output_placeholders is exactly that count: 0 if
-                # the engine has drained (e.g. pause_generation(keep) waited
-                # for idle), 1 for vanilla async mid-step, or 1 + spec/PP frames
-                # otherwise.
-                request.async_tokens_to_discard = request.num_output_placeholders
-                request.num_output_placeholders = 0
+                self._preempt_request(request, timestamp, drop_stale_output=True)
 
             # Clear scheduled request ids cache. Since we are forcing preemption
             # + resumption in the same step, we must act as if these requests were

vllm/v1/request.py

@@ -149,7 +149,12 @@ class Request:
 
         # Used in async scheduling.
         self.num_output_placeholders = 0
-        self.async_tokens_to_discard = 0
+        # Tokens of output in flight when the request was preempted: delivered
+        # on return, but must not mutate the reset counters.
+        self.num_stale_output_tokens = 0
+        # Drop the stale output instead, for same-step preempt + resume
+        # (reset_prefix_cache).
+        self.drop_stale_output = False
 
         # Tokens of steps whose output is not yet processed (async scheduling
         # and PP run ahead of the GPU); `num_computed_tokens` counts them