From 2f4a9759014e9294b4402aeae0fc25c11b47fc7d Mon Sep 17 00:00:00 2001 From: yzong-rh Date: Mon, 28 Sep 2026 15:26:57 -0400 Subject: [PATCH] [Bugfix] Fix resumable request + async scheduling handoff race (#58259) Signed-off-by: Yifan Zong --- tests/v1/core/test_async_scheduler.py | 47 ++++++++++++++++++++++++++- vllm/v1/core/sched/scheduler.py | 10 ++++++ 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 52167a5073..2995537aac 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -13,7 +13,7 @@ from vllm.v1.request import RequestStatus from vllm.v1.structured_output import StructuredOutputGrammar from vllm.v1.utils import ConstantList -from .utils import create_requests, create_scheduler, mock_kv +from .utils import EOS_TOKEN_ID, create_requests, create_scheduler, mock_kv pytestmark = pytest.mark.cpu_test @@ -766,3 +766,48 @@ def test_kv_pressure_preempt_mid_handoff(kv_role: str, defer_free: bool): else: assert handoff.is_finished() assert handoff.num_output_tokens == 1 + + +def test_resumable_request_handoff(): + """New streaming input resumes with the correct context and stale results + are ignored. + """ + scheduler = create_scheduler( + async_scheduling=True, + use_v2_model_runner=True, + ) + engine = PipelinedEngine(scheduler, queue_size=3) + engine._next_token = EOS_TOKEN_ID + + old_prompt_len = 10 + continuation_len = 12 + request_id = "resumable-race" + + request = create_requests( + 1, num_tokens=old_prompt_len, max_tokens=4, req_ids=[request_id] + )[0] + request.resumable = True + continuation = create_requests( + 1, num_tokens=continuation_len, max_tokens=4, req_ids=[request_id] + )[0] + continuation.resumable = True + + scheduler.add_request(request) + + for _ in range(3): + assert engine._schedule() + + engine._process_oldest_step() + scheduler.add_request(continuation) + assert engine._schedule() + new_step, _ = engine.queue[0] + assert new_step.scheduled_new_reqs[0].num_computed_tokens == old_prompt_len + + all_token_ids_before_stale = list(request.all_token_ids) + for _ in range(2): + engine._process_oldest_step() + assert list(request.output_token_ids) == [] + assert list(request.all_token_ids) == all_token_ids_before_stale + + engine._process_oldest_step() + assert len(request.output_token_ids) == 1 diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 8ebedef4ec..fa81626906 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -2403,6 +2403,16 @@ class Scheduler(SchedulerInterface): if not request.resumable: return True + # Drop the finished turn's in-flight work and resume from the + # materialized frontier. + safe_frontier = request.num_computed_tokens - request.num_output_placeholders + assert safe_frontier >= 0 + request.num_computed_tokens = safe_frontier + request.spec_token_ids = [] + request.drop_stale_output = True + request.num_stale_output_tokens = request.num_in_flight_tokens + request.num_output_placeholders = 0 + if request.streaming_queue: update = request.streaming_queue.popleft() if update is None: