[Bugfix] Fix resumable request + async scheduling handoff race (#58259)
Signed-off-by: Yifan Zong <yzong@redhat.com>
yzong-rh committed
Sep 28, 2026 at 15:26 UTC
2f4a9759014e9294b4402aeae0fc25c11b47fc7d
2 files changed
+56
-1
tests/v1/core/test_async_scheduler.py
+46
-1
@@ -13,7 +13,7 @@ from vllm.v1.request import RequestStatus
13
from vllm.v1.structured_output import StructuredOutputGrammar
14
from vllm.v1.utils import ConstantList
15
16
-from .utils import create_requests, create_scheduler, mock_kv
16
+from .utils import EOS_TOKEN_ID, create_requests, create_scheduler, mock_kv
17
18
pytestmark = pytest.mark.cpu_test
19
@@ -766,3 +766,48 @@ def test_kv_pressure_preempt_mid_handoff(kv_role: str, defer_free: bool):
766
else:
767
assert handoff.is_finished()
768
assert handoff.num_output_tokens == 1
769
+
770
+
771
+def test_resumable_request_handoff():
772
+ """New streaming input resumes with the correct context and stale results
773
+ are ignored.
774
+ """
775
+ scheduler = create_scheduler(
776
+ async_scheduling=True,
777
+ use_v2_model_runner=True,
778
+ )
779
+ engine = PipelinedEngine(scheduler, queue_size=3)
780
+ engine._next_token = EOS_TOKEN_ID
781
+
782
+ old_prompt_len = 10
783
+ continuation_len = 12
784
+ request_id = "resumable-race"
785
+
786
+ request = create_requests(
787
+ 1, num_tokens=old_prompt_len, max_tokens=4, req_ids=[request_id]
788
+ )[0]
789
+ request.resumable = True
790
+ continuation = create_requests(
791
+ 1, num_tokens=continuation_len, max_tokens=4, req_ids=[request_id]
792
+ )[0]
793
+ continuation.resumable = True
794
+
795
+ scheduler.add_request(request)
796
+
797
+ for _ in range(3):
798
+ assert engine._schedule()
799
+
800
+ engine._process_oldest_step()
801
+ scheduler.add_request(continuation)
802
+ assert engine._schedule()
803
+ new_step, _ = engine.queue[0]
804
+ assert new_step.scheduled_new_reqs[0].num_computed_tokens == old_prompt_len
805
+
806
+ all_token_ids_before_stale = list(request.all_token_ids)
807
+ for _ in range(2):
808
+ engine._process_oldest_step()
809
+ assert list(request.output_token_ids) == []
810
+ assert list(request.all_token_ids) == all_token_ids_before_stale
811
+
812
+ engine._process_oldest_step()
813
+ assert len(request.output_token_ids) == 1
vllm/v1/core/sched/scheduler.py
+10
@@ -2403,6 +2403,16 @@ class Scheduler(SchedulerInterface):
2403
if not request.resumable:
2404
return True
2405
2406
+ # Drop the finished turn's in-flight work and resume from the
2407
+ # materialized frontier.
2408
+ safe_frontier = request.num_computed_tokens - request.num_output_placeholders
2409
+ assert safe_frontier >= 0
2410
+ request.num_computed_tokens = safe_frontier
2411
+ request.spec_token_ids = []
2412
+ request.drop_stale_output = True
2413
+ request.num_stale_output_tokens = request.num_in_flight_tokens
2414
+ request.num_output_placeholders = 0
2415
+
2416
if request.streaming_queue:
2417
update = request.streaming_queue.popleft()
2418
if update is None: