[Bugfix][Core] Stop zero-progress preemption cascades for deferred KV frees (#49675)
Signed-off-by: LearningMachine621 <147347705+LearningMachine621@users.noreply.github.com>
焦恺文 committed
Sep 11, 2026 at 08:40 UTC
84030bbe3d74d99bad477a3d2e37a973ccd8865c
4 files changed
+130
-27
tests/v1/core/test_async_scheduler.py
+39
-15
@@ -701,16 +701,29 @@ def test_requires_kv_delivery_defaults_to_producer_role():
701
assert scheduler.requires_kv_delivery is expected, role
702
703
704
-@pytest.mark.parametrize("kv_role", ["kv_producer", "kv_consumer"])
705
-def test_kv_pressure_preempt_mid_handoff(kv_role: str):
706
- """P/D race: a request is KV-pressure preempted while the output of its
707
- final prefill chunk -- the hand-off token that would finish it -- is in
708
- flight.
709
-
710
- On a producer, that output must be dropped so the request recomputes;
711
- delivering it would finish the request and hand off blocks the preemption
712
- already freed, so the consumer pulls garbage. A consumer hands nothing off,
713
- so it keeps the lossless deliver-stale path.
704
+@pytest.mark.parametrize(
705
+ ("kv_role", "defer_free"),
706
+ [
707
+ ("kv_producer", False),
708
+ ("kv_consumer", False),
709
+ ("kv_consumer", True),
710
+ ],
711
+)
712
+def test_kv_pressure_preempt_mid_handoff(kv_role: str, defer_free: bool):
713
+ """P/D race: KV pressure hits while the output of a request's final
714
+ prefill chunk -- the hand-off token that would finish it -- is in flight.
715
+
716
+ When the victim's blocks free immediately, the request is preempted. On a
717
+ producer, that output must be dropped so the request recomputes; delivering
718
+ it would finish the request and hand off blocks the preemption already
719
+ freed, so the consumer pulls garbage. A consumer hands nothing off, so it
720
+ keeps the lossless deliver-stale path.
721
+
722
+ A consumer with overlapping batches (async scheduling or PP) instead fences
723
+ the victim's free behind its in-flight output, so the allocation retry
724
+ stops instead of preempting; the request then finishes from that output
725
+ once it lands. The gate depends on the platform (async scheduling is
726
+ force-disabled on CPU), so force the flag to cover both paths everywhere.
727
"""
728
is_producer = kv_role == "kv_producer"
729
scheduler = create_scheduler(
@@ -722,10 +735,13 @@ def test_kv_pressure_preempt_mid_handoff(kv_role: str):
735
max_num_batched_tokens=512,
736
)
737
assert scheduler.requires_kv_delivery is is_producer
738
+ # The production gate requires overlapping batches, which async
739
+ # scheduling only provides off-CPU.
740
+ scheduler.defer_block_free = defer_free
741
742
# 32-token prompts fill 2 blocks each, exhausting the usable pool, so the
727
- # next decode allocation preempts the tail of the running queue (the handoff
728
- # request) while its prefill output is still in flight.
743
+ # next decode allocation targets the tail of the running queue (the
744
+ # handoff request) while its prefill output is still in flight.
745
decoder = create_requests(
746
num_requests=1, num_tokens=32, max_tokens=8, req_ids=["decoder"]
747
)[0]
@@ -739,9 +755,17 @@ def test_kv_pressure_preempt_mid_handoff(kv_role: str):
755
assert handoff.num_output_placeholders == 1
756
757
scheduler.schedule()
742
- assert handoff.status == RequestStatus.PREEMPTED
743
- assert handoff.num_stale_output_tokens == handoff.num_prompt_tokens
744
- assert handoff.drop_stale_output is is_producer
758
+ if defer_free:
759
+ # The victim's blocks are fenced behind its in-flight output, so
760
+ # preempting it could not satisfy the allocation; the retry stops
761
+ # instead of preempting.
762
+ assert handoff.status == RequestStatus.RUNNING
763
+ assert handoff.num_stale_output_tokens == 0
764
+ else:
765
+ # Blocks free immediately, so the handoff request is preempted.
766
+ assert handoff.status == RequestStatus.PREEMPTED
767
+ assert handoff.num_stale_output_tokens == handoff.num_prompt_tokens
768
+ assert handoff.drop_stale_output is is_producer
769
770
scheduler.update_from_output(sched_output, _make_model_runner_output(sched_output))
771
tests/v1/core/test_deferred_block_free.py
+72
-2
@@ -48,13 +48,17 @@ def _make_model_runner_output(
48
)
49
50
51
-def _create_deferring_scheduler():
51
+def _create_deferring_scheduler(scheduling_policy="fcfs"):
52
"""Async scheduler with deferred block freeing forced on.
53
54
The production gate additionally requires a PD KV-consumer connector;
55
the mechanism itself is independent of it.
56
"""
57
- scheduler = create_scheduler(model=MODEL, async_scheduling=True)
57
+ scheduler = create_scheduler(
58
+ model=MODEL,
59
+ async_scheduling=True,
60
+ scheduling_policy=scheduling_policy,
61
+ )
62
scheduler.defer_block_free = True
63
return scheduler
64
@@ -79,6 +83,24 @@ def _setup_request_with_inflight_step(scheduler, max_tokens: int = 5):
83
return request, out0, out1
84
85
86
+def _fail_one_allocation(scheduler, call_number: int):
87
+ real_allocate = scheduler.kv_cache_manager.allocate_slots
88
+ calls = 0
89
+
90
+ def allocate(*args, **kwargs):
91
+ nonlocal calls
92
+ calls += 1
93
+ if calls == call_number:
94
+ return None
95
+ return real_allocate(*args, **kwargs)
96
+
97
+ return patch.object(
98
+ scheduler.kv_cache_manager,
99
+ "allocate_slots",
100
+ side_effect=allocate,
101
+ )
102
+
103
+
104
def test_gate_enabled_for_async_consumer():
105
# Overlapping batches + consumer-side connector enables the gate. Async
106
# scheduling (which would give >1 concurrent batches) is force-disabled on
@@ -224,6 +246,54 @@ def test_preempt_defers_free_and_clears_bookkeeping():
246
assert pool.get_num_free_blocks() == num_free_initially
247
248
249
+@pytest.mark.parametrize(
250
+ ("policy", "failed_call"),
251
+ [
252
+ ("fcfs", 1),
253
+ ("priority", 2),
254
+ ],
255
+)
256
+def test_allocation_retry_waits_for_fence_then_succeeds(policy, failed_call):
257
+ scheduler = _create_deferring_scheduler(policy)
258
+ requests = create_requests(
259
+ num_requests=3,
260
+ num_tokens=NUM_PROMPT_TOKENS,
261
+ max_tokens=10,
262
+ stop_token_ids=[STOP_TOKEN_ID],
263
+ )
264
+ for request in requests:
265
+ scheduler.add_request(request)
266
+ out0 = scheduler.schedule()
267
+ out1 = scheduler.schedule()
268
+
269
+ worst, trigger, tail = requests
270
+ worst.priority, trigger.priority, tail.priority = 9, 0, 1
271
+
272
+ with _fail_one_allocation(scheduler, failed_call) as allocate:
273
+ blocked = scheduler.schedule()
274
+ assert allocate.call_count == failed_call
275
+ assert not blocked.preempted_req_ids
276
+ retry_trigger = requests[failed_call - 1]
277
+ assert retry_trigger.request_id not in blocked.num_scheduled_tokens
278
+
279
+ for output in (out0, out1, blocked):
280
+ if output.total_num_scheduled_tokens:
281
+ scheduler.update_from_output(
282
+ output,
283
+ _make_model_runner_output(output),
284
+ )
285
+
286
+ expected_victim = worst if policy == "priority" else tail
287
+ with _fail_one_allocation(scheduler, failed_call) as allocate:
288
+ resumed = scheduler.schedule()
289
+
290
+ assert allocate.call_count > failed_call
291
+ assert expected_victim.request_id in resumed.preempted_req_ids
292
+ assert retry_trigger.request_id in resumed.num_scheduled_tokens
293
+ if policy == "priority":
294
+ assert tail.request_id in resumed.num_scheduled_tokens
295
+
296
+
297
def test_multiple_deferred_frees_drain_in_order():
298
scheduler = _create_deferring_scheduler()
299
pool = scheduler.kv_cache_manager.block_pool
tests/v1/core/utils.py
+3
@@ -17,6 +17,7 @@ from vllm.config import (
17
SpeculativeConfig,
18
VllmConfig,
19
)
20
+from vllm.config.scheduler import SchedulerPolicy
21
from vllm.multimodal.inputs import (
22
MultiModalFeatureSpec,
23
MultiModalKwargsItem,
@@ -74,6 +75,7 @@ def create_scheduler(
75
use_v2_model_runner: bool | None = None,
76
kv_cache_spec: KVCacheSpec | None = None,
77
per_request_spec_decode_metrics: str = "none",
78
+ scheduling_policy: SchedulerPolicy = "fcfs",
79
) -> Scheduler | AsyncScheduler:
80
"""Create scheduler under test.
81
@@ -113,6 +115,7 @@ def create_scheduler(
115
is_encoder_decoder=model_config.is_encoder_decoder,
116
# Ensure admission/preemption mechanics are deterministic
117
watermark=0.0,
118
+ policy=scheduling_policy,
119
)
120
# Cache config, optionally force APC
121
cache_config = CacheConfig(
vllm/v1/core/sched/scheduler.py
+16
-10
@@ -745,13 +745,16 @@ class Scheduler(SchedulerInterface):
745
self.running,
746
key=lambda r: (r.priority, r.arrival_time),
747
)
748
- # Record the index of the preemption victim to
749
- # maintain accurate loop state.
748
+ else:
749
+ preempted_req = self.running[-1]
750
+
751
+ # A deferred free will not help with immediate allocation.
752
+ if not self._request_blocks_can_be_freed(preempted_req):
753
+ break
754
+
755
+ if self.policy == SchedulingPolicy.PRIORITY:
756
victim_index = self.running.index(preempted_req)
757
del self.running[victim_index]
752
- # Decrement the loop cursor if the removed request
753
- # preceded the current iteration, preventing the
754
- # silent omission of the subsequent request.
758
if victim_index < req_index:
759
req_index -= 1
760
@@ -2605,15 +2608,18 @@ class Scheduler(SchedulerInterface):
2608
logger.info("setting pause state to %s", pause_state.name)
2609
self._pause_state = pause_state
2610
2611
+ def _request_blocks_can_be_freed(self, request: Request) -> bool:
2612
+ # We must defer freeing blocks if an async kv connector may
2613
+ # write to them immediately (not ordered with GPU stream).
2614
+ return not self.defer_block_free or (
2615
+ request.last_sched_seq <= self.processed_step_seq
2616
+ )
2617
+
2618
def _free_request_blocks(self, request: Request):
2619
"""Free the request's KV blocks, deferring the return to the block
2620
pool when an in-flight GPU step may still write them.
2621
"""
2612
- if not self.defer_block_free or (
2613
- # Last scheduled step already processed: no in-flight write remains
2614
- # (always the case for a normal finish), so free now.
2615
- request.last_sched_seq <= self.processed_step_seq
2616
- ):
2622
+ if self._request_blocks_can_be_freed(request):
2623
self.kv_cache_manager.free(request)
2624
return
2625
blocks = self.kv_cache_manager.pop_blocks_for_free(request)