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 5, 08:30 UTC
[KV Offload] Implement `shutdown()` in `OffloadingConnector` and related classes (#39182)
Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>
Ronen Schaffer committed
Apr 10, 2026 at 08:06 UTC
a5b17fba8ff3fcc076d73ba749a0819e0ec25f06
6 files changed
+39
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+3
@@ -428,3 +428,6 @@ class OffloadingConnectorScheduler:
428
medium=event.medium,
429
lora_name=None,
430
)
431
+
432
+ def shutdown(self) -> None:
433
+ self.manager.shutdown()
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+10
@@ -394,3 +394,13 @@ class OffloadingConnectorWorker:
394
kv_connector_stats = self.kv_connector_stats
395
self.kv_connector_stats = OffloadingConnectorStats()
396
return kv_connector_stats
397
+
398
+ def shutdown(self) -> None:
399
+ # Drop deferred store jobs: there is no point in submitting
400
+ # them during shutdown.
401
+ self._unsubmitted_store_jobs.clear()
402
+ self._jobs.clear()
403
+ self._load_job.clear()
404
+ self._store_jobs.clear()
405
+ self._finished_reqs_waiting_for_store.clear()
406
+ self.worker.shutdown()
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+6
@@ -64,6 +64,12 @@ class OffloadingConnector(KVConnectorBase_V1):
64
elif role == KVConnectorRole.WORKER:
65
self.connector_worker = OffloadingConnectorWorker(spec)
66
67
+ def shutdown(self) -> None:
68
+ if self.connector_worker is not None:
69
+ self.connector_worker.shutdown()
70
+ if self.connector_scheduler is not None:
71
+ self.connector_scheduler.shutdown()
72
+
73
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
74
assert self.connector_worker is not None
75
self.connector_worker.register_kv_caches(kv_caches)
vllm/v1/kv_offload/abstract.py
+4
@@ -178,3 +178,7 @@ class OffloadingManager(ABC):
178
New OffloadingEvents collected since the last call.
179
"""
180
return ()
181
+
182
+ def shutdown(self) -> None:
183
+ """Shutdown the manager and release any resources."""
184
+ return
vllm/v1/kv_offload/worker/cpu_gpu.py
+8
@@ -274,6 +274,14 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
274
if event is not None:
275
event.synchronize()
276
277
+ def shutdown(self) -> None:
278
+ while self._transfers:
279
+ transfer = self._transfers.popleft()
280
+ transfer.end_event.synchronize()
281
+ self._transfer_events.clear()
282
+ self._stream_pool.clear()
283
+ self._event_pool.clear()
284
+
285
286
class CpuGpuOffloadingHandlers:
287
def __init__(
vllm/v1/kv_offload/worker/worker.py
+8
@@ -69,6 +69,10 @@ class OffloadingHandler(ABC):
69
job_ids: The set of job IDs to wait for.
70
"""
71
72
+ def shutdown(self) -> None:
73
+ """Shutdown the handler and release any resources."""
74
+ return
75
+
76
77
class OffloadingWorker:
78
"""
@@ -166,3 +170,7 @@ class OffloadingWorker:
170
"""
171
for handler in self.handlers:
172
handler.wait(job_ids)
173
+
174
+ def shutdown(self) -> None:
175
+ for handler in self.handlers:
176
+ handler.shutdown()