From d6bb17f10e993b830153deb6a5a6bf628c03bf82 Mon Sep 17 00:00:00 2001 From: UnifiedCacheManager Date: Wed, 7 Jan 2026 11:30:33 +0800 Subject: [PATCH] [Bugfix]Add register_kv_cache in ucm_connector (#5657) ### What this PR does / why we need it? To adapt different shapes of the KV cache, UCM optimized the initialization of store by moving it into `register_kv_caches`. Therefore, this update adds `register_kv_caches` interface to UCMConnectorV1. ### How was this patch tested? - vLLM version: v0.13.0 - vLLM main: https://github.com/vllm-project/vllm/commit/2f4e6548efec402b913ffddc8726230d9311948d Signed-off-by: UnifiedCacheManager --- vllm_ascend/distributed/ucm_connector.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/vllm_ascend/distributed/ucm_connector.py b/vllm_ascend/distributed/ucm_connector.py index f44e5ee2..d38b6519 100644 --- a/vllm_ascend/distributed/ucm_connector.py +++ b/vllm_ascend/distributed/ucm_connector.py @@ -40,6 +40,15 @@ class UCMConnectorV1(KVConnectorBase_V1): # ============================== # Worker-side methods # ============================== + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: + """ + Initialize with the KV caches. Useful for pre-registering the + KV Caches in the KVConnector (e.g. for NIXL). + Args: + kv_caches: A dictionary mapping layer names to KV cache tensors. + """ + self._ucm_engine.register_kv_caches(kv_caches) + def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None: """