diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 0a203e6bd2..1b01d184dc 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/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 diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 2b2d241f57..5807574c64 100644 --- a/tests/v1/core/utils.py +++ b/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"}, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index c0eb247297..076d40a187 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/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", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index f3e89f6623..59dad186c6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/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" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index 2fe4bf6a5a..538dd341f9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/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, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py index f1dac13ca5..3ed34f169d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py +++ b/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, diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 7ba41e05e6..fa4a8daea9 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/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"