[Bugfix][V1][AsyncScheduler] Clamp delivered tokens to reserved output placeholders instead of killing EngineCore

tests/v1/core/test_async_scheduler.py

@@ -811,3 +811,61 @@ def test_resumable_request_handoff():
 
     engine._process_oldest_step()
     assert len(request.output_token_ids) == 1
+
+
+def test_prefill_chunk_token_delivery_does_not_underflow():
+    """A token delivered for a request the scheduler classed as a prefill chunk
+    must not drive num_output_placeholders negative.
+
+    Under async scheduling the scheduler reserves one output placeholder per
+    step in which it expects the runner to sample, and reserves none for a
+    prefill chunk -- it relies on the runner returning empty token ids while a
+    request is still being prefilled. The runner keeps its own copy of that
+    predicate (``discard_request_mask``, built from
+    ``num_computed_tokens + num_scheduled_tokens < num_tokens``) and the two
+    can disagree, e.g. under chunked prefill with concurrency, or when a
+    preemption has just zeroed the placeholders. The runner then hands back a
+    token no placeholder was reserved for, which used to trip
+    ``assert request.num_output_placeholders >= 0`` and kill EngineCore with
+    EngineDeadError.
+
+    Such a token was sampled from a prefix the scheduler never committed to,
+    so it must be dropped rather than counted or appended.
+    """
+    scheduler = create_scheduler(async_scheduling=True, max_num_batched_tokens=64)
+    (request,) = create_requests(
+        num_requests=1, num_tokens=200, max_tokens=16, req_ids=["r"]
+    )
+    scheduler.add_request(request)
+
+    sched_output = scheduler.schedule()
+    assert "r" in sched_output.num_scheduled_tokens
+    assert request.is_prefill_chunk, "prompt must not fit in a single chunk"
+    assert request.num_output_placeholders == 0
+
+    num_output_tokens = request.num_output_tokens
+
+    # The runner disagrees about the prefill state and delivers a token.
+    scheduler.update_from_output(sched_output, _make_model_runner_output(sched_output))
+
+    assert request.num_output_placeholders == 0
+    assert request.num_output_tokens == num_output_tokens
+
+
+def test_async_placeholder_consumed_by_decode_delivery():
+    """The healthy path: a decode step reserves a placeholder and the token
+    delivered for it is appended, not discarded by the underflow guard."""
+    scheduler = create_scheduler(async_scheduling=True)
+    (request,) = create_requests(
+        num_requests=1, num_tokens=8, max_tokens=4, req_ids=["r"]
+    )
+    scheduler.add_request(request)
+
+    sched_output = scheduler.schedule()
+    assert not request.is_prefill_chunk
+    assert request.num_output_placeholders == 1
+
+    scheduler.update_from_output(sched_output, _make_model_runner_output(sched_output))
+
+    assert request.num_output_placeholders == 0
+    assert request.num_output_tokens == 1

vllm/v1/core/sched/async_scheduler.py

@@ -60,6 +60,38 @@ def _update_request_with_output(
         self, request: Request, new_token_ids: list[int], is_stale: bool = False
     ) -> tuple[list[int], bool]:
         status_before_update = request.status
+
+        if not is_stale and len(new_token_ids) > request.num_output_placeholders:
+            # The runner may only return tokens for steps that the scheduler
+            # reserved output placeholders for, and must return none while a
+            # request is still being prefilled. If it returns one anyway, that
+            # token was sampled from an incomplete prefix and is not a valid
+            # continuation: drop it rather than appending garbage to the
+            # response and driving num_output_placeholders negative, which trips
+            # the invariant below and kills EngineCore.
+            #
+            # Dropping is a correctness fix, not only crash protection: without
+            # it these tokens are appended to the request's output and streamed
+            # to the client as if they were valid continuations of the prompt.
+            num_reserved = request.num_output_placeholders
+            logger.warning_once(
+                "Model runner returned more tokens than there were reserved output "
+                "placeholders; dropping the extra token(s). A token sampled while "
+                "the request is still prefilling is not a valid continuation of the "
+                "prompt, so it must not be appended to the response or streamed to "
+                "the client. Run with debug logging to see which requests are "
+                "affected.",
+            )
+            logger.debug(
+                "Request %s: model runner returned %d token(s) but only %d output "
+                "placeholder(s) were reserved; dropping %d token(s).",
+                request.request_id,
+                len(new_token_ids),
+                num_reserved,
+                len(new_token_ids) - num_reserved,
+            )
+            new_token_ids = new_token_ids[:num_reserved]
+
         new_token_ids, stopped = super()._update_request_with_output(
             request, new_token_ids
         )