A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 20:15 UTC
[kv_offload][BugFix] Fix store deferral (#41945)
Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com>
Martin Hickey committed
May 12, 2026 at 16:04 UTC
418ba8ef141c372bf8ce6835f4ebfc6298a0af37
2 files changed
+9
-5
tests/v1/kv_connector/unit/offloading_connector/utils.py
-3
@@ -459,9 +459,6 @@ class RequestRunner:
459
self.worker_connector.bind_connector_metadata(kv_connector_metadata)
460
self.worker_connector.start_load_kv(self._dummy_ctx)
461
462
- if scheduler_output.total_num_scheduled_tokens > 0:
463
- self.worker_connector.wait_for_save()
464
-
462
if complete_transfers:
463
self.offloading_spec.complete_transfers()
464
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+9
-2
@@ -104,12 +104,19 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
104
pass
105
106
def wait_for_save(self):
107
+ # Store deferral is handled in get_finished(), which always runs even
108
+ # when wait_for_save() is skipped (e.g. kv_connector_no_forward).
109
+ pass
110
+
111
+ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
112
assert self.connector_worker is not None
113
assert isinstance(self._connector_metadata, OffloadingConnectorMetadata)
114
+
115
+ # Defer store jobs to the next step's start_kv_transfers. Done here
116
+ # (rather than wait_for_save) so stores are queued even on steps where
117
+ # wait_for_save is skipped.
118
self.connector_worker.prepare_store_kv(self._connector_metadata)
119
111
- def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
112
- assert self.connector_worker is not None
120
return self.connector_worker.get_finished(finished_req_ids)
121
122
def build_connector_worker_meta(self) -> OffloadingWorkerMetadata | None: