diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index 89f05f1053..ec5ffef935 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -459,9 +459,6 @@ class RequestRunner: self.worker_connector.bind_connector_metadata(kv_connector_metadata) self.worker_connector.start_load_kv(self._dummy_ctx) - if scheduler_output.total_num_scheduled_tokens > 0: - self.worker_connector.wait_for_save() - if complete_transfers: self.offloading_spec.complete_transfers() 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 daa8042c8b..48dd61b50b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -104,12 +104,19 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA): pass def wait_for_save(self): - assert self.connector_worker is not None - assert isinstance(self._connector_metadata, OffloadingConnectorMetadata) - self.connector_worker.prepare_store_kv(self._connector_metadata) + # Store deferral is handled in get_finished(), which always runs even + # when wait_for_save() is skipped (e.g. kv_connector_no_forward). + pass 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: