diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 2995537aac66..c5c4274b2519 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/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 diff --git a/vllm/v1/core/sched/async_scheduler.py b/vllm/v1/core/sched/async_scheduler.py index 290204d58c39..28cad25caa27 100644 --- a/vllm/v1/core/sched/async_scheduler.py +++ b/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 )