diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index ad8d4fc362..cb766c6676 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -517,6 +517,7 @@ class RequestRunner: if complete_transfers: self.offloading_spec.complete_transfers() + self.worker_connector.wait_for_save() finished_sending, finished_recving = self.worker_connector.get_finished( scheduler_output.finished_req_ids ) diff --git a/tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py b/tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py index b6b18b092b..a1df2a964b 100644 --- a/tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py +++ b/tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py @@ -71,6 +71,7 @@ def test_kv_connector_mixin_clears_metadata(): assert connector._connector_metadata is None # Test connector wrapper records method calls assert connector.call_record.get("bind_connector_metadata", 0) == 1 + assert connector.call_record.get("wait_for_save", 0) == 1 assert connector.call_record.get("clear_connector_metadata", 0) == 1 finally: # Ensure we clean up the global connector between tests diff --git a/tests/v1/kv_connector/unit/test_lmcache_connector.py b/tests/v1/kv_connector/unit/test_lmcache_connector.py index a64f7f32d5..13f08a610c 100644 --- a/tests/v1/kv_connector/unit/test_lmcache_connector.py +++ b/tests/v1/kv_connector/unit/test_lmcache_connector.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Iterator from unittest.mock import MagicMock import pytest from vllm.distributed.kv_events import BlockStored +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata from vllm.distributed.kv_transfer.kv_connector.v1.lmcache_connector import ( LMCacheConnectorV1, LMCacheKVEvents, @@ -783,3 +785,32 @@ class TestIntegrationScenarios: assert aggregated_events[0].block_hashes == ["hash_common"] assert aggregated_events[0].parent_block_hash == "parent_common" assert aggregated_events[0].token_ids == [1, 2, 3] + + +def test_legacy_layerwise_engine_does_not_reuse_storers_on_empty_steps(): + """Empty steps must not advance generators exhausted by the last forward.""" + + class LegacyLayerwiseEngine: + layerwise_storers: list[Iterator[None]] + + def save_kv_layer(self, *args, **kwargs): + self.layerwise_storers.append(iter([None])) + + def wait_for_save(self): + for storer in self.layerwise_storers: + next(storer) + + connector = LMCacheConnectorV1.__new__(LMCacheConnectorV1) + connector._lmcache_engine = LegacyLayerwiseEngine() + metadata = KVConnectorMetadata() + connector.bind_connector_metadata(metadata) + connector.save_kv_layer("layer", MagicMock(), MagicMock()) + connector.wait_for_save() + connector.clear_connector_metadata() + + for _ in range(2): + connector.bind_connector_metadata(metadata) + connector.wait_for_save() + connector.clear_connector_metadata() + + assert connector._lmcache_engine.layerwise_storers == [] diff --git a/tests/v1/worker/test_gpu_kv_connector.py b/tests/v1/worker/test_gpu_kv_connector.py index 03c483dce1..5e56e8a317 100644 --- a/tests/v1/worker/test_gpu_kv_connector.py +++ b/tests/v1/worker/test_gpu_kv_connector.py @@ -96,4 +96,4 @@ def test_no_forward_starts_deferred_load_once(monkeypatch: pytest.MonkeyPatch): connector.no_forward(_scheduler_output(False)) # type: ignore[arg-type] - assert events == ["handle", "bind", "start", "clear"] + assert events == ["handle", "bind", "start", "wait", "clear"] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index 072018c039..bc3c32ba4c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -31,7 +31,7 @@ The class provides the following primitives: wait_for_layer_load() - blocks until layer i load is done save_kv_layer() - starts saving KV for layer i (maybe async) - wait_for_save() - blocks until all saves are done + wait_for_save() - submits saves and performs required synchronization get_transfer_results() - returns async send/receive completions and receive failures in one snapshot. @@ -356,11 +356,14 @@ class KVConnectorBase_V1(ABC): @abstractmethod def wait_for_save(self): - """Block until all the save operations is done. This is called - as the forward context exits to ensure that the async saving - from save_kv_layer is complete before finishing the forward. + """Submit saves and perform required synchronization. - This prevents overwrites of paged KV buffer before saving done. + Called once per step, including empty steps, after any target and draft + forwards. Metadata remains bound during this call. Do nothing if there + is no save work. + + Async saves must retain their source blocks until completion is reported + to the scheduler. Otherwise, wait for saves to finish before returning. """ pass diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py index 96191cd8f7..91552b7fe8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py @@ -142,6 +142,7 @@ class LMCacheConnectorV1(KVConnectorBase_V1): engine.current_layer = 0 if hasattr(engine, "layerwise_retrievers"): engine.layerwise_retrievers = [] + engine.layerwise_storers = [] def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None: """Start loading the KV cache from the connector to vLLM's paged diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py index a50e603db9..ed887f296b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py @@ -606,6 +606,7 @@ class LMCacheConnectorV1Impl: self.async_loading = config.enable_async_loading self.layerwise_retrievers: list[Generator[torch.Tensor | None, None, None]] = [] + self.layerwise_storers: list[Generator[Any, None, None]] = [] self._stats_monitor = LMCStatsMonitor.GetOrCreate() if role == KVConnectorRole.SCHEDULER: # Create lookup client using factory @@ -798,6 +799,7 @@ class LMCacheConnectorV1Impl: """ self.current_layer = 0 self.layerwise_retrievers = [] + self.layerwise_storers = [] @_lmcache_nvtx_annotate def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): 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 7357088175..def6e1e1f3 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -131,19 +131,14 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA): pass def wait_for_save(self): - # Store deferral is handled in get_finished(), which always runs even - # when wait_for_save() is skipped (e.g. kv_connector_no_forward). - pass + assert self.connector_worker is not None + assert isinstance(self._connector_metadata, OffloadingConnectorMetadata) + # Defer store jobs to the next step's start_kv_transfers. + self.connector_worker.prepare_store_kv(self._connector_metadata) def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: assert self.connector_worker is not None assert isinstance(self._connector_metadata, OffloadingConnectorMetadata) - - # Defer store jobs to the next step's start_kv_transfers. Done here - # (rather than wait_for_save) so stores are queued even on steps where - # wait_for_save is skipped. - self.connector_worker.prepare_store_kv(self._connector_metadata) - return self.connector_worker.get_finished(finished_req_ids) def build_connector_worker_meta(self) -> OffloadingWorkerMetadata | None: diff --git a/vllm/v1/simple_kv_offload/worker.py b/vllm/v1/simple_kv_offload/worker.py index 9fc16014c2..c120951eed 100644 --- a/vllm/v1/simple_kv_offload/worker.py +++ b/vllm/v1/simple_kv_offload/worker.py @@ -67,7 +67,6 @@ class SimpleCPUOffloadWorker: # Metadata for the current step self._connector_metadata: SimpleCPUOffloadMetadata | None = None - self._store_submitted = False # Compute-done event recorded before each store; reused across steps # (get_finished runs once per step, copy queue is FIFO). @@ -246,15 +245,12 @@ class SimpleCPUOffloadWorker: def bind_connector_metadata(self, metadata: SimpleCPUOffloadMetadata) -> None: self._connector_metadata = metadata - self._store_submitted = False if metadata.load_event >= 0: self._pending_load_event_indices.add(metadata.load_event) if metadata.store_event >= 0: self._pending_store_event_indices.add(metadata.store_event) def clear_connector_metadata(self) -> None: - # No-forward steps skip the normal wait_for_save hook. - self.wait_for_save() self._connector_metadata = None def start_load_kv(self) -> None: @@ -283,11 +279,7 @@ class SimpleCPUOffloadWorker: #45704 for the bug and #39306 for the srcAccessOrder rationale. """ metadata = self._connector_metadata - if ( - metadata is not None - and metadata.store_gpu_blocks - and not self._store_submitted - ): + if metadata is not None and metadata.store_gpu_blocks: backend = self._backend assert backend is not None if self._store_compute_done is None: @@ -301,7 +293,6 @@ class SimpleCPUOffloadWorker: events_list=self._store_events, wait_event=self._store_compute_done, ) - self._store_submitted = True def get_finished( self, finished_req_ids: set[str] diff --git a/vllm/v1/worker/gpu/kv_connector.py b/vllm/v1/worker/gpu/kv_connector.py index 88b7b0cf9d..1ba75beae2 100644 --- a/vllm/v1/worker/gpu/kv_connector.py +++ b/vllm/v1/worker/gpu/kv_connector.py @@ -35,9 +35,7 @@ class KVConnector: def finish_forward(self) -> None: pass - def post_forward( - self, finished_req_ids: set[str], wait_for_save: bool = True - ) -> KVConnectorOutput | None: + def post_forward(self, finished_req_ids: set[str]) -> KVConnectorOutput | None: return None def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput: @@ -97,9 +95,7 @@ class ActiveKVConnector(KVConnector): def reset_capture_state(self) -> None: self.kv_connector.reset_capture_state() - def post_forward( - self, finished_req_ids: set[str], wait_for_save: bool = True - ) -> KVConnectorOutput | None: + def post_forward(self, finished_req_ids: set[str]) -> KVConnectorOutput | None: if self._disabled: return None @@ -107,8 +103,7 @@ class ActiveKVConnector(KVConnector): self._start_load_kv() output = KVConnectorOutput() - if wait_for_save: - self.kv_connector.wait_for_save() + self.kv_connector.wait_for_save() transfer_results = self.kv_connector.get_transfer_results(finished_req_ids) output.finished_sending = transfer_results.finished_sending or None output.finished_recving = transfer_results.finished_recving or None @@ -129,7 +124,7 @@ class ActiveKVConnector(KVConnector): self.pre_forward(scheduler_output) self.finish_forward() finished_req_ids = scheduler_output.finished_req_ids - kv_connector_output = self.post_forward(finished_req_ids, wait_for_save=False) + kv_connector_output = self.post_forward(finished_req_ids) return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) def set_disabled(self, disabled: bool) -> None: diff --git a/vllm/v1/worker/kv_connector_model_runner_mixin.py b/vllm/v1/worker/kv_connector_model_runner_mixin.py index 26ffc159f3..40382b5c74 100644 --- a/vllm/v1/worker/kv_connector_model_runner_mixin.py +++ b/vllm/v1/worker/kv_connector_model_runner_mixin.py @@ -29,7 +29,7 @@ class KVConnectorModelRunnerMixin: with ( set_forward_context(None, vllm_config), KVConnectorModelRunnerMixin._get_kv_connector_output( - scheduler_output, wait_for_save=False + scheduler_output ) as kv_connector_output, ): pass @@ -66,7 +66,6 @@ class KVConnectorModelRunnerMixin: @contextmanager def _get_kv_connector_output( scheduler_output: "SchedulerOutput", - wait_for_save: bool = True, defer_finalize: bool = False, ) -> Generator[KVConnectorOutput, None, None]: output = KVConnectorOutput() @@ -89,7 +88,7 @@ class KVConnectorModelRunnerMixin: finally: if start_after_forward: kv_connector.start_load_kv(get_forward_context()) - if wait_for_save and not defer_finalize: + if not defer_finalize: kv_connector.wait_for_save() transfer_results = kv_connector.get_transfer_results(