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 10, 21:30 UTC
implements register kv caches in lmcache connector (#31397)
Signed-off-by: idellzheng <idellzheng@tencent.com>
chunxiaozheng committed
Dec 30, 2025 at 03:13 UTC
b12cb38398a57a4abc4b8fb120fd39e97029ffa1
2 files changed
+26
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py
+16
@@ -107,6 +107,22 @@ class LMCacheConnectorV1(KVConnectorBase_V1):
107
# ==============================
108
# Worker-side methods
109
# ==============================
110
+ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
111
+ """
112
+ Initialize with the KV caches. Useful for pre-registering the
113
+ KV Caches in the KVConnector (e.g. for NIXL).
114
+
115
+ Args:
116
+ kv_caches: dictionary of layer names, kv cache
117
+ """
118
+ if hasattr(self._lmcache_engine, "register_kv_caches"):
119
+ self._lmcache_engine.register_kv_caches(kv_caches)
120
+ else:
121
+ logger.warning(
122
+ "LMCache engine does not support register_kv_caches, "
123
+ "please check and use the latest version"
124
+ )
125
+
126
def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:
127
"""
128
Start loading the KV cache from the connector to vLLM's paged
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py
+10
@@ -782,6 +782,16 @@ class LMCacheConnectorV1Impl:
782
####################
783
# Worker side APIs
784
####################
785
+ @_lmcache_nvtx_annotate
786
+ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
787
+ logger.info("Registering KV caches")
788
+ # TODO(chunxiaozheng): `_init_kv_caches_from_forward_context` is
789
+ # not called, we should consider removing it.
790
+ assert len(self.kv_caches) == 0 and len(kv_caches) > 0
791
+ self.kv_caches = kv_caches
792
+ if self.lmcache_engine is not None:
793
+ kvcaches = list(self.kv_caches.values())
794
+ self.lmcache_engine.post_init(kvcaches=kvcaches)
795
796
@_lmcache_nvtx_annotate
797
def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None: