[BugFix] Fix P/D preemption race condition (#50297)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Nick Hill committed
Jul 29, 2026 at 18:11 UTC
437e0b7f8eb787e1cd7bdf090e47a015b54fd22c
7 files changed
+117
-5
tests/v1/core/test_async_scheduler.py
+73
@@ -638,3 +638,76 @@ def test_reset_prefix_cache_with_inflight_output_under_kv_pressure(pp_size: int)
638
_assert_positions_consistent(req, engine)
639
# All stale shares fully drained by the end.
640
assert getattr(req, "num_stale_output_tokens", 0) == 0
641
+
642
+
643
+def test_requires_kv_delivery_defaults_to_producer_role():
644
+ # No connector: nothing is handed off, so keep the lossless deliver-stale
645
+ # path on preemption.
646
+ assert create_scheduler(async_scheduling=True).requires_kv_delivery is False
647
+ # Only a producer hands KV off when a request completes.
648
+ for role, expected in (
649
+ ("kv_producer", True),
650
+ ("kv_both", True),
651
+ ("kv_consumer", False),
652
+ ):
653
+ scheduler = create_scheduler(
654
+ async_scheduling=True, use_kv_connector=True, kv_role=role
655
+ )
656
+ assert scheduler.requires_kv_delivery is expected, role
657
+
658
+
659
+@pytest.mark.parametrize("kv_role", ["kv_producer", "kv_consumer"])
660
+def test_kv_pressure_preempt_mid_handoff(kv_role: str):
661
+ """P/D race: a request is KV-pressure preempted while the output of its
662
+ final prefill chunk -- the hand-off token that would finish it -- is in
663
+ flight.
664
+
665
+ On a producer, that output must be dropped so the request recomputes;
666
+ delivering it would finish the request and hand off blocks the preemption
667
+ already freed, so the consumer pulls garbage. A consumer hands nothing off,
668
+ so it keeps the lossless deliver-stale path.
669
+ """
670
+ is_producer = kv_role == "kv_producer"
671
+ scheduler = create_scheduler(
672
+ async_scheduling=True,
673
+ use_kv_connector=True,
674
+ kv_role=kv_role,
675
+ num_blocks=5,
676
+ block_size=16,
677
+ max_num_batched_tokens=512,
678
+ )
679
+ assert scheduler.requires_kv_delivery is is_producer
680
+
681
+ # 32-token prompts fill 2 blocks each, exhausting the usable pool, so the
682
+ # next decode allocation preempts the tail of the running queue (the handoff
683
+ # request) while its prefill output is still in flight.
684
+ decoder = create_requests(
685
+ num_requests=1, num_tokens=32, max_tokens=8, req_ids=["decoder"]
686
+ )[0]
687
+ handoff = create_requests(
688
+ num_requests=1, num_tokens=32, max_tokens=1, req_ids=["handoff"]
689
+ )[0]
690
+ scheduler.add_request(decoder)
691
+ scheduler.add_request(handoff)
692
+ sched_output = scheduler.schedule()
693
+ assert handoff.status == RequestStatus.RUNNING
694
+ assert handoff.num_output_placeholders == 1
695
+
696
+ scheduler.schedule()
697
+ assert handoff.status == RequestStatus.PREEMPTED
698
+ assert handoff.num_stale_output_tokens == handoff.num_prompt_tokens
699
+ assert handoff.drop_stale_output is is_producer
700
+
701
+ scheduler.update_from_output(sched_output, _make_model_runner_output(sched_output))
702
+
703
+ assert handoff.num_stale_output_tokens == 0
704
+ if is_producer:
705
+ # Dropped: recomputed from the waiting queue, so the hand-off happens
706
+ # against real KV.
707
+ assert not handoff.is_finished()
708
+ assert handoff.status == RequestStatus.PREEMPTED
709
+ assert handoff.num_output_tokens == 0
710
+ assert handoff.request_id in scheduler.requests
711
+ else:
712
+ assert handoff.is_finished()
713
+ assert handoff.num_output_tokens == 1
tests/v1/core/utils.py
+4
-3
@@ -55,6 +55,7 @@ def create_scheduler(
55
long_prefill_token_threshold: int = 0,
56
disable_chunked_mm_input: bool = False,
57
use_kv_connector: None | bool | str | MockKVConfig = None,
58
+ kv_role: str = "kv_both",
59
num_blocks: int = 10000,
60
block_size: int = 16,
61
max_model_len: int | None = None,
@@ -115,7 +116,7 @@ def create_scheduler(
116
if isinstance(use_kv_connector, MockKVConfig):
117
kv_transfer_config = KVTransferConfig(
118
kv_connector="MockKVConnector",
118
- kv_role="kv_both",
119
+ kv_role=kv_role,
120
kv_connector_extra_config={
121
"matched_tokens": use_kv_connector.matched_tokens,
122
"is_async": use_kv_connector.is_async,
@@ -127,12 +128,12 @@ def create_scheduler(
128
elif isinstance(use_kv_connector, str):
129
kv_transfer_config = KVTransferConfig(
130
kv_connector=use_kv_connector,
130
- kv_role="kv_both",
131
+ kv_role=kv_role,
132
)
133
elif use_kv_connector:
134
kv_transfer_config = KVTransferConfig(
135
kv_connector="ExampleConnector",
135
- kv_role="kv_both",
136
+ kv_role=kv_role,
137
kv_connector_extra_config={"shared_storage_path": "local_storage"},
138
)
139
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+12
@@ -181,6 +181,18 @@ class KVConnectorBase_V1(ABC):
181
"""
182
return False
183
184
+ @property
185
+ def requires_kv_delivery(self) -> bool:
186
+ """Whether this connector hands off KV that must be reliably delivered.
187
+
188
+ If True, a request preempted while its hand-off is still pending is
189
+ recomputed rather than allowed to finish and hand off blocks that the
190
+ preemption already freed. Defaults to the producer role, since only a
191
+ producer hands KV off when a request completes. Best-effort caches
192
+ return False, as a dropped save is just a future cache miss.
193
+ """
194
+ return self._kv_transfer_config.is_kv_producer
195
+
196
def __init__(
197
self,
198
vllm_config: "VllmConfig",
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+4
@@ -209,6 +209,10 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
209
return False
210
return all(c.prefer_cross_layer_blocks for c in self._connectors)
211
212
+ @property
213
+ def requires_kv_delivery(self) -> bool:
214
+ return any(c.requires_kv_delivery for c in self._connectors)
215
+
216
@classmethod
217
def _get_connector_classes_and_configs(
218
cls, vllm_config: "VllmConfig"
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+6
@@ -51,6 +51,12 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
51
def prefer_cross_layer_blocks(self) -> bool:
52
return True
53
54
+ @property
55
+ def requires_kv_delivery(self) -> bool:
56
+ # Runs as kv_both, but is a best-effort cache: a dropped save is just a
57
+ # future cache miss, so opt out of the producer-role default.
58
+ return False
59
+
60
def __init__(
61
self,
62
vllm_config: VllmConfig,
vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
+6
@@ -45,6 +45,12 @@ DEFAULT_CPU_CAPACITY_BYTES = 8 * (1024**3)
45
class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):
46
"""CPU KV cache offloading with custom kernel transfers and BlockPool LRU."""
47
48
+ @property
49
+ def requires_kv_delivery(self) -> bool:
50
+ # Runs as kv_both, but is a best-effort cache: a dropped save is just a
51
+ # future cache miss, so opt out of the producer-role default.
52
+ return False
53
+
54
def __init__(
55
self,
56
vllm_config: VllmConfig,
vllm/v1/core/sched/scheduler.py
+12
-2
@@ -129,6 +129,9 @@ class Scheduler(SchedulerInterface):
129
self.connector_prefix_cache_stats: PrefixCacheStats | None = None
130
self.recompute_kv_load_failures = True
131
self.defer_block_free = False
132
+ # Whether a preempted request's in-flight output must be dropped; see
133
+ # KVConnectorBase_V1.requires_kv_delivery.
134
+ self.requires_kv_delivery = False
135
kv_transfer_config = self.vllm_config.kv_transfer_config
136
if kv_transfer_config is not None:
137
assert not self.is_encoder_decoder, (
@@ -152,6 +155,8 @@ class Scheduler(SchedulerInterface):
155
if multiple_inflight_batches and kv_transfer_config.is_kv_consumer:
156
self.defer_block_free = True
157
158
+ self.requires_kv_delivery = self.connector.requires_kv_delivery
159
+
160
self.kv_event_publisher = EventPublisherFactory.create(
161
self.kv_events_config,
162
self.parallel_config.data_parallel_index,
@@ -602,7 +607,11 @@ class Scheduler(SchedulerInterface):
607
else:
608
preempted_req = self.running.pop()
609
605
- self._preempt_request(preempted_req, scheduled_timestamp)
610
+ self._preempt_request(
611
+ preempted_req,
612
+ scheduled_timestamp,
613
+ drop_stale_output=self.requires_kv_delivery,
614
+ )
615
preempted_reqs.append(preempted_req)
616
if preempted_req == request:
617
# No more request to preempt. Cannot schedule this request.
@@ -1265,7 +1274,8 @@ class Scheduler(SchedulerInterface):
1274
1275
drop_stale_output: drop (rather than deliver) any in-flight output; used
1276
by reset_prefix_cache, whose same-step resume would otherwise deliver
1268
- tokens out of order.
1277
+ tokens out of order, and for connectors with a pending KV hand-off,
1278
+ which the preemption's block free would leave without valid KV.
1279
"""
1280
assert request.status == RequestStatus.RUNNING, (
1281
"Only running requests can be preempted"