[Bugfix][KVConnector] Finalize saves on steps without a forward (#57775)

Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: OpenAI Codex <noreply@openai.com>
Co-authored-by: Kimi Code <noreply@moonshot.ai>

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
             )

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

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

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"]

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
 

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

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]):

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:

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]

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:

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(