From 8b1f3bebcab8f501e13183f84b20635cf336ccc1 Mon Sep 17 00:00:00 2001 From: Cao Qian Date: Mon, 20 Apr 2026 13:42:49 -0700 Subject: [PATCH] [LMCache MP Connector] Add num_lmcache_extra_cached_token in KVTransferParams (#39843) Signed-off-by: aeon-x --- .../kv_connector/v1/lmcache_mp_connector.py | 21 ++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py index c6d46b49af..f55f04a082 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py @@ -930,12 +930,31 @@ class LMCacheMPConnector(KVConnectorBase_V1): Optional KVTransferParams to be included in the request outputs returned by the engine. """ + + params: dict[str, Any] | None = getattr(request, "kv_transfer_params", None) + return_params: dict[str, Any] | None = {} if params is not None else None + + if ( + params is not None + and return_params is not None + and "num_lmcache_extra_cached_tokens" in params + ): + request_tracker = self._get_request_tracker(request.request_id) + num_extra_cached_blocks = max( + 0, + request_tracker.num_lmcache_hit_blocks + - request_tracker.num_vllm_hit_blocks, + ) + return_params["num_lmcache_extra_cached_tokens"] = ( + num_extra_cached_blocks * self.vllm_block_size + ) + # Clean up request tracker to prevent memory leak self._cleanup_request_tracker(request.request_id) # Notify LMCache to end the session for this request self.scheduler_adapter.end_session(request.request_id) - return True, None + return True, return_params def take_events(self) -> Iterable["KVCacheEvent"]: """