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 11, 05:00 UTC
[BugFix]Reclaim resources to prevent memory leaks when use LMCacheMPConnector (#30745)
Signed-off-by: wz1qqx <ziqi.wang@novita.ai> Co-authored-by: wz1qqx <ziqi.wang@novita.ai>
wz1qqx committed
Dec 19, 2025 at 03:09 UTC
889f8bb250d498dcd38d7bcf58fb3c9e50d54d14
2 files changed
+24
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py
+8
@@ -147,6 +147,14 @@ class LMCacheMPSchedulerAdapter:
147
"""
148
return self.blocks_in_chunk
149
150
+ def _cleanup_lookup_result(self, request_id: str) -> None:
151
+ """
152
+ Clean up lookup future for a finished request to prevent memory leak.
153
+ Args:
154
+ request_id: The ID of the finished request.
155
+ """
156
+ self.lookup_futures.pop(request_id, None)
157
+
158
# Helper functions
159
def _create_key(self, block_hash: bytes) -> IPCCacheEngineKey:
160
"""Convert a block hash to an IPC cache engine key"""
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py
+16
@@ -701,6 +701,8 @@ class LMCacheMPConnector(KVConnectorBase_V1):
701
if condition
702
else LMCacheMPRequestState.READY
703
)
704
+ # Clean up lookup future in scheduler adapter
705
+ self.scheduler_adapter._cleanup_lookup_result(request.request_id)
706
707
def build_connector_meta(
708
self, scheduler_output: SchedulerOutput
@@ -754,6 +756,8 @@ class LMCacheMPConnector(KVConnectorBase_V1):
756
Optional KVTransferParams to be included in the request outputs
757
returned by the engine.
758
"""
759
+ # Clean up request tracker to prevent memory leak
760
+ self._cleanup_request_tracker(request.request_id)
761
return True, None
762
763
def take_events(self) -> Iterable["KVCacheEvent"]:
@@ -915,3 +919,15 @@ class LMCacheMPConnector(KVConnectorBase_V1):
919
new_tracker = LMCacheMPRequestTracker(request)
920
self.request_trackers[request_id] = new_tracker
921
return self.request_trackers[request_id]
922
+
923
+ def _cleanup_request_tracker(self, request_id: str) -> None:
924
+ """
925
+ Clean up request tracker and associated lookup future for a request.
926
+ This should be called when a request is finished to prevent memory leak.
927
+ """
928
+ # Clean up request tracker
929
+ if self.request_trackers.pop(request_id, None):
930
+ logger.debug(
931
+ "[KVConnector] Cleaned up request_tracker for request %s",
932
+ request_id,
933
+ )