@tduncan / Nvidia-Optimization-Stack-VLLM / commits / a0c092ee72

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

Nick Hill committed Jul 29, 2026 at 06:41 UTC a0c092ee72c0dcefbb3b3e74f97ac62d842e5f4b

5 files changed +385 -77

tests/v1/core/test_async_scheduler.py

+301 -35

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from collections import deque

3 +from collections import defaultdict, deque

4 +from collections.abc import Callable

5 from unittest.mock import Mock

7 import pytest

@@ -329,46 +330,311 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():

330 assert not scheduler.running

331

332

332 -def test_no_placeholder_underflow_on_discarded_spec_frame():

333 - num_spec = 5

333 +class PipelinedEngine:

334 + """Drive a real AsyncScheduler like EngineCore.step_with_batch_queue:

335 + schedule until the batch queue is full, then process the oldest step's

336 + output. Async PP runs pp_size+1 concurrent batches, so up to pp_size

337 + steps are in flight at each schedule() call -- the window in which

338 + preemption must handle output that has not yet returned. (Single-GPU e2e

339 + tests can never create this window: at PP=1, exactly one step is in

340 + flight and it is processed before a preempted request can resume.)

341 +

342 + The model runner is emulated with the V2 runner's own bookkeeping, from

343 + only what the scheduler serializes to it: slots flushed on

344 + preempted_req_ids, resumed requests re-added from the NewRequestData

345 + snapshot, sampling when a step reaches the end of the runner's own view

346 + of the sequence. This makes preemption races observable: a stale token

347 + delivered after a resume is scheduled extends the scheduler's sequence

348 + but not the runner's.

349 +

350 + Every sample emits a globally unique token tagged with its sampled

351 + position, so tests can assert exact delivery.

352 + """

353 +

354 + def __init__(

355 + self,

356 + scheduler: AsyncScheduler,

357 + queue_size: int,

358 + accept_drafts: Callable[[int, str, int], int] | None = None,

359 + ):

360 + self.scheduler = scheduler

361 + self.queue_size = queue_size

362 + self.accept_drafts = accept_drafts

363 + # In-flight steps: (scheduler_output, new_reqs snapshot) in FIFO order.

364 + self.queue: deque[tuple[SchedulerOutput, list[tuple[str, int, int]]]] = deque()

365 + # Runner-side request state: req_id -> [seq_len, num_computed] as the

366 + # runner sees them (its own sampled tokens, not the scheduler's).

367 + self.runner_view: dict[str, list[int]] = {}

368 + # All tokens the fake runner ever sampled, per request, in order.

369 + self.emitted: dict[str, list[int]] = defaultdict(list)

370 + # Sequence position each (globally unique) token was sampled for.

371 + self.emitted_position: dict[int, int] = {}

372 + self.step_idx = 0

373 + self._next_token = 1000

374 +

375 + def _schedule(self) -> bool:

376 + scheduler_output = self.scheduler.schedule()

377 + self.step_idx += 1

378 + # Snapshot what NewRequestData serializes at schedule time (both new

379 + # and resumed requests for the V2 runner).

380 + new_reqs = [

381 + (r.req_id, len(r.prefill_token_ids), r.num_computed_tokens)

382 + for r in scheduler_output.scheduled_new_reqs

383 + ]

384 + # Enqueue empty steps too (the engine executes them), so the runner

385 + # still observes their preempted/finished request ids in step order.

386 + self.queue.appendleft((scheduler_output, new_reqs))

387 + return True

388 +

389 + def _process_oldest_step(self) -> None:

390 + scheduler_output, new_reqs = self.queue.pop()

391 + # Worker-side state updates, in step order: flush preempted/finished

392 + # slots, then (re-)add new/resumed requests.

393 + for req_id in scheduler_output.preempted_req_ids or ():

394 + self.runner_view.pop(req_id, None)

395 + for req_id in scheduler_output.finished_req_ids or ():

396 + self.runner_view.pop(req_id, None)

397 + for req_id, seq_len, num_computed in new_reqs:

398 + self.runner_view[req_id] = [seq_len, num_computed]

