[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>

tests/v1/core/test_async_scheduler.py

@@ -638,3 +638,76 @@ def test_reset_prefix_cache_with_inflight_output_under_kv_pressure(pp_size: int)
         _assert_positions_consistent(req, engine)
         # All stale shares fully drained by the end.
         assert getattr(req, "num_stale_output_tokens", 0) == 0
+
+
+def test_requires_kv_delivery_defaults_to_producer_role():
+    # No connector: nothing is handed off, so keep the lossless deliver-stale
+    # path on preemption.
+    assert create_scheduler(async_scheduling=True).requires_kv_delivery is False
+    # Only a producer hands KV off when a request completes.
+    for role, expected in (
+        ("kv_producer", True),
+        ("kv_both", True),
+        ("kv_consumer", False),
+    ):
+        scheduler = create_scheduler(
+            async_scheduling=True, use_kv_connector=True, kv_role=role
+        )
+        assert scheduler.requires_kv_delivery is expected, role
+
+
+@pytest.mark.parametrize("kv_role", ["kv_producer", "kv_consumer"])
+def test_kv_pressure_preempt_mid_handoff(kv_role: str):
+    """P/D race: a request is KV-pressure preempted while the output of its
+    final prefill chunk -- the hand-off token that would finish it -- is in
+    flight.
+
+    On a producer, that output must be dropped so the request recomputes;
+    delivering it would finish the request and hand off blocks the preemption
+    already freed, so the consumer pulls garbage. A consumer hands nothing off,
+    so it keeps the lossless deliver-stale path.
+    """
+    is_producer = kv_role == "kv_producer"
+    scheduler = create_scheduler(
+        async_scheduling=True,
+        use_kv_connector=True,
+        kv_role=kv_role,
+        num_blocks=5,
+        block_size=16,
+        max_num_batched_tokens=512,
+    )
+    assert scheduler.requires_kv_delivery is is_producer
+
+    # 32-token prompts fill 2 blocks each, exhausting the usable pool, so the
+    # next decode allocation preempts the tail of the running queue (the handoff
+    # request) while its prefill output is still in flight.
+    decoder = create_requests(
+        num_requests=1, num_tokens=32, max_tokens=8, req_ids=["decoder"]
+    )[0]
+    handoff = create_requests(
+        num_requests=1, num_tokens=32, max_tokens=1, req_ids=["handoff"]
+    )[0]
+    scheduler.add_request(decoder)
+    scheduler.add_request(handoff)
+    sched_output = scheduler.schedule()
+    assert handoff.status == RequestStatus.RUNNING
+    assert handoff.num_output_placeholders == 1
+
+    scheduler.schedule()
+    assert handoff.status == RequestStatus.PREEMPTED
+    assert handoff.num_stale_output_tokens == handoff.num_prompt_tokens
+    assert handoff.drop_stale_output is is_producer
+
+    scheduler.update_from_output(sched_output, _make_model_runner_output(sched_output))
+
+    assert handoff.num_stale_output_tokens == 0
+    if is_producer:
+        # Dropped: recomputed from the waiting queue, so the hand-off happens
+        # against real KV.
+        assert not handoff.is_finished()
+        assert handoff.status == RequestStatus.PREEMPTED
+        assert handoff.num_output_tokens == 0
+        assert handoff.request_id in scheduler.requests
+    else:
+        assert handoff.is_finished()
+        assert handoff.num_output_tokens == 1

tests/v1/core/utils.py

@@ -55,6 +55,7 @@ def create_scheduler(
     long_prefill_token_threshold: int = 0,
     disable_chunked_mm_input: bool = False,
     use_kv_connector: None | bool | str | MockKVConfig = None,
+    kv_role: str = "kv_both",
     num_blocks: int = 10000,
     block_size: int = 16,
     max_model_len: int | None = None,
@@ -115,7 +116,7 @@ def create_scheduler(
     if isinstance(use_kv_connector, MockKVConfig):
         kv_transfer_config = KVTransferConfig(
             kv_connector="MockKVConnector",
-            kv_role="kv_both",
+            kv_role=kv_role,
             kv_connector_extra_config={
                 "matched_tokens": use_kv_connector.matched_tokens,
                 "is_async": use_kv_connector.is_async,
@@ -127,12 +128,12 @@ def create_scheduler(
     elif isinstance(use_kv_connector, str):
         kv_transfer_config = KVTransferConfig(
             kv_connector=use_kv_connector,
-            kv_role="kv_both",
+            kv_role=kv_role,
         )
     elif use_kv_connector:
         kv_transfer_config = KVTransferConfig(
             kv_connector="ExampleConnector",
-            kv_role="kv_both",
+            kv_role=kv_role,
             kv_connector_extra_config={"shared_storage_path": "local_storage"},
         )
 

vllm/distributed/kv_transfer/kv_connector/v1/base.py

@@ -181,6 +181,18 @@ class KVConnectorBase_V1(ABC):
         """
         return False
 
+    @property
+    def requires_kv_delivery(self) -> bool:
+        """Whether this connector hands off KV that must be reliably delivered.
+
+        If True, a request preempted while its hand-off is still pending is
+        recomputed rather than allowed to finish and hand off blocks that the
+        preemption already freed. Defaults to the producer role, since only a
+        producer hands KV off when a request completes. Best-effort caches
+        return False, as a dropped save is just a future cache miss.
+        """
+        return self._kv_transfer_config.is_kv_producer
+
     def __init__(
         self,
         vllm_config: "VllmConfig",

vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py

@@ -209,6 +209,10 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
             return False
         return all(c.prefer_cross_layer_blocks for c in self._connectors)
 
+    @property
+    def requires_kv_delivery(self) -> bool:
+        return any(c.requires_kv_delivery for c in self._connectors)
+
     @classmethod
     def _get_connector_classes_and_configs(
         cls, vllm_config: "VllmConfig"

vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py

@@ -51,6 +51,12 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
     def prefer_cross_layer_blocks(self) -> bool:
         return True
 
+    @property
+    def requires_kv_delivery(self) -> bool:
+        # Runs as kv_both, but is a best-effort cache: a dropped save is just a
+        # future cache miss, so opt out of the producer-role default.
+        return False
+
     def __init__(
         self,
         vllm_config: VllmConfig,

vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py

@@ -45,6 +45,12 @@ DEFAULT_CPU_CAPACITY_BYTES = 8 * (1024**3)
 class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):
     """CPU KV cache offloading with custom kernel transfers and BlockPool LRU."""
 
+    @property
+    def requires_kv_delivery(self) -> bool:
+        # Runs as kv_both, but is a best-effort cache: a dropped save is just a
+        # future cache miss, so opt out of the producer-role default.
+        return False
+
     def __init__(
         self,
         vllm_config: VllmConfig,

vllm/v1/core/sched/scheduler.py

@@ -129,6 +129,9 @@ class Scheduler(SchedulerInterface):
         self.connector_prefix_cache_stats: PrefixCacheStats | None = None
         self.recompute_kv_load_failures = True
         self.defer_block_free = False
+        # Whether a preempted request's in-flight output must be dropped; see
+        # KVConnectorBase_V1.requires_kv_delivery.
+        self.requires_kv_delivery = False
         kv_transfer_config = self.vllm_config.kv_transfer_config
         if kv_transfer_config is not None:
             assert not self.is_encoder_decoder, (
@@ -152,6 +155,8 @@ class Scheduler(SchedulerInterface):
             if multiple_inflight_batches and kv_transfer_config.is_kv_consumer:
                 self.defer_block_free = True
 
+            self.requires_kv_delivery = self.connector.requires_kv_delivery
+
         self.kv_event_publisher = EventPublisherFactory.create(
             self.kv_events_config,
             self.parallel_config.data_parallel_index,
@@ -602,7 +607,11 @@ class Scheduler(SchedulerInterface):
                     else:
                         preempted_req = self.running.pop()
 
-                    self._preempt_request(preempted_req, scheduled_timestamp)
+                    self._preempt_request(
+                        preempted_req,
+                        scheduled_timestamp,
+                        drop_stale_output=self.requires_kv_delivery,
+                    )
                     preempted_reqs.append(preempted_req)
                     if preempted_req == request:
                         # No more request to preempt. Cannot schedule this request.
@@ -1265,7 +1274,8 @@ class Scheduler(SchedulerInterface):
 
         drop_stale_output: drop (rather than deliver) any in-flight output; used
         by reset_prefix_cache, whose same-step resume would otherwise deliver
-        tokens out of order.
+        tokens out of order, and for connectors with a pending KV hand-off,
+        which the preemption's block free would leave without valid KV.
         """
         assert request.status == RequestStatus.RUNNING, (
             "Only running requests can be preempted"