399 +

400 + req_ids = list(scheduler_output.num_scheduled_tokens.keys())

401 + sampled_token_ids: list[list[int]] = []

402 + for req_id in req_ids:

403 + num_scheduled = scheduler_output.num_scheduled_tokens[req_id]

404 + view = self.runner_view.get(req_id)

405 + if view is None:

406 + # Slot already flushed (request finished/aborted mid-flight).

407 + sampled_token_ids.append([])

408 + continue

409 + seq_len, num_computed = view

410 + end = num_computed + num_scheduled

411 + if end < seq_len:

412 + # Partial prefill by the runner's own bookkeeping: no sample.

413 + view[1] = end

414 + sampled_token_ids.append([])

415 + continue

416 + drafts = scheduler_output.scheduled_spec_decode_tokens.get(req_id, ())

417 + num_accepted = (

418 + min(self.accept_drafts(self.step_idx, req_id, len(drafts)), len(drafts))

419 + if drafts and self.accept_drafts

420 + else 0

421 + )

422 + num_rejected = len(drafts) - num_accepted

423 + tokens = list(range(self._next_token, self._next_token + 1 + num_accepted))

424 + self._next_token += 1 + num_accepted

425 + self.emitted[req_id].extend(tokens)

426 + sampled_token_ids.append(tokens)

427 + # Rejected drafts roll back computed; the sampled tokens extend

428 + # the runner's sequence.

429 + view[1] = end - num_rejected

430 + view[0] = view[1] + 1

431 + for offset, token in enumerate(tokens):

432 + self.emitted_position[token] = view[0] - len(tokens) + offset

433 + model_runner_output = ModelRunnerOutput(

434 + req_ids=req_ids,

435 + req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)},

436 + sampled_token_ids=sampled_token_ids,

437 + logprobs=None,

438 + prompt_logprobs_dict={},

439 + pooler_output=[],

440 + )

441 + self.scheduler.update_from_output(scheduler_output, model_runner_output)

442 +

443 + def run(

444 + self,

445 + max_steps: int = 2000,

446 + before_step: Callable[[int, "PipelinedEngine"], None] | None = None,

447 + ) -> None:

448 + for i in range(max_steps):

449 + if not self.scheduler.has_requests() and not self.queue:

450 + return

451 + if before_step is not None:

452 + before_step(i, self)

453 + scheduled = (

454 + self.scheduler.has_requests()

455 + and len(self.queue) < self.queue_size

456 + and self._schedule()

457 + )

458 + if scheduled and len(self.queue) < self.queue_size:

459 + # Queue not yet full: the engine returns without blocking.

460 + continue

461 + if self.queue:

462 + self._process_oldest_step()

463 + raise AssertionError("engine loop did not converge")

464 +

465 +

466 +def _create_async_pp_scheduler(

467 + num_spec: int, pp_size: int = 3, num_blocks: int = 5

468 +) -> AsyncScheduler:

469 scheduler = create_scheduler(

470 async_scheduling=True,

336 - num_speculative_tokens=num_spec,

337 - speculative_method="ngram_gpu",

471 + num_speculative_tokens=num_spec or None,

472 + speculative_method="ngram_gpu" if num_spec else None,

473 + use_v2_model_runner=True,

474 + num_blocks=num_blocks,

475 + block_size=16,

476 + max_num_batched_tokens=512,

477 )

339 - req = create_requests(num_requests=1, max_tokens=20)[0]

340 - req.num_computed_tokens = req.num_tokens

341 - scheduler.requests[req.request_id] = req

342 - scheduler.running.append(req)

343 - req.status = RequestStatus.RUNNING

478 + # Emulate PP at the scheduler level; constructing with

479 + # pipeline_parallel_size>1 requires that many visible GPUs. Drive with

480 + # queue_size=pp_size+1 (V2 async PP runs pp_size+1 concurrent batches).

481 + scheduler.pp_size = pp_size

482 + scheduler.use_pp = pp_size > 1

483 + return scheduler

484 +

485 +

486 +def _assert_ordered_subset(delivered: list[int], emitted: list[int]) -> None:

487 + """Delivered tokens must be an order-preserving subset of the emitted

488 + tokens with no duplicates (tokens are globally unique)."""

489 + it = iter(emitted)

490 + for token in delivered:

491 + assert token in it, f"token {token} delivered out of order or twice"

492 +

493 +

494 +def _assert_positions_consistent(req, engine: PipelinedEngine) -> None:

495 + """The i-th delivered output token must be one the runner sampled for

496 + exactly sequence position prompt_len + i: catches a preempted request's

497 + stale output landing on a position the resumed request resampled (or

498 + vice versa), which token-stream equality alone cannot see."""

499 + for i, token in enumerate(req.output_token_ids):

500 + expected = req.num_prompt_tokens + i

501 + actual = engine.emitted_position[token]

502 + assert actual == expected, (

503 + f"output {i} of {req.request_id}: token sampled for position "

504 + f"{actual}, delivered as position {expected}"

505 + )

506

345 - req.num_output_placeholders = 1

346 - req.async_tokens_to_discard = num_spec

347 - computed_before = req.num_computed_tokens

507

349 - scheduler_output = SchedulerOutput(

350 - scheduled_new_reqs=[],

351 - scheduled_cached_reqs=CachedRequestData.make_empty(),

352 - num_scheduled_tokens={req.request_id: num_spec + 1},

353 - total_num_scheduled_tokens=num_spec + 1,

354 - scheduled_encoder_inputs={},

355 - scheduled_spec_decode_tokens={req.request_id: [10] * num_spec},

356 - num_common_prefix_blocks=[],

357 - finished_req_ids=set(),

358 - free_encoder_mm_hashes=[],

508 +@pytest.mark.parametrize("num_spec", [0, 3])

509 +def test_kv_pressure_preemption_with_inflight_output(num_spec: int):

510 + """KV-pressure preemption of requests with in-flight async output.

511 +

512 + PP=3 + async scheduling (batch queue of 4), a block pool small enough

513 + that decodes contend and preempt mid-flight, and staggered arrivals so

514 + the batch queue actually pipelines. A preempted request's in-flight steps

515 + still return: their tokens must be delivered exactly once, their stale

516 + spec-rejection counts must not corrupt the rolled-back counters, and the

517 + resume must not resample a position that output later delivers.

518 +

519 + Regression for the num_output_placeholders underflow EngineCore crash:

520 + with the fix reverted, the num_spec=3 variant fails with exactly

521 + ``assert request.num_output_placeholders >= 0`` when a stale spec output

522 + returns after the preempted request was resumed and sampled.

523 + """

524 + max_tokens = 24

525 + scheduler = _create_async_pp_scheduler(num_spec)

526 + requests = create_requests(

527 + num_requests=8, num_tokens=8, max_tokens=max_tokens, ignore_eos=True

528 )

360 - model_runner_output = ModelRunnerOutput(

361 - req_ids=[req.request_id],

362 - req_id_to_index={req.request_id: 0},

363 - sampled_token_ids=[[999]],

364 - logprobs=None,

365 - prompt_logprobs_dict={},

366 - pooler_output=[],

529 + pending = list(requests)

530 + for _ in range(2):

531 + scheduler.add_request(pending.pop(0))

532 +

533 + # Observe that the scenario under test actually occurs.

534 + preempts_with_inflight_output = 0

535 + orig_preempt = scheduler._preempt_request

536 +

537 + def counting_preempt(request, timestamp, **kwargs):

538 + nonlocal preempts_with_inflight_output

539 + if request.num_in_flight_tokens > 0:

540 + preempts_with_inflight_output += 1

541 + return orig_preempt(request, timestamp, **kwargs)

542 +

543 + scheduler._preempt_request = counting_preempt

544 +

545 + def add_requests(step: int, engine: PipelinedEngine):

546 + if pending:

547 + scheduler.add_request(pending.pop(0))

548 +

549 + engine = PipelinedEngine(

550 + scheduler,

551 + queue_size=4,

552 + # Deterministically vary spec acceptance so stale outputs carry

553 + # nonzero rejection counts.

554 + accept_drafts=lambda step, req_id, n: (step + int(req_id)) % (n + 1),

555 )

556 + engine.run(before_step=add_requests)

557

369 - scheduler.update_from_output(scheduler_output, model_runner_output)

558 + assert preempts_with_inflight_output > 0, (

559 + "test did not exercise preemption with in-flight output"

560 + )

561 + for req in requests:

562 + assert req.is_finished()

563 + assert req.num_output_tokens == max_tokens

564 + # Lossless: delivered tokens are exactly the sampled tokens, in order

565 + # (the excluded tail was emitted after the request finished).

566 + emitted = engine.emitted[req.request_id]

567 + assert list(req.output_token_ids) == emitted[:max_tokens]

568 + _assert_positions_consistent(req, engine)

569 +

570 +

571 +@pytest.mark.parametrize("pp_size", [1, 3])

572 +def test_reset_prefix_cache_with_inflight_output_under_kv_pressure(pp_size: int):

573 + """reset_prefix_cache(reset_running_requests=True) resumes requests in

574 + the same step it preempts them, so in-flight output must be dropped (the

575 + resume resamples those positions).

576 +

577 + pp_size=1: regression for the frame-based discard this fix replaces,

578 + which with spec decode drained one *token* count per output frame and

579 + over-discarded, corrupting the fresh frames after the resume.

580 + pp_size=3: back-to-back resets, so the second re-preempts requests whose

581 + dropped stale share is still in flight -- it must be recorded once (not

582 + accumulated) and stay dropped.

583 + """

584 + max_tokens = 24

585 + scheduler = _create_async_pp_scheduler(num_spec=3, pp_size=pp_size)

586 + requests = create_requests(

587 + num_requests=8, num_tokens=8, max_tokens=max_tokens, ignore_eos=True

588 + )

589 + pending = list(requests)

590 + for _ in range(2):

591 + scheduler.add_request(pending.pop(0))

592 +

593 + # Observe re-preemptions with an undrained stale share (the

594 + # double-count hazard).

595 + repreempts_with_stale = 0

596 + orig_preempt = scheduler._preempt_request

597 +

598 + def counting_preempt(request, timestamp, **kwargs):

599 + nonlocal repreempts_with_stale

600 + if getattr(request, "num_stale_output_tokens", 0) > 0:

601 + repreempts_with_stale += 1

602 + return orig_preempt(request, timestamp, **kwargs)

603 +

604 + scheduler._preempt_request = counting_preempt

605 +

606 + resets = 0

607 + reset_steps = {6, 14} if pp_size == 1 else {6, 7, 18, 19}

608 +

609 + def before_step(step: int, engine: PipelinedEngine):

610 + nonlocal resets

611 + if pending:

612 + scheduler.add_request(pending.pop(0))

613 + if step in reset_steps and (engine.queue or scheduler.running):

614 + scheduler.reset_prefix_cache(reset_running_requests=True)

615 + resets += 1

616 +

617 + engine = PipelinedEngine(

618 + scheduler,

619 + queue_size=pp_size + 1,

620 + accept_drafts=lambda step, req_id, n: (step + int(req_id)) % (n + 1),

621 + )

622 + engine.run(before_step=before_step)

623

371 - assert req.num_output_placeholders == 1

372 - assert req.num_computed_tokens == computed_before

373 - assert req.async_tokens_to_discard == num_spec - 1

374 - assert req.status == RequestStatus.RUNNING

624 + assert resets > 0, "test did not exercise reset_prefix_cache"

625 + if pp_size > 1:

626 + # The re-preempt-while-stale-pending window needs pipeline depth.

627 + assert repreempts_with_stale > 0, (

628 + "test did not exercise re-preemption with an undrained stale share"

629 + )

630 + for req in requests:

631 + assert req.is_finished()

632 + assert req.num_output_tokens == max_tokens

633 + # Dropped tokens are never delivered; order must be preserved with

634 + # no duplicates.

635 + _assert_ordered_subset(

636 + list(req.output_token_ids), engine.emitted[req.request_id]

637 + )

638 + _assert_positions_consistent(req, engine)

639 + # All stale shares fully drained by the end.

640 + assert getattr(req, "num_stale_output_tokens", 0) == 0

tests/v1/core/test_scheduler.py

+16 -1

@@ -1107,7 +1107,7 @@ def test_preemption_re_records_prefix_cache_query():

1107 request = create_requests(num_requests=1)[0]

1108 scheduler.add_request(request)

1109

1110 - scheduler.schedule()

1110 + scheduler_output = scheduler.schedule()

1111 stats = scheduler.kv_cache_manager.prefix_cache_stats

1112 assert stats is not None

1113 assert (stats.requests, stats.preempted_requests) == (1, 0)

@@ -1116,6 +1116,21 @@ def test_preemption_re_records_prefix_cache_query():

1116 scheduler._preempt_request(request, 0.0)

1117 assert request.status == RequestStatus.PREEMPTED

1118

1119 + scheduler.update_from_output(

1120 + scheduler_output,

1121 + ModelRunnerOutput(

1122 + req_ids=[request.request_id],

1123 + req_id_to_index={request.request_id: 0},

1124 + sampled_token_ids=[[1000]],

1125 + logprobs=None,

1126 + prompt_logprobs_dict={},

1127 + pooler_output=[],

1128 + ),

1129 + )

1130 + assert request.num_stale_output_tokens == 0

1131 + stats = scheduler.kv_cache_manager.prefix_cache_stats

1132 + assert stats is not None

1133 +

1134 scheduler.schedule()

1135 assert request.status == RequestStatus.RUNNING

1136 assert stats.preempted_requests == 1

vllm/v1/core/sched/async_scheduler.py

+6 -11

@@ -49,23 +49,18 @@ class AsyncScheduler(Scheduler):

49 request.next_decode_eligible_step = self.current_step + self.pp_size

51 def _update_request_with_output(

52 - self, request: Request, new_token_ids: list[int]

52 + self, request: Request, new_token_ids: list[int], is_stale: bool = False

53 ) -> tuple[list[int], bool]:

54 - if request.async_tokens_to_discard > 0:

55 - # The request was force-preempted in reset_prefix_cache; drop one

56 - # stale in-flight async output frame per call until the counter

57 - # is drained.

58 - request.async_tokens_to_discard -= 1

59 - return [], False

60 -

54 status_before_update = request.status

55 new_token_ids, stopped = super()._update_request_with_output(

56 request, new_token_ids

57 )

66 - # Update the number of output placeholders.

67 - request.num_output_placeholders -= len(new_token_ids)

68 - assert request.num_output_placeholders >= 0

59 + # Placeholders were zeroed at preemption; a stale delivery must not

60 + # decrement them (it would underflow).

61 + if not is_stale:

62 + request.num_output_placeholders -= len(new_token_ids)

63 + assert request.num_output_placeholders >= 0

65 # Cache the new tokens. Preempted requests should be skipped.

66 if status_before_update == RequestStatus.RUNNING:

vllm/v1/core/sched/scheduler.py

+56 -29

@@ -694,6 +694,17 @@ class Scheduler(SchedulerInterface):

694 step_skipped_waiting.prepend_request(request)

695 continue

696

697 + if (

698 + request.num_stale_output_tokens > 0

699 + and not request.drop_stale_output

700 + ):

701 + # Deliverable stale output still in flight: resuming now

702 + # could resample a position that output later delivers.

703 + # It drains within the pipeline depth.

704 + request_queue.pop_request()

705 + step_skipped_waiting.prepend_request(request)

706 + continue

707 +

708 # Check that adding the request still respects the max_loras

709 # constraint.

710 if (

@@ -1244,11 +1255,17 @@ class Scheduler(SchedulerInterface):

1255

1256 return new_block_ids_to_zero or None

1257

1247 - def _preempt_request(self, request: Request, timestamp: float) -> None:

1258 + def _preempt_request(

1259 + self, request: Request, timestamp: float, drop_stale_output: bool = False

1260 + ) -> None:

1261 """Preempt a request and put it back to the waiting queue.

1262

1263 NOTE: The request should be popped from the running queue outside of this

1264 method.

1265 +

1266 + drop_stale_output: drop (rather than deliver) any in-flight output; used

1267 + by reset_prefix_cache, whose same-step resume would otherwise deliver

1268 + tokens out of order.

1269 """

1270 assert request.status == RequestStatus.RUNNING, (

1271 "Only running requests can be preempted"

@@ -1260,6 +1277,18 @@ class Scheduler(SchedulerInterface):

1277 request.num_computed_tokens = 0

1278 if request.spec_token_ids:

1279 request.spec_token_ids = []

1280 + # Async scheduling: mark all in-flight output as stale. Its tokens are

1281 + # still delivered on return (dropping them would perturb spec-decode

1282 + # acceptance) but must not mutate the reset counters; each step drains

1283 + # its share in update_from_output. num_in_flight_tokens already

1284 + # includes any undrained stale share, so assign rather than accumulate.

1285 + # An undrained drop-mode share stays dropped: its positions have

1286 + # already been resampled.

1287 + request.drop_stale_output = drop_stale_output or (

1288 + request.drop_stale_output and request.num_stale_output_tokens > 0

1289 + )

1290 + request.num_stale_output_tokens = request.num_in_flight_tokens

1291 + request.num_output_placeholders = 0

1292 request.num_preemptions += 1

1293 if self.log_stats:

1294 request.record_event(EngineCoreEventType.PREEMPTED, timestamp)

@@ -1687,8 +1716,14 @@ class Scheduler(SchedulerInterface):

1716 for req_id, num_tokens_scheduled in num_scheduled_tokens.items():

1717 assert num_tokens_scheduled > 0

1718 request = self.requests.get(req_id)

1719 + output_is_stale = False

1720 if request is not None:

1721 request.num_in_flight_tokens -= num_tokens_scheduled

1722 + # Drain any stale share (see _preempt_request) in lockstep.

1723 + if request.num_stale_output_tokens > 0:

1724 + output_is_stale = True

1725 + request.num_stale_output_tokens -= num_tokens_scheduled

1726 + assert request.num_stale_output_tokens >= 0

1727 if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids:

1728 # skip failed or rescheduled requests from KV load failure

1729 continue

@@ -1702,6 +1737,10 @@ class Scheduler(SchedulerInterface):

1737 # In this case, we use is_finished() to check.

1738 continue

1739

1740 + # Drop-mode stale output (same-step resume) is discarded entirely.

1741 + if output_is_stale and request.drop_stale_output:

1742 + continue

1743 +

1744 req_index = model_runner_output.req_id_to_index[req_id]

1745 generated_token_ids = (

1746 sampled_token_ids[req_index] if sampled_token_ids else []

@@ -1710,28 +1749,22 @@ class Scheduler(SchedulerInterface):

1749 scheduled_spec_token_ids = (

1750 scheduler_output.scheduled_spec_decode_tokens.get(req_id)

1751 )

1713 - # Skip a stale frame still pending discard (async_tokens_to_discard

1714 - # > 0): its pre-reset rejection count would underflow the counters.

1715 - if (

1716 - scheduled_spec_token_ids

1717 - and (generated_token_ids or self.num_sampled_tokens_per_step == 0)

1718 - and request.async_tokens_to_discard == 0

1752 + if scheduled_spec_token_ids and (

1753 + generated_token_ids or self.num_sampled_tokens_per_step == 0

1754 ):

1755 num_draft_tokens = len(scheduled_spec_token_ids)

1756 num_sampled = self.num_sampled_tokens_per_step

1757 num_accepted = max(len(generated_token_ids) - num_sampled, 0)

1758 num_rejected = num_draft_tokens - num_accepted

1724 - # num_computed_tokens represents the number of tokens

1725 - # processed in the current step, considering scheduled

1726 - # tokens and rejections. If some tokens are rejected,

1727 - # num_computed_tokens is decreased by the number of rejected

1728 - # tokens.

1729 - if request.num_computed_tokens > 0:

1730 - request.num_computed_tokens -= num_rejected

1731 - # If async scheduling, num_output_placeholders also includes

1732 - # the scheduled spec tokens count and so is similarly adjusted.

1733 - if request.num_output_placeholders > 0:

1734 - request.num_output_placeholders -= num_rejected

1759 + # Rejections roll back num_computed_tokens (and, under async

1760 + # scheduling, num_output_placeholders, which covers the spec

1761 + # tokens). A stale rejection count predates the preemption

1762 + # rollback and must not apply.

1763 + if not output_is_stale:

1764 + if request.num_computed_tokens > 0:

1765 + request.num_computed_tokens -= num_rejected

1766 + if request.num_output_placeholders > 0:

1767 + request.num_output_placeholders -= num_rejected

1768 spec_decoding_stats = self.make_spec_decoding_stats(

1769 spec_decoding_stats,

1770 num_draft_tokens=num_draft_tokens,

@@ -1757,7 +1790,7 @@ class Scheduler(SchedulerInterface):

1790 # Check for stop and update request status.

1791 if new_token_ids:

1792 new_token_ids, stopped = self._update_request_with_output(

1760 - request, new_token_ids

1793 + request, new_token_ids, is_stale=output_is_stale

1794 )

1795 elif request.pooling_params and pooler_output is not None:

1796 # Pooling stops as soon as there is output.

@@ -1899,6 +1932,7 @@ class Scheduler(SchedulerInterface):

1932 if stopped_preempted_reqs:

1933 # This is a rare case and unlikely to impact performance.

1934 self.waiting.remove_requests(stopped_preempted_reqs)

1935 + self.skipped_waiting.remove_requests(stopped_preempted_reqs)

1936

1937 error_req_ids = set(self.grammar_compile_error_reqs)

1938 self.grammar_compile_error_reqs.clear()

@@ -2041,8 +2075,9 @@ class Scheduler(SchedulerInterface):

2075 return False

2076

2077 def _update_request_with_output(

2044 - self, request: Request, new_token_ids: list[int]

2078 + self, request: Request, new_token_ids: list[int], is_stale: bool = False

2079 ) -> tuple[list[int], bool]:

2080 + # is_stale is only used by the AsyncScheduler override.

2081 # Append generated tokens and check for stop. Note that if

2082 # a request is still being prefilled, we expect the model runner

2083 # to return empty token ids for the request.

@@ -2388,15 +2423,7 @@ class Scheduler(SchedulerInterface):

2423 # running queue in FIFO order.

2424 while self.running:

2425 request = self.running.pop()

2391 - self._preempt_request(request, timestamp)

2392 - # For async scheduling, any output frames already in flight at

2393 - # preemption time are now stale and must be discarded when they

2394 - # return. num_output_placeholders is exactly that count: 0 if

2395 - # the engine has drained (e.g. pause_generation(keep) waited

2396 - # for idle), 1 for vanilla async mid-step, or 1 + spec/PP frames

2397 - # otherwise.

2398 - request.async_tokens_to_discard = request.num_output_placeholders

2399 - request.num_output_placeholders = 0

2426 + self._preempt_request(request, timestamp, drop_stale_output=True)

2427

2428 # Clear scheduled request ids cache. Since we are forcing preemption

2429 # + resumption in the same step, we must act as if these requests were

vllm/v1/request.py

+6 -1

@@ -149,7 +149,12 @@ class Request:

149

150 # Used in async scheduling.

151 self.num_output_placeholders = 0

152 - self.async_tokens_to_discard = 0

152 + # Tokens of output in flight when the request was preempted: delivered

153 + # on return, but must not mutate the reset counters.

154 + self.num_stale_output_tokens = 0

155 + # Drop the stale output instead, for same-step preempt + resume

156 + # (reset_prefix_cache).

157 + self.drop_stale_output = False

158

159 # Tokens of steps whose output is not yet processed (async scheduling

160 # and PP run ahead of the GPU); `num_computed_tokens` counts them