init v0.23.0

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>

vllm_ascend/distributed/__init__.py

@@ -1,33 +0,0 @@
-#
-# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
-# This file is a part of the vllm-ascend project.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-
-from vllm.distributed.kv_transfer.kv_connector.factory import \
-    KVConnectorFactory
-
-KVConnectorFactory.register_connector(
-    "LLMDataDistCMgrConnector",
-    "vllm_ascend.distributed.llmdatadist_c_mgr_connector",
-    "LLMDataDistCMgrConnector")
-
-KVConnectorFactory.register_connector(
-    "MooncakeConnectorV1", "vllm_ascend.distributed.mooncake_connector",
-    "MooncakeConnector")
-
-KVConnectorFactory.register_connector(
-    "MooncakeConnectorStoreV1",
-    "vllm_ascend.distributed.mooncake.mooncake_store_connector_v1",
-    "MooncakeConnectorV1")

vllm_ascend/distributed/device_communicators/npu_communicator.py

@@ -0,0 +1,68 @@
+#
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+import torch
+import torch.distributed as dist
+from vllm.distributed.device_communicators.base_device_communicator import DeviceCommunicatorBase
+
+
+class NPUCommunicator(DeviceCommunicatorBase):
+    def __init__(
+        self,
+        cpu_group: dist.ProcessGroup,
+        device: torch.device | None = None,
+        device_group: dist.ProcessGroup | None = None,
+        unique_name: str = "",
+    ):
+        super().__init__(cpu_group, device, device_group, unique_name)
+        # TODO(hz): Refer to CudaCommunicator's implementation to integrate PyHcclCommunicator
+        # init device according to rank
+        self.device = torch.npu.current_device()
+
+        # For compatibility (mainly for reusing graph capturing code in vllm),
+        # init custom all-reduce implementation interface as in CUDACommunicator.
+        self.ca_comm = None
+
+    def all_to_all(
+        self,
+        input_: torch.Tensor,
+        scatter_dim: int = 0,
+        gather_dim: int = -1,
+        scatter_sizes: list[int] | None = None,
+        gather_sizes: list[int] | None = None,
+    ) -> torch.Tensor:
+        if scatter_dim < 0:
+            scatter_dim += input_.dim()
+        if gather_dim < 0:
+            gather_dim += input_.dim()
+
+        if scatter_sizes is not None and gather_sizes is not None:
+            input_list = [t.contiguous() for t in torch.split(input_, scatter_sizes, scatter_dim)]
+            output_list = []
+            tensor_shape_base = input_list[self.rank].size()
+            for i in range(self.world_size):
+                tensor_shape = list(tensor_shape_base)
+                tensor_shape[gather_dim] = gather_sizes[i]
+                output_list.append(torch.empty(tensor_shape, dtype=input_.dtype, device=input_.device))
+
+        else:
+            input_list = [t.contiguous() for t in torch.tensor_split(input_, self.world_size, scatter_dim)]
+            output_list = [torch.empty_like(input_list[i]) for i in range(self.world_size)]
+
+        dist.all_to_all(output_list, input_list, group=self.device_group)
+        output_tensor = torch.cat(output_list, dim=gather_dim).contiguous()
+        return output_tensor

vllm_ascend/distributed/device_communicators/pyhccl.py

@@ -15,7 +15,6 @@
 # limitations under the License.
 #
 
-from typing import Optional, Union
 
 import torch
 import torch.distributed as dist
@@ -24,18 +23,23 @@ from vllm.distributed.utils import StatelessProcessGroup
 from vllm.logger import logger
 
 from vllm_ascend.distributed.device_communicators.pyhccl_wrapper import (
-    HCCLLibrary, aclrtStream_t, buffer_type, hcclComm_t, hcclDataTypeEnum,
-    hcclRedOpTypeEnum, hcclUniqueId)
+    HCCLLibrary,
+    aclrtStream_t,
+    buffer_type,
+    hcclComm_t,
+    hcclDataTypeEnum,
+    hcclRedOpTypeEnum,
+    hcclUniqueId,
+)
 from vllm_ascend.utils import current_stream
 
 
 class PyHcclCommunicator:
-
     def __init__(
         self,
-        group: Union[ProcessGroup, StatelessProcessGroup],
-        device: Union[int, str, torch.device],
-        library_path: Optional[str] = None,
+        group: ProcessGroup | StatelessProcessGroup,
+        device: int | str | torch.device,
+        library_path: str | None = None,
     ):
         """
         Args:
@@ -52,7 +56,8 @@ class PyHcclCommunicator:
         if not isinstance(group, StatelessProcessGroup):
             assert dist.is_initialized()
             assert dist.get_backend(group) != dist.Backend.HCCL, (
-                "PyHcclCommunicator should be attached to a non-HCCL group.")
+                "PyHcclCommunicator should be attached to a non-HCCL group."
+            )
             # note: this rank is the rank in the group
             self.rank = dist.get_rank(group)
             self.world_size = dist.get_world_size(group)
@@ -113,8 +118,7 @@ class PyHcclCommunicator:
         # `torch.npu.device` is a context manager that changes the
         # current npu device to the specified one
         with torch.npu.device(device):
-            self.comm: hcclComm_t = self.hccl.hcclCommInitRank(
-                self.world_size, self.unique_id, self.rank)
+            self.comm: hcclComm_t = self.hccl.hcclCommInitRank(self.world_size, self.unique_id, self.rank)
 
             stream = current_stream()
             # A small all_reduce for warmup.
@@ -123,43 +127,45 @@ class PyHcclCommunicator:
             stream.synchronize()
             del data
 
-    def all_reduce(self,
-                   in_tensor: torch.Tensor,
-                   op: ReduceOp = ReduceOp.SUM,
-                   stream=None) -> torch.Tensor:
+    def all_reduce(self, in_tensor: torch.Tensor, op: ReduceOp = ReduceOp.SUM, stream=None) -> torch.Tensor:
         if self.disabled:
             return None
         # hccl communicator created on a specific device
         # will only work on tensors on the same device
         # otherwise it will cause "illegal memory access"
         assert in_tensor.device == self.device, (
-            f"this hccl communicator is created to work on {self.device}, "
-            f"but the input tensor is on {in_tensor.device}")
+            f"this hccl communicator is created to work on {self.device}, but the input tensor is on {in_tensor.device}"
+        )
 
         out_tensor = torch.empty_like(in_tensor)
 
         if stream is None:
             stream = current_stream()
-        self.hccl.hcclAllReduce(buffer_type(in_tensor.data_ptr()),
-                                buffer_type(out_tensor.data_ptr()),
-                                in_tensor.numel(),
-                                hcclDataTypeEnum.from_torch(in_tensor.dtype),
-                                hcclRedOpTypeEnum.from_torch(op), self.comm,
-                                aclrtStream_t(stream.npu_stream))
+        self.hccl.hcclAllReduce(
+            buffer_type(in_tensor.data_ptr()),
+            buffer_type(out_tensor.data_ptr()),
+            in_tensor.numel(),
+            hcclDataTypeEnum.from_torch(in_tensor.dtype),
+            hcclRedOpTypeEnum.from_torch(op),
+            self.comm,
+            aclrtStream_t(stream.npu_stream),
+        )
         return out_tensor
 
     def broadcast(self, tensor: torch.Tensor, src: int, stream=None):
         if self.disabled:
             return
         assert tensor.device == self.device, (
-            f"this hccl communicator is created to work on {self.device}, "
-            f"but the input tensor is on {tensor.device}")
+            f"this hccl communicator is created to work on {self.device}, but the input tensor is on {tensor.device}"
+        )
         if stream is None:
             stream = current_stream()
-        if src == self.rank:
-            buffer = buffer_type(tensor.data_ptr())
-        else:
-            buffer = buffer_type(tensor.data_ptr())
-        self.hccl.hcclBroadcast(buffer, tensor.numel(),
-                                hcclDataTypeEnum.from_torch(tensor.dtype), src,
-                                self.comm, aclrtStream_t(stream.npu_stream))
+        buffer = buffer_type(tensor.data_ptr())
+        self.hccl.hcclBroadcast(
+            buffer,
+            tensor.numel(),
+            hcclDataTypeEnum.from_torch(tensor.dtype),
+            src,
+            self.comm,
+            aclrtStream_t(stream.npu_stream),
+        )

vllm_ascend/distributed/device_communicators/pyhccl_wrapper.py

@@ -18,7 +18,7 @@
 import ctypes
 import platform
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 import torch
 from torch.distributed import ReduceOp
@@ -107,69 +107,74 @@ class hcclRedOpTypeEnum:
 class Function:
     name: str
     restype: Any
-    argtypes: List[Any]
+    argtypes: list[Any]
 
 
 class HCCLLibrary:
     exported_functions = [
         # const char* HcclGetErrorString(HcclResult code);
         Function("HcclGetErrorString", ctypes.c_char_p, [hcclResult_t]),
-
         # HcclResult HcclGetRootInfo(HcclRootInfo *rootInfo);
-        Function("HcclGetRootInfo", hcclResult_t,
-                 [ctypes.POINTER(hcclUniqueId)]),
-
+        Function("HcclGetRootInfo", hcclResult_t, [ctypes.POINTER(hcclUniqueId)]),
         # HcclResult HcclCommInitRootInfo(
         #   uint32_t nRanks, const HcclRootInfo *rootInfo, uint32_t rank, HcclComm *comm);
         # note that HcclComm is a pointer type, so the last argument is a pointer to a pointer
-        Function("HcclCommInitRootInfo", hcclResult_t, [
-            ctypes.c_int,
-            ctypes.POINTER(hcclUniqueId),
-            ctypes.c_int,
-            ctypes.POINTER(hcclComm_t),
-        ]),
-
+        Function(
+            "HcclCommInitRootInfo",
+            hcclResult_t,
+            [
+                ctypes.c_int,
+                ctypes.POINTER(hcclUniqueId),
+                ctypes.c_int,
+                ctypes.POINTER(hcclComm_t),
+            ],
+        ),
         # HcclResult HcclAllReduce(
         #   void *sendBuf, void *recvBuf, uint64_t count,
         #   HcclDataType dataType, HcclReduceOp op, HcclComm comm,
         #   aclrtStream stream);
-        Function("HcclAllReduce", hcclResult_t, [
-            buffer_type,
-            buffer_type,
-            ctypes.c_size_t,
-            hcclDataType_t,
-            hcclRedOp_t,
-            hcclComm_t,
-            aclrtStream_t,
-        ]),
-
+        Function(
+            "HcclAllReduce",
+            hcclResult_t,
+            [
+                buffer_type,
+                buffer_type,
+                ctypes.c_size_t,
+                hcclDataType_t,
+                hcclRedOp_t,
+                hcclComm_t,
+                aclrtStream_t,
+            ],
+        ),
         # HcclResult HcclBroadcast(
         #   void *buf, uint64_t count,
         #   HcclDataType dataType, uint32_t root,
         #   HcclComm comm, aclrtStream stream);
-        Function("HcclBroadcast", hcclResult_t, [
-            buffer_type,
-            ctypes.c_size_t,
-            hcclDataType_t,
-            ctypes.c_int,
-            hcclComm_t,
-            aclrtStream_t,
-        ]),
-
+        Function(
+            "HcclBroadcast",
+            hcclResult_t,
+            [
+                buffer_type,
+                ctypes.c_size_t,
+                hcclDataType_t,
+                ctypes.c_int,
+                hcclComm_t,
+                aclrtStream_t,
+            ],
+        ),
         # HcclResult HcclCommDestroy(HcclComm comm);
         Function("HcclCommDestroy", hcclResult_t, [hcclComm_t]),
     ]
 
     # class attribute to store the mapping from the path to the library
     # to avoid loading the same library multiple times
-    path_to_library_cache: Dict[str, Any] = {}
+    path_to_library_cache: dict[str, Any] = {}
 
     # class attribute to store the mapping from library path
-    # to the correspongding directory
-    path_to_dict_mapping: Dict[str, Dict[str, Any]] = {}
-
-    def __init__(self, so_file: Optional[str] = None):
+    # to the corresponding directory
+    path_to_dict_mapping: dict[str, dict[str, Any]] = {}
 
+    def __init__(self, so_file: str | None = None):
         so_file = so_file or find_hccl_library()
 
         try:
@@ -179,18 +184,21 @@ class HCCLLibrary:
             self.lib = HCCLLibrary.path_to_library_cache[so_file]
         except Exception as e:
             logger.error(
-                "Failed to load HCCL library from %s. "
-                "It is expected if you are not running on Ascend NPUs."
-                "Otherwise, the hccl library might not exist, be corrupted "
+                "Failed to load HCCL library. "
+                "so_file=%s, error=%s. "
+                "The hccl library might not exist, be corrupted "
                 "or it does not support the current platform %s. "
                 "If you already have the library, please set the "
                 "environment variable HCCL_SO_PATH"
-                " to point to the correct hccl library path.", so_file,
-                platform.platform())
+                " to point to the correct hccl library path.",
+                so_file,
+                e,
+                platform.platform(),
+            )
             raise e
 
         if so_file not in HCCLLibrary.path_to_dict_mapping:
-            _funcs: Dict[str, Any] = {}
+            _funcs: dict[str, Any] = {}
             for func in HCCLLibrary.exported_functions:
                 f = getattr(self.lib, func.name)
                 f.restype = func.restype
@@ -209,34 +217,37 @@ class HCCLLibrary:
 
     def hcclGetUniqueId(self) -> hcclUniqueId:
         unique_id = hcclUniqueId()
-        self.HCCL_CHECK(self._funcs["HcclGetRootInfo"](
-            ctypes.byref(unique_id)))
+        self.HCCL_CHECK(self._funcs["HcclGetRootInfo"](ctypes.byref(unique_id)))
         return unique_id
 
-    def hcclCommInitRank(self, world_size: int, unique_id: hcclUniqueId,
-                         rank: int) -> hcclComm_t:
+    def hcclCommInitRank(self, world_size: int, unique_id: hcclUniqueId, rank: int) -> hcclComm_t:
         comm = hcclComm_t()
-        self.HCCL_CHECK(self._funcs["HcclCommInitRootInfo"](
-            world_size, ctypes.byref(unique_id), rank, ctypes.byref(comm)))
+        self.HCCL_CHECK(
+            self._funcs["HcclCommInitRootInfo"](world_size, ctypes.byref(unique_id), rank, ctypes.byref(comm))
+        )
         return comm
 
-    def hcclAllReduce(self, sendbuff: buffer_type, recvbuff: buffer_type,
-                      count: int, datatype: int, op: int, comm: hcclComm_t,
-                      stream: aclrtStream_t) -> None:
+    def hcclAllReduce(
+        self,
+        sendbuff: buffer_type,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        op: int,
+        comm: hcclComm_t,
+        stream: aclrtStream_t,
+    ) -> None:
         # `datatype` actually should be `hcclDataType_t`
         # and `op` should be `hcclRedOp_t`
         # both are aliases of `ctypes.c_int`
         # when we pass int to a function, it will be converted to `ctypes.c_int`
         # by ctypes automatically
-        self.HCCL_CHECK(self._funcs["HcclAllReduce"](sendbuff, recvbuff, count,
-                                                     datatype, op, comm,
-                                                     stream))
+        self.HCCL_CHECK(self._funcs["HcclAllReduce"](sendbuff, recvbuff, count, datatype, op, comm, stream))
 
-    def hcclBroadcast(self, buf: buffer_type, count: int, datatype: int,
-                      root: int, comm: hcclComm_t,
-                      stream: aclrtStream_t) -> None:
-        self.HCCL_CHECK(self._funcs["HcclBroadcast"](buf, count, datatype,
-                                                     root, comm, stream))
+    def hcclBroadcast(
+        self, buf: buffer_type, count: int, datatype: int, root: int, comm: hcclComm_t, stream: aclrtStream_t
+    ) -> None:
+        self.HCCL_CHECK(self._funcs["HcclBroadcast"](buf, count, datatype, root, comm, stream))
 
     def hcclCommDestroy(self, comm: hcclComm_t) -> None:
         self.HCCL_CHECK(self._funcs["HcclCommDestroy"](comm))

vllm_ascend/distributed/kv_transfer/__init__.py

@@ -0,0 +1,87 @@
+#
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm-ascend project.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
+
+
+def register_connector():
+    # override multi_connector as ascend_multi_connector
+    if "MultiConnector" in KVConnectorFactory._registry:
+        KVConnectorFactory._registry.pop("MultiConnector")
+    KVConnectorFactory.register_connector(
+        "MultiConnector", "vllm_ascend.distributed.kv_transfer.ascend_multi_connector", "AscendMultiConnector"
+    )
+
+    KVConnectorFactory.register_connector(
+        "MooncakeConnectorV1", "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake_connector", "MooncakeConnector"
+    )
+
+    KVConnectorFactory.register_connector(
+        "MooncakeHybridConnector",
+        "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake_hybrid_connector",
+        "MooncakeConnector",
+    )
+
+    KVConnectorFactory.register_connector(
+        "MooncakeConnectorStoreV1",
+        "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector",
+        "AscendStoreConnector",
+    )
+
+    KVConnectorFactory.register_connector(
+        "AscendStoreConnector",
+        "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector",
+        "AscendStoreConnector",
+    )
+
+    KVConnectorFactory.register_connector(
+        "MooncakeLayerwiseConnector",
+        "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake_layerwise_connector",
+        "MooncakeLayerwiseConnector",
+    )
+
+    KVConnectorFactory.register_connector(
+        "UCMConnector", "vllm_ascend.distributed.kv_transfer.kv_pool.ucm_connector", "UCMConnectorV1"
+    )
+
+    KVConnectorFactory.register_connector(
+        "LMCacheAscendConnector",
+        "vllm_ascend.distributed.kv_transfer.kv_pool.lmcache_ascend_connector",
+        "LMCacheConnectorV1",
+    )
+
+    # Override the upstream SimpleCPUOffloadConnector with the NPU
+    # adaptation that uses aclrtMemcpyBatchAsync + torch.npu streams.
+    # Only override if the upstream module exists in this vLLM version.
+    try:
+        import vllm.v1.simple_kv_offload  # noqa: F401
+    except ImportError:
+        pass
+    else:
+        if "SimpleCPUOffloadConnector" in KVConnectorFactory._registry:
+            KVConnectorFactory._registry.pop("SimpleCPUOffloadConnector")
+        KVConnectorFactory.register_connector(
+            "SimpleCPUOffloadConnector",
+            "vllm_ascend.distributed.kv_transfer.kv_pool.simple_cpu_offload.simple_cpu_offload_connector",  # noqa: E501
+            "AscendSimpleCPUOffloadConnector",
+        )
+
+    KVConnectorFactory.register_connector(
+        "RecomputeCPUOffloadConnector",
+        "vllm_ascend.distributed.kv_transfer.kv_pool.recompute_cpu_offload.recompute_cpu_offload_connector",
+        "RecomputeCPUOffloadConnectorV1",
+    )

vllm_ascend/distributed/kv_transfer/ascend_multi_connector.py

@@ -0,0 +1,102 @@
+from typing import TYPE_CHECKING, Any, cast
+
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorRole,
+    SupportsHMA,
+    supports_hma,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector
+
+from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake_layerwise_connector import MooncakeLayerwiseConnector
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.request import Request
+
+
+class AscendMultiConnector(MultiConnector, SupportsHMA):
+    def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole, kv_cache_config: "KVCacheConfig"):
+        super().__init__(
+            vllm_config=vllm_config,
+            role=role,
+            kv_cache_config=kv_cache_config,
+        )
+
+        self._all_support_hma = all(supports_hma(c) for c in self._connectors)
+        assert vllm_config.scheduler_config.disable_hybrid_kv_cache_manager or self._all_support_hma, (
+            "HMA should not be enabled unless all sub-connectors support it"
+        )
+
+    def update_state_after_alloc(self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int):
+        chosen_connector = self._requests_to_connector.get(request.request_id, -1)
+        empty_blocks = blocks.new_empty()
+        for i, c in enumerate(self._connectors):
+            if i == chosen_connector or isinstance(c, MooncakeLayerwiseConnector):
+                # Forward call to the chosen connector (if any).
+                c.update_state_after_alloc(request, blocks, num_external_tokens)
+            else:
+                # Call with empty blocks for other connectors.
+                c.update_state_after_alloc(request, empty_blocks, 0)
+
+    def get_num_new_matched_tokens(
+        self,
+        request: "Request",
+        num_computed_tokens: int,
+    ) -> tuple[int | None, bool]:
+        # Recompute offload may contain an unhashed partial block that other
+        # prefix-cache connectors cannot restore. Give its request state
+        # priority regardless of connector ordering.
+        for i, connector in enumerate(self._connectors):
+            has_preempted_request = getattr(connector, "has_preempted_request", None)
+            if has_preempted_request is None or not has_preempted_request(request.request_id):
+                continue
+            tokens, load_async = connector.get_num_new_matched_tokens(request, num_computed_tokens)
+            if tokens is None:
+                return None, False
+            if tokens > 0:
+                self._requests_to_connector[request.request_id] = i
+                return tokens, load_async
+            break
+
+        return super().get_num_new_matched_tokens(request, num_computed_tokens)
+
+    def update_state_before_preempt(
+        self,
+        request: "Request",
+        block_ids: tuple[list[int], ...],
+        num_computed_tokens: int,
+    ) -> bool:
+        offloaded = False
+        for c in self._connectors:
+            hook = getattr(c, "update_state_before_preempt", None)
+            if hook is not None:
+                offloaded = bool(hook(request, block_ids, num_computed_tokens)) or offloaded
+        return offloaded
+
+    def request_finished_all_groups(
+        self,
+        request: "Request",
+        block_ids: tuple[list[int], ...],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        if not self._all_support_hma:
+            assert len(block_ids) == 1, "HMA with multiple kv_cache_groups requires all sub-connectors to support HMA"
+            return super().request_finished(request, block_ids[0])
+
+        async_saves = 0
+        kv_txfer_params = None
+        for c in self._connectors:
+            async_save, txfer_params = cast(SupportsHMA, c).request_finished_all_groups(request, block_ids)
+            if async_save:
+                async_saves += 1
+            if txfer_params is not None:
+                if kv_txfer_params is not None:
+                    raise RuntimeError("Only one connector can produce KV transfer params")
+                kv_txfer_params = txfer_params
+        if async_saves > 1:
+            self._extra_async_saves[request.request_id] = async_saves - 1
+
+        self._requests_to_connector.pop(request.request_id, None)
+
+        return async_saves > 0, kv_txfer_params

vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/ascend_store_connector.py

@@ -0,0 +1,339 @@
+import threading
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Any
+
+import torch
+import zmq
+from vllm.config import VllmConfig
+from vllm.distributed.kv_events import (
+    KVCacheEvent,
+    KVConnectorKVEvents,
+    KVEventAggregator,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorBase_V1,
+    KVConnectorMetadata,
+    KVConnectorRole,
+    SupportsHMA,
+)
+from vllm.forward_context import ForwardContext
+from vllm.logger import logger
+from vllm.utils.network_utils import make_zmq_socket
+from vllm.v1.attention.backend import AttentionMetadata  # type: ignore
+from vllm.v1.core.block_pool import BlockPool
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.outputs import KVConnectorOutput
+from vllm.v1.request import Request
+from vllm.v1.serial_utils import MsgpackDecoder
+
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.config_data import AscendStoreKVConnectorWorkerMetadata
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.pool_scheduler import (
+    KVPoolScheduler,
+    get_zmq_rpc_path_lookup,
+)
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.pool_worker import KVPoolWorker
+
+if TYPE_CHECKING:
+    from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorHandshakeMetadata
+
+
+class AscendStoreKVEvents(KVConnectorKVEvents):
+    def __init__(self, num_workers: int) -> None:
+        self._aggregator = KVEventAggregator(num_workers)
+
+    def add_events(self, events: list[KVCacheEvent]) -> None:
+        self._aggregator.add_events(events)
+
+    def aggregate(self) -> "AscendStoreKVEvents":
+        """
+        Aggregate KV events and retain only common events.
+        """
+        common_events = self._aggregator.get_common_events()
+        self._aggregator.clear_events()
+        self._aggregator.add_events(common_events)
+        self._aggregator.reset_workers()
+        return self
+
+    def increment_workers(self, count: int = 1) -> None:
+        self._aggregator.increment_workers(count)
+
+    def get_all_events(self) -> list[KVCacheEvent]:
+        return self._aggregator.get_all_events()
+
+    def get_number_of_workers(self) -> int:
+        return self._aggregator.get_number_of_workers()
+
+    def clear_events(self) -> None:
+        self._aggregator.clear_events()
+        self._aggregator.reset_workers()
+
+    def __repr__(self) -> str:
+        return f"<AscendStoreKVEvents events={self.get_all_events()}>"
+
+
+class AscendStoreConnector(KVConnectorBase_V1, SupportsHMA):
+    @classmethod
+    def requires_piecewise_for_cudagraph(cls, extra_config: dict[str, Any]) -> bool:
+        """
+        AscendStore requires PIECEWISE CUDA graph mode when layerwise
+        operations are enabled.
+        """
+        return extra_config.get("use_layerwise", False)
+
+    def __init__(self, vllm_config: VllmConfig, role: KVConnectorRole, kv_cache_config: KVCacheConfig | None = None):
+        super().__init__(vllm_config=vllm_config, role=role, kv_cache_config=kv_cache_config)
+        self.kv_role = vllm_config.kv_transfer_config.kv_role
+
+        self.use_layerwise = vllm_config.kv_transfer_config.kv_connector_extra_config.get("use_layerwise", False)
+        backend_name = vllm_config.kv_transfer_config.kv_connector_extra_config.get("backend", "mooncake")
+        self.backend_name = backend_name.lower()
+        self.use_gva_layerwise = self.use_layerwise and self.backend_name == "memcache"
+        self.consumer_is_to_put = vllm_config.kv_transfer_config.kv_connector_extra_config.get(
+            "consumer_is_to_put", False
+        )
+
+        connector_name = vllm_config.kv_transfer_config.kv_connector
+        if connector_name == "MooncakeConnectorStoreV1":
+            logger.warning(
+                "It is recommended to use the AscendStoreConnector, "
+                "as the MoonCakeStoreConnector will be removed in the future."
+            )
+
+        self.kv_caches: dict[str, torch.Tensor] = {}
+        self._kv_cache_events: AscendStoreKVEvents | None = None
+
+        self._current_step_has_real_forward = False
+
+        if role == KVConnectorRole.SCHEDULER:
+            assert kv_cache_config is not None
+            page_size_bytes = kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes
+            self.connector_scheduler = KVPoolScheduler(
+                vllm_config, self.use_layerwise, kv_cache_config, page_size_bytes=page_size_bytes
+            )
+        else:
+            self.connector_worker = KVPoolWorker(
+                vllm_config,
+                self.use_layerwise,
+                kv_cache_config,
+            )
+            assert self.connector_worker is not None
+            if not self.use_layerwise and vllm_config.parallel_config.rank == 0:
+                self.lookup_server = LookupKeyServer(self.connector_worker, vllm_config)
+
+    ############################################################
+    # Scheduler Side Methods
+    ############################################################
+
+    def set_xfer_handshake_metadata_pp_aware(
+        self,
+        metadata: dict[tuple[int, int], "KVConnectorHandshakeMetadata"],
+    ) -> None:
+        """Ignore P/D handshake metadata because AscendStore handles PP via pool keys."""
+        pass
+
+    def get_num_new_matched_tokens(self, request: "Request", num_computed_tokens: int) -> tuple[int, bool]:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.get_num_new_matched_tokens(request, num_computed_tokens)
+
+    def update_state_after_alloc(self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int):
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.update_state_after_alloc(request, blocks, num_external_tokens)
+
+    def build_connector_meta(
+        self,
+        scheduler_output: SchedulerOutput,
+    ) -> KVConnectorMetadata:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.build_connector_meta(scheduler_output)
+
+    def request_finished(
+        self,
+        request: "Request",
+        block_ids: list[int],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.request_finished(request, block_ids)
+
+    def request_finished_all_groups(
+        self,
+        request: "Request",
+        block_ids: tuple[list[int], ...],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.request_finished_all_groups(request, block_ids)
+
+    def update_connector_output(self, connector_output: KVConnectorOutput):
+        """
+        Update KVConnector state from worker-side connectors output.
+
+        Args:
+            connector_output (KVConnectorOutput): the worker-side connectors output.
+        """
+        if self.connector_scheduler is not None:
+            self.connector_scheduler.update_connector_output(connector_output)
+
+        # Get the KV events
+        kv_cache_events = connector_output.kv_cache_events
+        if not kv_cache_events or not isinstance(kv_cache_events, AscendStoreKVEvents):
+            return
+
+        if self._kv_cache_events is None:
+            self._kv_cache_events = kv_cache_events
+        else:
+            self._kv_cache_events.add_events(kv_cache_events.get_all_events())
+            self._kv_cache_events.increment_workers(kv_cache_events.get_number_of_workers())
+        return
+
+    def take_events(self) -> Iterable["KVCacheEvent"]:
+        """
+        Take the KV cache events from the connector.
+
+        Yields:
+            New KV cache events since the last call.
+        """
+        if self._kv_cache_events is not None:
+            self._kv_cache_events.aggregate()
+            kv_cache_events = self._kv_cache_events.get_all_events()
+            yield from kv_cache_events
+            self._kv_cache_events.clear_events()
+            self._kv_cache_events = None
+
+    ############################################################
+    # Worker Side Methods
+    ############################################################
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+        assert self.connector_worker is not None
+        self.connector_worker.register_kv_caches(kv_caches)
+
+    def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
+        assert self.connector_worker is not None
+        metadata = self._get_connector_metadata()
+        self._current_step_has_real_forward = forward_context is not None
+        logger.debug(
+            "KV pool connector start_load_kv metadata_requests=%d specs=%s",
+            len(metadata.requests),
+            [
+                (
+                    request.req_id,
+                    None if request.load_spec is None else request.load_spec.can_load,
+                    None if request.load_spec is None else request.load_spec.vllm_cached_tokens,
+                    None if request.load_spec is None else request.load_spec.kvpool_cached_tokens,
+                )
+                for request in metadata.requests
+            ],
+        )
+        self.connector_worker.start_load_kv(metadata)
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        if not self.use_layerwise:
+            return
+        self.connector_worker.wait_for_layer_load()
+
+    def save_kv_layer(
+        self, layer_name: str, kv_layer: torch.Tensor, attn_metadata: "AttentionMetadata", **kwargs
+    ) -> None:
+        if not self.use_layerwise:
+            return
+
+        if self.kv_role == "kv_consumer":
+            # Don't do save if the role is kv_consumer
+            return
+        self.connector_worker.save_kv_layer(self._get_connector_metadata())
+
+    def wait_for_save(self):
+        if self.kv_role == "kv_consumer" and not self.consumer_is_to_put:
+            # Don't do save if the role is kv_consumer
+            return
+
+        if self.use_layerwise:
+            return
+
+        self.connector_worker.wait_for_save(self._get_connector_metadata())
+
+    def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
+        """Get the finished recving and sending requests."""
+        assert self.connector_worker is not None
+        metadata = self._get_connector_metadata()
+        if self._current_step_has_real_forward:
+            try:
+                self.connector_worker.ensure_store_initialized()
+            finally:
+                self._current_step_has_real_forward = False
+        done_sending, done_recving = self.connector_worker.get_finished(finished_req_ids, metadata)
+        return done_sending, done_recving
+
+    def get_block_ids_with_load_errors(self) -> set[int]:
+        """Return KV block IDs that failed to load on the worker."""
+        assert self.connector_worker is not None
+        return self.connector_worker.get_block_ids_with_load_errors()
+
+    def get_kv_connector_kv_cache_events(self) -> AscendStoreKVEvents | None:
+        """
+        Get the KV connector kv cache events collected during the last interval.
+        """
+        events = self.connector_worker.get_kv_events()
+        if not events:
+            return None
+
+        ascend_store_kv_events = AscendStoreKVEvents(num_workers=1)
+        ascend_store_kv_events.add_events(events)
+        return ascend_store_kv_events
+
+    def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None:
+        assert self.connector_scheduler is not None
+        self.connector_scheduler.bind_gpu_block_pool(gpu_block_pool)
+
+    def build_connector_worker_meta(self) -> AscendStoreKVConnectorWorkerMetadata | None:
+        assert self.connector_worker is not None
+        return self.connector_worker.build_connector_worker_meta()
+
+
+class LookupKeyServer:
+    def __init__(
+        self,
+        pool_worker: KVPoolWorker,
+        vllm_config: "VllmConfig",
+    ):
+        self.decoder = MsgpackDecoder()
+        self.ctx = zmq.Context()  # type: ignore[attr-defined]
+        socket_path = get_zmq_rpc_path_lookup(vllm_config)
+        self.socket = make_zmq_socket(
+            self.ctx,
+            socket_path,
+            zmq.REP,  # type: ignore[attr-defined]
+            bind=True,
+        )
+
+        self.pool_worker = pool_worker
+        self.running = True
+
+        def process_request():
+            while self.running:
+                all_frames = self.socket.recv_multipart(copy=False)
+                token_len = int.from_bytes(all_frames[0], byteorder="big")
+                kv_group_ids = self.decoder.decode([all_frames[1]])
+                hbm_hit_tokens = int.from_bytes(all_frames[2], byteorder="big")
+                hashes_str = self.decoder.decode(all_frames[3:])
+                result = self.pool_worker.lookup_scheduler(
+                    token_len,
+                    hashes_str,
+                    kv_group_ids,
+                    use_layerwise=False,
+                    hbm_hit_tokens=hbm_hit_tokens,
+                )
+                logger.debug(
+                    "KV pool lookup response token_len=%d groups=%s hit_tokens=%d",
+                    token_len,
+                    kv_group_ids,
+                    result,
+                )
+                response = result.to_bytes(4, "big")
+                self.socket.send(response)
+
+        self.thread = threading.Thread(target=process_request, daemon=True)
+        self.thread.start()
+
+    def close(self):
+        self.socket.close(linger=0)

vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/backend/__init__.py

@@ -0,0 +1,30 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+backend_map = {
+    "mooncake": {
+        "name": "MooncakeBackend",
+        "path": "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.mooncake_backend",
+    },
+    "memcache": {
+        "name": "MemcacheBackend",
+        "path": "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.memcache_backend",
+    },
+    "yuanrong": {
+        "name": "YuanrongBackend",
+        "path": "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.yuanrong_backend",
+    },
+}

vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/backend/backend.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from vllm.config import ParallelConfig
+
+
+class Backend(ABC):
+    store: Any | None = None
+
+    @abstractmethod
+    def __init__(self, parallel_config: ParallelConfig):
+        pass
+
+    @classmethod
+    def create_scheduler_client(cls, parallel_config: ParallelConfig):
+        return cls(parallel_config)
+
+    @abstractmethod
+    def set_device(self):
+        pass
+
+    @abstractmethod
+    def register_buffer(self, ptrs: list[int], lengths: list[int]):
+        pass
+
+    @abstractmethod
+    def exists(self, keys: list[str]) -> list[int]:
+        pass
+
+    def batch_is_exist(self, keys: list[str]) -> list[int]:
+        return self.exists(keys)
+
+    def batch_get_key_info(self, keys: list[str]):
+        raise NotImplementedError(f"{type(self).__name__} does not support batch_get_key_info")
+
+    def batch_alloc(self, keys: list[str], sizes: list[int]) -> list[int]:
+        raise NotImplementedError(f"{type(self).__name__} does not support batch_alloc")
+
+    def batch_add_lease(self, keys: list[str], lease_ttl_ms: int = 0) -> list[int]:
+        raise NotImplementedError(f"{type(self).__name__} does not support batch_add_lease")
+
+    def batch_remove_lease(self, keys: list[str]) -> int:
+        raise NotImplementedError(f"{type(self).__name__} does not support batch_remove_lease")
+
+    def batch_write_finish(self, keys: list[str], results: list[int]) -> list[int]:
+        raise NotImplementedError(f"{type(self).__name__} does not support batch_write_finish")
+
+    @abstractmethod
+    def put(self, keys: list[str], addrs: list[list[int]], sizes: list[list[int]]):
+        pass
+
+    @abstractmethod
+    def get(self, keys: list[str], addrs: list[list[int]], sizes: list[list[int]]):
+        pass

vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/backend/memcache_backend.py

@@ -0,0 +1,233 @@
+# Standard
+import os
+import threading
+import time
+from enum import Enum
+from typing import Any
+
+import torch
+from vllm.config import ParallelConfig
+from vllm.distributed.parallel_state import get_world_group
+from vllm.logger import logger
+
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.backend import Backend
+
+MEMCACHE_THREAD_START_WAIT_S = 0.1
+
+
+def _is_device_sdma() -> bool:
+    config_path = os.getenv("MMC_LOCAL_CONFIG_PATH")
+    if not config_path:
+        raise ValueError("The environment variable 'MMC_LOCAL_CONFIG_PATH' is not set.")
+    with open(config_path, encoding="utf-8") as config_file:
+        for line in config_file:
+            line = line.strip()
+            if not line or line.startswith(("#", ";")):
+                continue
+            key, separator, value = line.partition("=")
+            if separator and key.strip() == "ock.mmc.local_service.protocol":
+                return value.strip() == "device_sdma"
+    return False
+
+
+class MmcDirect(Enum):
+    COPY_L2G = 0
+    COPY_G2L = 1
+    COPY_G2H = 2
+    COPY_H2G = 3
+
+
+class MemcacheBackend(Backend):
+    def __init__(
+        self,
+        parallel_config: ParallelConfig,
+        local_rank: int | None = None,
+        init_bm: bool = True,
+        lazy_init: bool = False,
+    ):
+        self.local_rank = local_rank if local_rank is not None else get_world_group().local_rank
+        self._init_bm = init_bm
+        self._lazy_init = lazy_init and _is_device_sdma()
+
+        self.store: Any | None = None
+        self._store_initialized = False
+        self._store_init_lock = threading.Lock()
+        self._pending_buffers: tuple[list[int], list[int]] | None = None
+
+        if not self._lazy_init:
+            self.store = self._setup_store()
+            self._store_initialized = True
+
+    def ensure_initialized(self):
+        if self._store_initialized:
+            return
+
+        with self._store_init_lock:
+            if self._store_initialized:
+                return
+
+            logger.info("Initializing Memcache store. local_rank=%d", self.local_rank)
+            self.store = self._setup_store()
+            self._store_initialized = True
+            self._register_buffers_if_needed()
+
+    def _setup_store(self):
+        try:
+            from memcache_hybrid import DistributedObjectStore  # type: ignore
+        except ImportError as e:
+            raise ImportError(
+                "Please install memcache by following the instructions at "
+                "https://gitee.com/ascend/memfabric_hybrid "  # noqa: E501
+                "to run vLLM with MemcacheConnector."
+            ) from e
+
+        store = DistributedObjectStore()
+
+        try:
+            res = store.init(self.local_rank, init_bm=self._init_bm)
+        except ValueError as e:
+            logger.error("Configuration loading failed. error=%s. Check memcache config and environment.", e)
+            raise
+        except Exception as exc:
+            logger.error("Store initialization failed. error=%s. Check memcache setup and dependencies.", exc)
+            raise
+
+        assert res == 0
+        time.sleep(MEMCACHE_THREAD_START_WAIT_S)
+        return store
+
+    @classmethod
+    def create_scheduler_client(cls, parallel_config: ParallelConfig):
+        # The scheduler is a single metadata client. It is initialized before
+        # the world group exists and must not initialize memcache storage, so
+        # keep the old device_id=0/init_bm=False behavior here.
+        return cls(parallel_config, local_rank=0, init_bm=False)
+
+    def init_store(self, init_bm: bool = True):
+        if self.store is not None:
+            return
+        self._init_bm = init_bm
+        self.store = self._setup_store()
+        self._store_initialized = True
+        self._register_buffers_if_needed()
+
+    def set_device(self):
+        device = torch.device(f"npu:{self.local_rank}")
+        torch.npu.set_device(device)
+
+    def register_buffer(self, ptrs: list[int], sizes: list[int]):
+        self._pending_buffers = (list(ptrs), list(sizes))
+        self._register_buffers_if_needed()
+
+    def _register_buffers_if_needed(self):
+        if self._pending_buffers is None or not self._store_initialized:
+            return
+        assert self.store is not None
+        ptrs, sizes = self._pending_buffers
+        for ptr, size in zip(ptrs, sizes):
+            self.store.register_buffer(ptr, size)
+        self._pending_buffers = None
+
+    def exists(self, keys: list[str]) -> list[int]:
+        if self._lazy_init and not self._store_initialized:
+            logger.debug(
+                "MemcacheBackend.exists called before store initialization; treating %d keys as missing.",
+                len(keys),
+            )
+            return [0] * len(keys)
+        assert self.store is not None
+        return self.store.batch_is_exist(keys)
+
+    def batch_get_key_info(self, keys: list[str]) -> list[Any]:
+        if self._lazy_init and not self._store_initialized:
+            logger.debug(
+                "MemcacheBackend.batch_get_key_info called before store initialization; "
+                "returning empty list for %d keys.",
+                len(keys),
+            )
+            return []
+        assert self.store is not None
+        return self.store.batch_get_key_info(keys)
+
+    def batch_alloc(self, keys: list[str], sizes: list[int]) -> list[int]:
+        self.ensure_initialized()
+        assert self.store is not None
+        return self.store.batch_alloc(keys, sizes)
+
+    def batch_add_lease(self, keys: list[str], lease_ttl_ms: int = 0) -> list[int]:
+        assert self.store is not None
+        return self.store.batch_add_lease(keys, lease_ttl_ms)
+
+    def batch_remove_lease(self, keys: list[str]) -> int:
+        assert self.store is not None
+        return self.store.batch_remove_lease(keys)
+
+    def batch_write_finish(self, keys: list[str], results: list[int]) -> list[int]:
+        assert self.store is not None
+        return self.store.batch_write_finish(keys, results)
+
+    def get(self, key: list[str], addr: list[list[int]], size: list[list[int]]):
+        if self._lazy_init and not self._store_initialized:
+            logger.error(
+                "Failed to get %d keys out of %d. Store is not initialized; "
+                "call put() first to trigger initialization.",
+                len(key),
+                len(key),
+            )
+            logger.debug("Failed to get key details. keys=%s", key)
+            return
+        assert self.store is not None
+        try:
+            res = self.store.batch_get_into_layers(key, addr, size, MmcDirect.COPY_G2L.value)
+            failed_codes = [int(value) for value in res if value != 0]
+            failed_count = len(failed_codes)
+            if failed_count:
+                error_codes = sorted(set(failed_codes))
+                logger.error(
+                    "Failed to get %d keys out of %d. error_codes=%s. Check key existence and memory state.",
+                    failed_count,
+                    len(key),
+                    error_codes,
+                )
+                logger.debug("Failed to get key details. keys=%s, result=%s", key, res)
+            return res
+        except Exception as e:
+            logger.error(
+                "Failed to get %d keys out of %d. type=%s, error=%s. Check store state and network.",
+                len(key),
+                len(key),
+                type(e).__name__,
+                e,
+            )
+            logger.debug("Failed to get key details. keys=%s", key)
+            return None
+
+    def put(self, key: list[str], addr: list[list[int]], size: list[list[int]]):
+        self.ensure_initialized()
+        assert self.store is not None
+        try:
+            res = self.store.batch_put_from_layers(key, addr, size, MmcDirect.COPY_L2G.value)
+            failed_codes = [int(value) for value in res if value != 0]
+            failed_count = len(failed_codes)
+            if failed_count:
+                error_codes = sorted(set(failed_codes))
+                logger.error(
+                    "Failed to put %d keys out of %d. error_codes=%s. Check memory and store capacity.",
+                    failed_count,
+                    len(key),
+                    error_codes,
+                )
+                logger.debug("Failed to put key details. keys=%s, result=%s", key, res)
+                if self._lazy_init:
+                    logger.warning("First DSV4(compress) request failure is expected. This is normal behavior.")
+        except Exception as e:
+            logger.error(
+                "Failed to put %d keys out of %d. type=%s, error=%s. Check store state and memory.",
+                len(key),
+                len(key),
+                type(e).__name__,
+                e,
+            )
+            logger.debug("Failed to put key details. keys=%s", key)
+            if self._lazy_init:
+                logger.warning("First DSV4(compress) request failure is expected. This is normal behavior.")

vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/backend/mooncake_backend.py

@@ -0,0 +1,402 @@
+# Standard
+import functools
+import json
+import os
+import threading
+from dataclasses import dataclass
+from typing import Any
+
+import regex as re
+import torch
+
+# Third Party
+from mooncake.store import ReplicateConfig  # type: ignore
+from vllm.config import ParallelConfig
+from vllm.distributed.parallel_state import get_world_group
+from vllm.logger import logger
+from vllm.utils.network_utils import get_ip
+
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.backend import Backend
+from vllm_ascend.distributed.kv_transfer.utils.mooncake_transfer_engine import global_te
+from vllm_ascend.distributed.parallel_state import get_global_rank
+
+DEFAULT_GLOBAL_SEGMENT_SIZE = 1073741824  # 1.0 GiB
+DEFAULT_LOCAL_BUFFER_SIZE = 1073741824  # 1.0 GiB
+
+
+@functools.lru_cache(maxsize=1)
+def _mooncake_setup_supports_ssd_offload() -> bool:
+    """True when installed Mooncake exposes SSD kwargs on setup() (v0.3.11+)."""
+    from mooncake.store import MooncakeDistributedStore  # type: ignore
+
+    setup = MooncakeDistributedStore.setup
+    try:
+        import inspect
+
+        sig = inspect.signature(setup)
+        return "enable_ssd_offload" in sig.parameters
+    except (TypeError, ValueError):
+        # pybind11 overloaded bindings often reject inspect.signature
+        doc = setup.__doc__ or ""
+        return "enable_ssd_offload" in doc
+
+
+def _ssd_setup_kwargs(config: "MooncakeStoreConfig") -> dict[str, object]:
+    """Keyword args for store.setup(); empty on old Mooncake or when SSD is off."""
+    if not config.enable_ssd_offload:
+        return {}
+    if not _mooncake_setup_supports_ssd_offload():
+        raise RuntimeError(
+            "mooncake.json has enable_ssd_offload=true, but the installed "
+            "Mooncake does not support enable_ssd_offload/ssd_offload_path in "
+            "MooncakeDistributedStore.setup(). Upgrade Mooncake to v0.3.11 or "
+            "later (see Mooncake ssd-offload.md Step 3A), or set "
+            "enable_ssd_offload to false."
+        )
+    return {
+        "enable_ssd_offload": config.enable_ssd_offload,
+        "ssd_offload_path": config.ssd_offload_path,
+    }
+
+
+class MooncakeBackend(Backend):
+    def __init__(self, parallel_config: ParallelConfig, lazy_init: bool = False, contribute_memory: bool = True):
+        self.parallel_config = parallel_config
+        self.config = MooncakeStoreConfig.load_from_env()
+        if self.config.protocol != "ascend":
+            raise NotImplementedError(f"MooncakeBackend does not support protocol {self.config.protocol!r}.")
+
+        self.store: Any | None = None
+        self.local_seg: str | None = None
+        self._use_fabric_mem = os.getenv("ASCEND_ENABLE_USE_FABRIC_MEM", "0") == "1"
+        self._lazy_init = lazy_init and self._use_fabric_mem
+        self._contribute_memory = contribute_memory
+        self._store_initialized = False
+        self._store_init_lock = threading.Lock()
+
+        if not self._lazy_init:
+            self.store = self._setup_store()
+            self._store_initialized = True
+
+    def ensure_initialized(self):
+        if self._store_initialized:
+            return
+
+        with self._store_init_lock:
+            if self._store_initialized:
+                return
+
+            logger.info("Initializing Mooncake store. metadata_server=%s", self.config.metadata_server)
+            self.store = self._setup_store()
+            self._store_initialized = True
+
+    def _setup_store(self):
+        try:
+            from mooncake.store import MooncakeDistributedStore  # type: ignore
+        except ImportError as e:
+            raise ImportError(
+                "Please install mooncake by following the instructions at "
+                "https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md "  # noqa: E501
+                "to run vLLM with MooncakeConnector."
+            ) from e
+
+        store = MooncakeDistributedStore()
+        local_hostname = get_ip()
+        ssd_kwargs = _ssd_setup_kwargs(self.config)
+        # Scheduler-only clients (contribute_memory=False) do not contribute
+        # KV cache memory and therefore do not need SSD offload. Passing
+        # enable_ssd_offload=True for them would cause Mooncake to register
+        # an extra active client on the master, inflating both the client
+        # count and the reported SSD storage usage.
+        if ssd_kwargs and not self._contribute_memory:
+            ssd_kwargs = {}
+        # Each rank that contributes memory to the pool uses its own SSD
+        # directory to avoid bucket file collisions. Key by the globally unique
+        # rank so that DP/TP/PP/CP replicas never share a directory (dense and
+        # MoE alike); only ranks that contribute memory need an offload dir.
+        if ssd_kwargs and ssd_kwargs.get("ssd_offload_path"):
+            global_rank = get_global_rank(self.parallel_config)
+            rank_path = os.path.join(str(ssd_kwargs["ssd_offload_path"]), f"rank_{global_rank}")
+            try:
+                os.makedirs(rank_path, exist_ok=True)
+            except OSError as e:
+                raise RuntimeError(f"Failed to create per-rank SSD offload directory: {rank_path!r} ({e})")
+            ssd_kwargs["ssd_offload_path"] = rank_path
+        # ASCEND_ENABLE_USE_FABRIC_MEM: Enable unified memory address direct transmission scheme
+        # and only can be used for 800 I/T A3 series.
+        # Required supporting hardware versions are as follows:
+        if not self._use_fabric_mem:
+            transfer_engine = global_te.get_transfer_engine(local_hostname, device_name=None)
+            self.local_seg = local_hostname + ":" + str(transfer_engine.get_rpc_port())
+            ret = store.setup(
+                local_hostname=self.local_seg,
+                metadata_server=self.config.metadata_server,
+                global_segment_size=self.config.global_segment_size if self._contribute_memory else 0,
+                local_buffer_size=self.config.local_buffer_size if self._contribute_memory else 0,
+                protocol=self.config.protocol,
+                rdma_devices=self.config.device_name,
+                master_server_addr=self.config.master_server_address,
+                engine=transfer_engine.get_engine(),
+                **ssd_kwargs,
+            )
+        else:
+            self.local_seg = local_hostname
+            ret = store.setup(
+                local_hostname=self.local_seg,
+                metadata_server=self.config.metadata_server,
+                global_segment_size=self.config.global_segment_size if self._contribute_memory else 0,
+                local_buffer_size=0,
+                protocol=self.config.protocol,
+                rdma_devices=self.config.device_name,
+                master_server_addr=self.config.master_server_address,
+                **ssd_kwargs,
+            )
+
+        if ret != 0:
+            msg = "Initialize mooncake failed."
+            logger.error(
+                "Initialize mooncake failed. ret=%d, metadata_server=%s. Check mooncake config and network.",
+                ret,
+                self.config.metadata_server,
+            )
+            raise RuntimeError(msg)
+        if ssd_kwargs:
+            logger.info(
+                "Mooncake SSD offload enabled (Mode A): path=%s",
+                self.config.ssd_offload_path,
+            )
+        return store
+
+    @classmethod
+    def create_scheduler_client(cls, parallel_config: ParallelConfig):
+        torch.npu.set_device(0)
+        return cls(parallel_config, contribute_memory=False)
+
+    def set_device(self):
+        local_rank = get_world_group().local_rank
+        device = torch.device(f"npu:{local_rank}")
+        torch.npu.set_device(device)
+
+    def register_buffer(self, ptrs: list[int], lengths: list[int]):
+        if not self._use_fabric_mem:
+            local_hostname = get_ip()
+            global_te.get_transfer_engine(local_hostname, device_name=None)
+            global_te.register_buffer(ptrs, lengths)
+
+    def exists(self, keys: list[str]) -> list[int]:
+        if self._lazy_init and not self._store_initialized:
+            logger.debug(
+                "MooncakeBackend.exists called before store initialization; treating %d keys as missing.",
+                len(keys),
+            )
+            return [0] * len(keys)
+        assert self.store is not None
+        return self.store.batch_is_exist(keys)
+
+    def put(self, keys: list[str], addrs: list[list[int]], sizes: list[list[int]]):
+        self.ensure_initialized()
+        assert self.store is not None
+        try:
+            config = ReplicateConfig()
+            if self.config.preferred_segment:
+                config.preferred_segment = self.local_seg
+            config.prefer_alloc_in_same_node = self.config.prefer_alloc_in_same_node
+            res = self.store.batch_put_from_multi_buffers(keys, addrs, sizes, config)
+            failed_codes = [int(value) for value in res if value < 0]
+            failed_count = len(failed_codes)
+            if failed_count:
+                error_codes = sorted(set(failed_codes))
+                logger.error(
+                    "Failed to put %d keys out of %d. error_codes=%s. Check memory and store capacity.",
+                    failed_count,
+                    len(keys),
+                    error_codes,
+                )
+                logger.debug("Failed to put key details. keys=%s, result=%s", keys, res)
+                if self._lazy_init:
+                    logger.warning("First DSV4(compress) request failure is expected. This is normal behavior.")
+        except Exception as e:
+            logger.error(
+                "Failed to put %d keys out of %d. type=%s, error=%s. Check store state and memory.",
+                len(keys),
+                len(keys),
+                type(e).__name__,
+                e,
+            )
+            logger.debug("Failed to put key details. keys=%s", keys)
+            if self._lazy_init:
+                logger.warning("First DSV4(compress) request failure is expected. This is normal behavior.")
+
+    def get(self, keys: list[str], addrs: list[list[int]], sizes: list[list[int]]):
+        if self._lazy_init and not self._store_initialized:
+            logger.error(
+                "Failed to get %d keys out of %d. Store is not initialized; "
+                "call put() first to trigger initialization.",
+                len(keys),
+                len(keys),
+            )
+            logger.debug("Failed to get key details. keys=%s", keys)
+            return
+        assert self.store is not None
+        logger.debug(
+            "MooncakeBackend.get enter keys=%d sample_keys=%s",
+            len(keys),
+            keys[:3],
+        )
+        try:
+            res = self.store.batch_get_into_multi_buffers(keys, addrs, sizes)
+            res_list = list(res)
+            failed_codes = [int(value) for value in res_list if value < 0]
+            failed_count = len(failed_codes)
+            error_codes = sorted(set(failed_codes))
+            if failed_count:
+                logger.error(
+                    "Failed to get %d keys out of %d. error_codes=%s. Check key existence and memory state.",
+                    failed_count,
+                    len(keys),
+                    error_codes,
+                )
+                logger.debug("Failed to get key details. keys=%s, result=%s", keys, res_list)
+            for i, value in enumerate(res_list):
+                if value > 0:
+                    res_list[i] = 0
+            return res_list
+        except Exception as e:
+            logger.error(
+                "Failed to get %d keys out of %d. type=%s, error=%s. Check store state and network.",
+                len(keys),
+                len(keys),
+                type(e).__name__,
+                e,
+            )
+            logger.debug("Failed to get key details. keys=%s", keys)
+            return None
+
+
+@dataclass
+class MooncakeStoreConfig:
+    metadata_server: str
+    global_segment_size: int | str
+    local_buffer_size: int
+    protocol: str
+    device_name: str
+    master_server_address: str
+    preferred_segment: bool
+    prefer_alloc_in_same_node: bool
+    enable_ssd_offload: bool = False
+    ssd_offload_path: str = ""
+
+    def __post_init__(self) -> None:
+        if not self.enable_ssd_offload:
+            return
+        if not self.ssd_offload_path:
+            raise ValueError(
+                "enable_ssd_offload is true but ssd_offload_path is empty. Set ssd_offload_path in mooncake.json."
+            )
+        if not os.path.isabs(self.ssd_offload_path):
+            raise ValueError(f"ssd_offload_path must be an absolute path, got: {self.ssd_offload_path!r}")
+
+    @staticmethod
+    def from_file(file_path: str) -> "MooncakeStoreConfig":
+        with open(file_path) as file:
+            config = json.load(file)
+        master_server_address = os.getenv("MOONCAKE_MASTER", None)
+        global_segment_size_env = os.getenv("MOONCAKE_GLOBAL_SEGMENT_SIZE", None)
+        return MooncakeStoreConfig(
+            metadata_server=config.get("metadata_server"),
+            global_segment_size=_parse_global_segment_size(
+                global_segment_size_env
+                if global_segment_size_env is not None
+                else config.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE)
+            ),
+            local_buffer_size=_parse_global_segment_size(config.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE)),
+            protocol=config.get("protocol", "ascend"),
+            device_name=config.get("device_name", ""),
+            master_server_address=master_server_address
+            if master_server_address is not None
+            else config.get("master_server_address"),
+            preferred_segment=config.get("preferred_segment", False),
+            prefer_alloc_in_same_node=config.get("prefer_alloc_in_same_node", True),
+            enable_ssd_offload=bool(config.get("enable_ssd_offload", False)),
+            ssd_offload_path=config.get("ssd_offload_path", ""),
+        )
+
+    @staticmethod
+    def load_from_env() -> "MooncakeStoreConfig":
+        config_path = os.getenv("MOONCAKE_CONFIG_PATH")
+        if not config_path:
+            raise ValueError("The environment variable 'MOONCAKE_CONFIG_PATH' is not set.")
+        return MooncakeStoreConfig.from_file(config_path)
+
+
+def _parse_global_segment_size(value) -> int:
+    """
+    Parse storage size strings with support for units: GB, MB, KB, B
+
+    Args:
+        value: Input value (int, str, or other convertible types)
+
+    Returns:
+        int: Size in bytes
+
+    Raises:
+        ValueError: For invalid format, missing number, or negative values
+        TypeError: For unsupported input types
+    """
+
+    if isinstance(value, int):
+        return value
+    elif not isinstance(value, str):
+        try:
+            return int(value)
+        except (TypeError, ValueError) as e:
+            raise TypeError(f"Unsupported type for global_segment_size: {type(value)}") from e
+
+    cleaned_input = value.strip().lower()
+    if not cleaned_input:
+        raise ValueError("global segment size cannot be empty.")
+
+    UNIT_MULTIPLIERS = {
+        "gb": 1024**3,  # 1 GB = 1024^3 bytes
+        "mb": 1024**2,  # 1 MB = 1024^2 bytes
+        "kb": 1024,  # 1 KB = 1024 bytes
+        "b": 1,  # 1 B = 1 byte
+    }
+    pattern = r"^\s*([\d.]+)\s*(gb|mb|kb|b)?\s*$"
+    match = re.match(pattern, cleaned_input)
+
+    if not match:
+        raise ValueError(f"Invalid format: '{value}'")
+
+    number_str = match.group(1)
+    unit = match.group(2) or "b"
+
+    multiplier = UNIT_MULTIPLIERS[unit]
+    return _convert_to_bytes(number_str, multiplier, value)
+
+
+def _convert_to_bytes(number_str: str, multiplier: int, original_input: str) -> int:
+    """
+    Convert numeric string to byte count
+
+    Args:
+        number_str: Numeric portion of input
+        multiplier: Unit conversion factor
+        original_input: Original input string (for error messages)
+
+    Returns:
+        int: Byte count
+
+    Raises:
+        ValueError: For invalid numbers or negative results
+    """
+    try:
+        numeric_value = float(number_str)
+    except ValueError:
+        raise ValueError(f"Invalid numeric value '{number_str}' in: '{original_input}'")
+    # Calculate byte count
+    try:
+        byte_count = int(numeric_value * multiplier)
+    except OverflowError:
+        raise ValueError(f"Storage size too large: '{original_input}'")
+    return byte_count

vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/backend/yuanrong_backend.py

@@ -0,0 +1,238 @@
+import hashlib
+import os
+from dataclasses import dataclass
+from typing import Any
+
+import regex as re
+import torch
+from vllm.config import ParallelConfig
+from vllm.distributed.parallel_state import get_world_group
+from vllm.logger import logger
+from vllm.utils.network_utils import split_host_port
+
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.backend import Backend
+from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type
+
+
+def _iter_slices(total: int, batch_size: int):
+    for start in range(0, total, batch_size):
+        end = min(start + batch_size, total)
+        yield start, end
+
+
+@dataclass
+class YuanrongConfig:
+    worker_addr: str
+    enable_exclusive_connection: bool
+    enable_remote_h2d: bool
+
+    @staticmethod
+    def load_from_env() -> "YuanrongConfig":
+        worker_addr = os.getenv("DS_WORKER_ADDR")
+        if not worker_addr:
+            raise ValueError("Environment variable DS_WORKER_ADDR is required, expected format '<host>:<port>'.")
+
+        return YuanrongConfig(
+            worker_addr=worker_addr,
+            enable_exclusive_connection=bool(int(os.getenv("DS_ENABLE_EXCLUSIVE_CONNECTION", "0"))),
+            enable_remote_h2d=bool(int(os.getenv("DS_ENABLE_REMOTE_H2D", "0"))),
+        )
+
+
+class YuanrongHelper:
+    _DS_KEY_MAX_LEN = 1024
+    _DS_KEY_ALLOWED_PATTERN = re.compile(r"^[a-zA-Z0-9\-_!@#%\^\*\(\)\+\=\:;]+$")
+    _DS_KEY_INVALID_CHAR_PATTERN = re.compile(r"[^a-zA-Z0-9\-_!@#%\^\*\(\)\+\=\:;]")
+    _DS_KEY_HASH_SUFFIX_LEN = 16
+
+    def __init__(self, blob_cls, blob_list_cls):
+        self._blob_cls = blob_cls
+        self._blob_list_cls = blob_list_cls
+        self._device_id: int | None = None
+
+    def normalize_keys(self, keys: list[str]) -> list[str]:
+        normalized: list[str] = []
+        for key in keys:
+            if len(key) <= self._DS_KEY_MAX_LEN and self._DS_KEY_ALLOWED_PATTERN.match(key):
+                normalized.append(key)
+                continue
+
+            sanitized = self._DS_KEY_INVALID_CHAR_PATTERN.sub("_", key)
+            hash_digest = hashlib.sha256(key.encode("utf-8")).hexdigest()
+            suffix = f"__{hash_digest[: self._DS_KEY_HASH_SUFFIX_LEN]}"
+            max_prefix_len = self._DS_KEY_MAX_LEN - len(suffix)
+            normalized.append(sanitized[:max_prefix_len] + suffix)
+        return normalized
+
+    def make_blob_lists(self, addrs_list: list[list[int]], sizes_list: list[list[int]]) -> list[Any]:
+        total = len(addrs_list)
+        if total != len(sizes_list):
+            raise ValueError("Address list and size list length mismatch.")
+
+        device_id = self._device_id
+        if device_id is None:
+            logger.error("Device id is not set. Check device initialization and configuration.")
+            raise RuntimeError("Yuanrong backend device id is not initialized.")
+
+        blob_lists: list[Any] = []
+        for addrs, sizes in zip(addrs_list, sizes_list):
+            if len(addrs) != len(sizes):
+                raise ValueError("Address list and size list length mismatch.")
+            blobs = [
+                self._blob_cls(addr, size)  # type: ignore[misc]
+                for addr, size in zip(addrs, sizes)
+            ]
+            blob_lists.append(
+                self._blob_list_cls(device_id, blobs)  # type: ignore[misc]
+            )
+        return blob_lists
+
+
+class YuanrongBackend(Backend):
+    _DS_MAX_BATCH_KEYS = 10000
+
+    def __init__(self, parallel_config: ParallelConfig):
+        try:
+            from yr.datasystem.hetero_client import Blob, DeviceBlobList, HeteroClient  # type: ignore[import-not-found]
+            from yr.datasystem.kv_client import SetParam  # type: ignore[import-not-found]
+            from yr.datasystem.object_client import WriteMode  # type: ignore[import-not-found]
+        except ImportError as exc:
+            raise ImportError("Please install openyuanrong-datasystem to use the yuanrong backend.") from exc
+
+        self._helper = YuanrongHelper(Blob, DeviceBlobList)
+        self._ds_set_param = SetParam()
+        self._ds_set_param.write_mode = WriteMode.NONE_L2_CACHE_EVICT
+
+        self.config = YuanrongConfig.load_from_env()
+        try:
+            host, port = split_host_port(self.config.worker_addr)
+        except Exception as exc:
+            raise ValueError(f"Invalid DS_WORKER_ADDR '{self.config.worker_addr}', expected '<host>:<port>'.") from exc
+        self._hetero_client = HeteroClient(
+            host,
+            int(port),
+            enable_exclusive_connection=self.config.enable_exclusive_connection,
+            enable_remote_h2d=self.config.enable_remote_h2d,
+        )
+        self._hetero_client.init()
+        self._is_a2 = get_ascend_device_type() in {AscendDeviceType.A2}
+        self._registered_buffers: tuple[list[int], list[int]] | None = None
+        self._buffers_registered = False
+
+    def _ensure_device_ready(self):
+        if self._helper._device_id is None:
+            self.set_device()
+
+    def set_device(self):
+        local_rank = get_world_group().local_rank
+        device = torch.device(f"npu:{local_rank}")
+        torch.npu.set_device(device)
+        self._helper._device_id = int(torch.npu.current_device())
+
+    def register_buffer(self, ptrs: list[int], lengths: list[int]):
+        self._registered_buffers = (list(ptrs), list(lengths))
+        self._register_buffers_if_needed()
+
+    def _register_buffers_if_needed(self):
+        if self._is_a2:
+            return
+        if not self.config.enable_remote_h2d:
+            return
+        if self._registered_buffers is None or self._buffers_registered:
+            return
+        ptrs, lengths = self._registered_buffers
+        self._hetero_client.pre_register_device_memory(ptrs, lengths)  # type: ignore[union-attr]
+        self._buffers_registered = True
+
+    def exists(self, keys: list[str]) -> list[int]:
+        if len(keys) == 0:
+            return []
+        try:
+            keys = self._helper.normalize_keys(keys)
+            if len(keys) <= self._DS_MAX_BATCH_KEYS:
+                exists = self._hetero_client.exist(keys)  # type: ignore[union-attr]
+                return [1 if value else 0 for value in exists]
+            results: list[int] = []
+            for start, end in _iter_slices(len(keys), self._DS_MAX_BATCH_KEYS):
+                exists = self._hetero_client.exist(keys[start:end])  # type: ignore[union-attr]
+                results.extend(1 if value else 0 for value in exists)
+            return results
+        except Exception as exc:
+            logger.error(
+                "Failed to check keys. keys_count=%d, type=%s, error=%s. Check network and yuanrong service.",
+                len(keys),
+                type(exc).__name__,
+                exc,
+            )
+            return [0] * len(keys)
+
+    def get(self, keys: list[str], addrs: list[list[int]], sizes: list[list[int]]) -> list[int] | None:
+        if len(keys) == 0:
+            return []
+        failed_keys_for_log = keys
+        try:
+            self._ensure_device_ready()
+            keys = self._helper.normalize_keys(keys)
+            failed_keys_for_log = keys
+            blob_lists = self._helper.make_blob_lists(addrs, sizes)
+            failed_keys: list[str] = []
+            if len(keys) <= self._DS_MAX_BATCH_KEYS:
+                failed_keys = self._hetero_client.mget_h2d(  # type: ignore[union-attr]
+                    keys, blob_lists, 0
+                )
+            else:
+                for start, end in _iter_slices(len(keys), self._DS_MAX_BATCH_KEYS):
+                    failed_keys_for_log = keys[start:end]
+                    failed_keys.extend(
+                        self._hetero_client.mget_h2d(  # type: ignore[union-attr]
+                            keys[start:end], blob_lists[start:end], 0
+                        )
+                    )
+            if failed_keys:
+                logger.error(
+                    "Failed to get %d keys out of %d. Check key existence and memory state.",
+                    len(failed_keys),
+                    len(keys),
+                )
+                logger.debug("Failed to get key details. failed_keys=%s", failed_keys)
+            failed_set = set(failed_keys)
+            return [1 if k in failed_set else 0 for k in keys]
+        except Exception as exc:
+            logger.error(
+                "Failed to get %d keys out of %d. type=%s, error=%s. Check network and yuanrong service.",
+                len(failed_keys_for_log),
+                len(keys),
+                type(exc).__name__,
+                exc,
+            )
+            logger.debug("Failed to get key details. keys=%s", failed_keys_for_log)
+            return None
+
+    def put(self, keys: list[str], addrs: list[list[int]], sizes: list[list[int]]):
+        if len(keys) == 0:
+            return
+        failed_keys_for_log = keys
+        try:
+            self._ensure_device_ready()
+            keys = self._helper.normalize_keys(keys)
+            failed_keys_for_log = keys
+            blob_lists = self._helper.make_blob_lists(addrs, sizes)
+            if len(keys) <= self._DS_MAX_BATCH_KEYS:
+                self._hetero_client.mset_d2h(  # type: ignore[union-attr]
+                    keys, blob_lists, self._ds_set_param
+                )
+            else:
+                for start, end in _iter_slices(len(keys), self._DS_MAX_BATCH_KEYS):
+                    failed_keys_for_log = keys[start:end]
+                    self._hetero_client.mset_d2h(  # type: ignore[union-attr]
+                        keys[start:end], blob_lists[start:end], self._ds_set_param
+                    )
+        except Exception as exc:
+            logger.error(
+                "Failed to put %d keys out of %d. type=%s, error=%s. Check network and yuanrong service.",
+                len(failed_keys_for_log),
+                len(keys),
+                type(exc).__name__,
+                exc,
+            )
+            logger.debug("Failed to put key details. keys=%s", failed_keys_for_log)

vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/coordinator.py

@@ -0,0 +1,450 @@
+from __future__ import annotations
+
+from dataclasses import replace
+from importlib import import_module
+from typing import Any, cast
+
+from vllm.logger import logger
+from vllm.v1.core.block_pool import BlockPool
+from vllm.v1.core.kv_cache_utils import BlockHash, BlockHashList, KVCacheBlock
+from vllm.v1.core.single_type_kv_cache_manager import SingleTypeKVCacheManager
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheGroupSpec,
+    KVCacheSpec,
+    UniformTypeKVCacheSpecs,
+)
+
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.config_data import (
+    GroupedBlockHashCache,
+    block_hash_to_bytes,
+    get_block_hashes,
+)
+
+_CACHE_MISSING = object()
+_MANAGER_CLASS_CACHE_ATTR = "_manager_class_cache"
+
+
+class ExternalCachedBlockPool:
+    """Duck-typed BlockPool backed by external AscendStore key existence."""
+
+    def __init__(self, exists: set[tuple[int, bytes]] | None = None) -> None:
+        # exists=None is used for load/store masks where hit length has already
+        # been decided and each manager only needs to apply its own reachability.
+        self._exists = exists
+        self.null_block = KVCacheBlock(block_id=0)
+        self._present_block = KVCacheBlock(block_id=1)
+
+    def get_cached_block(
+        self,
+        block_hash: BlockHash,
+        group_ids: list[int],
+    ) -> list[KVCacheBlock] | None:
+        if self._exists is None:
+            return [self._present_block] * len(group_ids)
+        h = block_hash_to_bytes(block_hash)
+        if all((group_id, h) in self._exists for group_id in group_ids):
+            return [self._present_block] * len(group_ids)
+        return None
+
+
+class AscendStoreCoordinator:
+    """Hybrid cache-hit/mask coordinator for AscendStore external KV Pool.
+
+    This mirrors vLLM MooncakeStoreCoordinator but uses AscendStore's external
+    key granularity. For DSV4 compressed groups, keys are generated over the
+    raw-token span ``group_block_size * compress_ratio`` while transfer
+    addresses remain in cache-domain blocks.
+    """
+
+    def __init__(
+        self,
+        kv_cache_groups: list[KVCacheGroupSpec],
+        scheduler_block_size: int,
+        hash_block_size: int,
+        group_block_sizes: list[int],
+        group_cache_families: list[str],
+        use_eagle: bool = False,
+        retention_interval: int | None = None,
+    ) -> None:
+        assert len(kv_cache_groups) == len(group_block_sizes)
+        assert len(kv_cache_groups) == len(group_cache_families)
+        assert scheduler_block_size % hash_block_size == 0, (
+            f"scheduler_block_size ({scheduler_block_size}) must be a multiple of hash_block_size ({hash_block_size})"
+        )
+
+        self.kv_cache_groups = kv_cache_groups
+        self.hash_block_size = hash_block_size
+        self.lcm_block_size = scheduler_block_size
+        self.use_eagle = use_eagle
+        self.retention_interval = retention_interval
+        self.group_block_sizes = group_block_sizes
+        self.group_cache_families = group_cache_families
+        self.group_effective_block_sizes = [
+            _cache_family_granularity(block_size, family)
+            for block_size, family in zip(group_block_sizes, group_cache_families, strict=True)
+        ]
+        for effective_block_size in self.group_effective_block_sizes:
+            assert effective_block_size % hash_block_size == 0, "block_size must be divisible by hash_block_size"
+            assert scheduler_block_size % effective_block_size == 0, (
+                "scheduler_block_size must be a multiple of each group's effective block_size"
+            )
+
+        self.eagle_group_ids = {i for i, group in enumerate(kv_cache_groups) if group.is_eagle_group}
+        if use_eagle and not self.eagle_group_ids:
+            self.eagle_group_ids = set(range(len(kv_cache_groups)))
+
+        self._verify_and_split_kv_cache_groups()
+
+    def _verify_and_split_kv_cache_groups(self) -> None:
+        attention_groups: list[tuple[KVCacheSpec, list[int], type[SingleTypeKVCacheManager]]] = []
+        self.group_effective_specs: list[KVCacheSpec] = []
+
+        for group_id, group in enumerate(self.kv_cache_groups):
+            spec = _unwrap_spec(group.kv_cache_spec)
+            effective_spec = _copy_spec_with_block_size(spec, self.group_effective_block_sizes[group_id])
+            if (
+                not _uses_reachable_mask(self.group_cache_families[group_id])
+                and getattr(effective_spec, "compress_ratio", 1) > 1
+            ):
+                # The cache family already folds the compression ratio into
+                # the external key granularity. Avoid applying it again inside
+                # CompressAttentionManager.find_longest_cache_hit().
+                effective_spec = replace(effective_spec, compress_ratio=1)
+            self.group_effective_specs.append(effective_spec)
+            manager_cls = _get_manager_class(spec)
+
+            for existing_spec, group_ids, existing_cls in attention_groups:
+                if existing_spec == effective_spec:
+                    assert manager_cls is existing_cls, "Expected same manager class for identical KV cache specs."
+                    group_ids.append(group_id)
+                    break
+            else:
+                attention_groups.append((effective_spec, [group_id], manager_cls))
+
+        self.attention_groups = sorted(
+            attention_groups,
+            key=lambda item: not isinstance(item[0], FullAttentionSpec),
+        )
+        self.eagle_attn_group_indices: set[int] = {
+            index
+            for index, (_, group_ids, _) in enumerate(self.attention_groups)
+            if any(group_id in self.eagle_group_ids for group_id in group_ids)
+        }
+        if self.use_eagle and not self.eagle_attn_group_indices:
+            self.eagle_attn_group_indices = set(range(len(self.attention_groups)))
+        self.eagle_reachable_group_ids: set[int] = {
+            group_id for index in self.eagle_attn_group_indices for group_id in self.attention_groups[index][1]
+        }
+
+    def find_longest_cache_hit(
+        self,
+        block_hashes: list[BlockHash],
+        max_length: int,
+        cached_block_pool: ExternalCachedBlockPool,
+        *,
+        apply_eagle: bool = True,
+        grouped_hash_cache: GroupedBlockHashCache | None = None,
+    ) -> tuple[tuple[list[bool], ...], int]:
+        blocks_per_group, hit_length = self._find_hit_blocks(
+            block_hashes,
+            max_length,
+            cached_block_pool,
+            apply_eagle=apply_eagle,
+            grouped_hash_cache=grouped_hash_cache,
+        )
+        masks = tuple([block is not cached_block_pool.null_block for block in blocks] for blocks in blocks_per_group)
+        return masks, hit_length
+
+    def load_mask(
+        self,
+        block_hashes: list[BlockHash],
+        token_len: int,
+        grouped_hash_cache: GroupedBlockHashCache | None = None,
+    ) -> tuple[list[bool], ...]:
+        masks, _ = self.find_longest_cache_hit(
+            block_hashes,
+            token_len,
+            ExternalCachedBlockPool(),
+            apply_eagle=False,
+            grouped_hash_cache=grouped_hash_cache,
+        )
+        return tuple(
+            [True] * _num_chunks(token_len, self.group_effective_block_sizes[group_id])
+            if not _uses_reachable_mask(self.group_cache_families[group_id])
+            else mask
+            for group_id, mask in enumerate(masks)
+        )
+
+    def _reachable_masks(
+        self,
+        aligned_token_len: int,
+        retention_interval: int | None,
+        num_prompt_tokens: int | None,
+    ) -> list[tuple[int, list[bool] | None]]:
+        assert aligned_token_len % self.lcm_block_size == 0, (
+            f"aligned_token_len ({aligned_token_len}) must be a multiple of lcm_block_size ({self.lcm_block_size})"
+        )
+        masks: list[tuple[int, list[bool] | None]] = []
+        for group_id, spec in enumerate(self.group_effective_specs):
+            num_chunks = aligned_token_len // self.group_effective_block_sizes[group_id]
+            if not _uses_reachable_mask(self.group_cache_families[group_id]):
+                masks.append((num_chunks, None))
+                continue
+            manager_cls = _get_manager_class(_unwrap_spec(self.kv_cache_groups[group_id].kv_cache_spec))
+            mask = _reachable_block_mask(
+                manager_cls,
+                start_block=0,
+                end_block=num_chunks,
+                alignment_tokens=self.lcm_block_size,
+                kv_cache_spec=spec,
+                use_eagle=group_id in self.eagle_reachable_group_ids,
+                retention_interval=retention_interval,
+                num_prompt_tokens=num_prompt_tokens,
+            )
+            masks.append((num_chunks, mask))
+        return masks
+
+    def store_mask(
+        self,
+        aligned_token_len: int,
+        num_prompt_tokens: int | None = None,
+    ) -> tuple[list[bool], ...]:
+        masks = self._reachable_masks(aligned_token_len, self.retention_interval, num_prompt_tokens)
+        return tuple([True] * num_chunks if mask is None else mask for num_chunks, mask in masks)
+
+    def lookup_mask(
+        self,
+        aligned_token_len: int,
+    ) -> tuple[list[bool] | None, ...]:
+        masks = self._reachable_masks(aligned_token_len, None, None)
+        for num_chunks, mask in masks:
+            if mask is not None:
+                assert len(mask) == num_chunks
+        return tuple(None if mask is None or all(mask) else mask for _, mask in masks)
+
+    def block_hashes_for_spec(
+        self,
+        block_hashes: list[BlockHash],
+        spec: KVCacheSpec,
+        grouped_hash_cache: GroupedBlockHashCache | None = None,
+    ) -> BlockHashList:
+        if spec.block_size == self.hash_block_size:
+            return block_hashes
+        return cast(
+            BlockHashList,
+            get_block_hashes(
+                block_hashes,
+                spec.block_size,
+                self.hash_block_size,
+                grouped_hash_cache=grouped_hash_cache,
+            ),
+        )
+
+    def _find_hit_blocks(
+        self,
+        block_hashes: list[BlockHash],
+        max_length: int,
+        cached_block_pool: ExternalCachedBlockPool,
+        *,
+        apply_eagle: bool = True,
+        grouped_hash_cache: GroupedBlockHashCache | None = None,
+    ) -> tuple[tuple[list[KVCacheBlock], ...], int]:
+        eagle_indices = self.eagle_attn_group_indices if apply_eagle else set()
+        if len(self.attention_groups) == 1:
+            spec, group_ids, manager_cls = self.attention_groups[0]
+            hashes = self.block_hashes_for_spec(block_hashes, spec, grouped_hash_cache)
+            hit_blocks = _find_longest_cache_hit(
+                manager_cls,
+                block_hashes=hashes,
+                max_length=max_length,
+                kv_cache_group_ids=group_ids,
+                block_pool=cast(BlockPool, cached_block_pool),
+                kv_cache_spec=spec,
+                drop_eagle_block=0 in eagle_indices,
+                alignment_tokens=spec.block_size,
+            )
+            blocks_by_group: list[list[KVCacheBlock]] = [[] for _ in range(len(self.kv_cache_groups))]
+            for group_id, blocks in zip(group_ids, hit_blocks, strict=True):
+                blocks_by_group[group_id] = blocks
+            return tuple(blocks_by_group), len(hit_blocks[0]) * spec.block_size
+
+        hit_length = max_length
+        hit_blocks_by_group: list[list[KVCacheBlock] | None] = [None] * len(self.kv_cache_groups)
+        is_simple_hybrid = len(self.attention_groups) == 2 and isinstance(
+            self.attention_groups[0][0], FullAttentionSpec
+        )
+        eagle_verified: set[int] = set()
+
+        while True:
+            curr_hit_length = hit_length
+
+            for index, (spec, group_ids, manager_cls) in enumerate(self.attention_groups):
+                cached = hit_blocks_by_group[group_ids[0]]
+                if isinstance(spec, FullAttentionSpec) and cached is not None:
+                    curr_hit_length = curr_hit_length // spec.block_size * spec.block_size
+                    continue
+
+                drop_eagle_block = index in eagle_indices and index not in eagle_verified
+                max_group_length = curr_hit_length
+                if drop_eagle_block:
+                    max_group_length = min(curr_hit_length + spec.block_size, max_length)
+                hashes = self.block_hashes_for_spec(block_hashes, spec, grouped_hash_cache)
+                hit_blocks = _find_longest_cache_hit(
+                    manager_cls,
+                    block_hashes=hashes,
+                    max_length=max_group_length,
+                    kv_cache_group_ids=group_ids,
+                    block_pool=cast(BlockPool, cached_block_pool),
+                    kv_cache_spec=spec,
+                    drop_eagle_block=drop_eagle_block,
+                    alignment_tokens=self.lcm_block_size,
+                )
+                new_hit_length = len(hit_blocks[0]) * spec.block_size
+                if drop_eagle_block:
+                    eagle_verified.add(index)
+                elif new_hit_length < curr_hit_length:
+                    eagle_verified.clear()
+                curr_hit_length = new_hit_length
+                for group_id, blocks in zip(group_ids, hit_blocks, strict=True):
+                    hit_blocks_by_group[group_id] = blocks
+
+            if curr_hit_length >= hit_length:
+                break
+            hit_length = curr_hit_length
+            if is_simple_hybrid:
+                break
+
+        spec0, group_ids0, _ = self.attention_groups[0]
+        if isinstance(spec0, FullAttentionSpec):
+            num_blocks = hit_length // spec0.block_size
+            for group_id in group_ids0:
+                full_blocks = hit_blocks_by_group[group_id]
+                assert full_blocks is not None
+                del full_blocks[num_blocks:]
+
+        return (
+            tuple(blocks if blocks is not None else [] for blocks in hit_blocks_by_group),
+            hit_length,
+        )
+
+
+def _unwrap_spec(spec: KVCacheSpec) -> KVCacheSpec:
+    if isinstance(spec, UniformTypeKVCacheSpecs):
+        return next(iter(spec.kv_cache_specs.values()))
+    return spec
+
+
+def _copy_spec_with_block_size(spec: KVCacheSpec, block_size: int) -> KVCacheSpec:
+    if spec.block_size == block_size:
+        return spec
+    copy_with_new_block_size = getattr(spec, "copy_with_new_block_size", None)
+    if copy_with_new_block_size is not None:
+        return copy_with_new_block_size(block_size)
+    return replace(spec, block_size=block_size)
+
+
+def _get_manager_class_cache() -> dict[str, Any]:
+    cache = getattr(_get_manager_class, _MANAGER_CLASS_CACHE_ATTR, None)
+    if not isinstance(cache, dict):
+        cache = {}
+        setattr(_get_manager_class, _MANAGER_CLASS_CACHE_ATTR, cache)
+    return cast(dict[str, Any], cache)
+
+
+def _get_manager_class(spec: KVCacheSpec) -> type[SingleTypeKVCacheManager]:
+    cache = _get_manager_class_cache()
+    compress_ratio = getattr(spec, "compress_ratio", None)
+    if compress_ratio is not None and compress_ratio > 1:
+        compress_manager = cache.get("compress_manager", _CACHE_MISSING)
+        if compress_manager is _CACHE_MISSING:
+            try:
+                from vllm_ascend.core.single_type_kv_cache_manager import CompressAttentionManager
+            except ImportError:
+                compress_manager = None
+            else:
+                compress_manager = CompressAttentionManager
+            cache["compress_manager"] = compress_manager
+        if compress_manager is not None:
+            return cast(type[SingleTypeKVCacheManager], compress_manager)
+
+    registry = cache.get("registry", _CACHE_MISSING)
+    if registry is _CACHE_MISSING:
+        try:
+            registry_module = import_module("vllm.v1.kv_cache_spec_registry")
+            registry = getattr(registry_module, "KVCacheSpecRegistry", None)
+        except ImportError:
+            registry = None
+        cache["registry"] = registry
+
+    if registry is not None:
+        manager_cls = registry.get_manager_class(spec)
+        if manager_cls is not None:
+            return manager_cls
+
+    spec_manager_map = cache.get("spec_manager_map", _CACHE_MISSING)
+    if spec_manager_map is _CACHE_MISSING:
+        try:
+            manager_module = import_module("vllm.v1.core.single_type_kv_cache_manager")
+            spec_manager_map = vars(manager_module)["spec_manager_map"]
+        except Exception as exc:
+            raise AssertionError(f"No manager registered for KVCacheSpec {type(spec)}") from exc
+        cache["spec_manager_map"] = spec_manager_map
+
+    try:
+        manager_cls = spec_manager_map[type(spec)]
+    except Exception as exc:
+        raise AssertionError(f"No manager registered for KVCacheSpec {type(spec)}") from exc
+    return manager_cls
+
+
+def _find_longest_cache_hit(
+    manager_cls: type[SingleTypeKVCacheManager],
+    **kwargs: Any,
+) -> tuple[list[KVCacheBlock], ...]:
+    try:
+        return manager_cls.find_longest_cache_hit(**kwargs)
+    except TypeError as exc:
+        if "drop_eagle_block" not in str(exc):
+            raise
+        kwargs["use_eagle"] = kwargs.pop("drop_eagle_block")
+        return manager_cls.find_longest_cache_hit(**kwargs)
+
+
+def _reachable_block_mask(
+    manager_cls: type[SingleTypeKVCacheManager],
+    **kwargs: Any,
+) -> list[bool] | None:
+    reachable_block_mask = getattr(manager_cls, "reachable_block_mask", None)
+    if reachable_block_mask is None:
+        return None
+    try:
+        return reachable_block_mask(**kwargs)
+    except TypeError as exc:
+        if "retention_interval" not in str(exc) and "num_prompt_tokens" not in str(exc):
+            logger.debug("KV cache manager does not support reachable_block_mask kwargs: %s", exc)
+            return reachable_block_mask(
+                start_block=kwargs["start_block"],
+                end_block=kwargs["end_block"],
+                alignment_tokens=kwargs["alignment_tokens"],
+                kv_cache_spec=kwargs["kv_cache_spec"],
+                use_eagle=kwargs["use_eagle"],
+            )
+        kwargs.pop("retention_interval", None)
+        kwargs.pop("num_prompt_tokens", None)
+        return reachable_block_mask(**kwargs)
+
+
+def _cache_family_granularity(block_size: int, cache_family: str | None) -> int:
+    if not cache_family or not cache_family.startswith("c"):
+        return block_size
+    ratio = cache_family[1:]
+    return block_size * int(ratio) if ratio.isdigit() else block_size
+
+
+def _uses_reachable_mask(cache_family: str | None) -> bool:
+    return cache_family in (None, "default", "c1")
+
+
+def _num_chunks(token_len: int, block_size: int) -> int:
+    return (token_len + block_size - 1) // block_size

vllm_ascend/distributed/kv_transfer/kv_pool/cpu_offload/cpu_kv_cache_manager.py

@@ -0,0 +1,199 @@
+import time
+from collections import defaultdict
+
+from vllm.logger import logger
+from vllm.utils.hashing import sha256
+from vllm.v1.core.block_pool import BlockPool
+from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock
+from vllm.v1.kv_cache_interface import KVCacheSpec
+from vllm.v1.metrics.stats import CachingMetrics, PrefixCacheStats
+from vllm.v1.request import Request
+
+from vllm_ascend.core.single_type_kv_cache_manager import get_manager_for_kv_cache_spec
+
+
+class CPUCacheStats:
+    def __init__(self, enable_prefix_caching: bool, log_stats: bool = False):
+        self.enable_prefix_caching = enable_prefix_caching
+        self.log_stats = log_stats
+        self.prefix_cache_stats = PrefixCacheStats() if log_stats else None
+        self.cpu_prefix_cache_metrics = CachingMetrics()
+        self.time_sec = int(time.time())
+
+    def log(self):
+        current_time_sec = int(time.time())
+        # Log the prefix cache hit rate every 10 seconds.
+        if current_time_sec - self.time_sec >= 10:
+            self.time_sec = current_time_sec
+            logger.info("CPU Prefix cache hit rate: %.1f%%", self.cpu_prefix_cache_metrics.hit_rate * 100)
+
+    def make_prefix_cache_stats(self) -> PrefixCacheStats | None:
+        """Get (and reset) the prefix cache stats.
+        Returns:
+            The current prefix caching stats, or None if logging is disabled.
+        """
+        if not self.log_stats:
+            return None
+        stats = self.prefix_cache_stats
+        self.prefix_cache_stats = PrefixCacheStats()
+        return stats
+
+    def update(self, num_tokens, num_computed_tokens):
+        # Note the function is called by scheduler
+        if self.log_stats and self.enable_prefix_caching:
+            assert self.prefix_cache_stats is not None
+            self.prefix_cache_stats.requests += 1
+            self.prefix_cache_stats.queries += num_tokens
+            self.prefix_cache_stats.hits += num_computed_tokens
+
+    def set_cache_stats(self, num_tokens, num_computed_tokens):
+        assert self.prefix_cache_stats is not None
+        self.prefix_cache_stats.hits = num_computed_tokens
+        self.prefix_cache_stats.queries = num_tokens
+        self.prefix_cache_stats.requests = 1
+
+
+class CPUKVCacheManager:
+    def __init__(
+        self,
+        kv_cache_spec: KVCacheSpec,
+        num_cpu_blocks: int,
+        caching_hash_algo: str = "builtin",
+        use_eagle: bool = False,
+        enable_kv_cache_events: bool = False,
+    ) -> None:
+        self.block_size = kv_cache_spec.block_size
+        self.num_cpu_blocks = num_cpu_blocks
+        self.caching_hash_fn = sha256 if caching_hash_algo == "sha256" else hash
+        self.use_eagle = use_eagle
+        self.block_pool = BlockPool(self.num_cpu_blocks, True, self.block_size, enable_kv_cache_events)
+        max_model_len = self.num_cpu_blocks * self.block_size
+        manager_kwargs = dict(
+            kv_cache_spec=kv_cache_spec,
+            block_pool=self.block_pool,
+            enable_caching=True,
+            kv_cache_group_id=0,
+            max_num_batched_tokens=max_model_len,
+            max_model_len=max_model_len,
+        )
+        manager_kwargs["scheduler_block_size"] = kv_cache_spec.block_size
+        self.single_type_manager = get_manager_for_kv_cache_spec(**manager_kwargs)
+        # Record kv block hashes, avoid redundant computation.
+        self.req_to_block_hashes: defaultdict[str, list[BlockHash]] = defaultdict(list)
+        # Record blocks touched in get_matched_num_and_touch().
+        self.req_to_computed_blocks: defaultdict[str, list[KVCacheBlock]] = defaultdict(list)
+        # Record the request that failed to allocate.
+        self.req_failed_to_allocate: defaultdict[str, bool] = defaultdict(bool)
+        self.req_to_num_tokens: defaultdict[str, int] = defaultdict(int)
+        self.cpu_cache_stats = CPUCacheStats(enable_prefix_caching=True, log_stats=True)
+        # Record request that will be free after finish sending
+        self.req_to_free: defaultdict[str, Request] = defaultdict(Request)
+
+    def get_matched_num_and_touch(self, request: Request) -> tuple[int, bool]:
+        # When the request requires prompt logprobs, we skip prefix caching.
+        if request.sampling_params.prompt_logprobs is not None:
+            return 0, False
+        request_id = request.request_id
+        # The block hashes for the request may already be computed
+        # if the scheduler has tried to schedule the request before.
+        block_hashes = self.req_to_block_hashes[request_id]
+        if not block_hashes:
+            block_hashes = request.block_hashes
+            self.req_to_block_hashes[request_id] = block_hashes
+        max_cache_hit_length = request.num_tokens - 1
+        eagle_kwarg = {"drop_eagle_block": self.use_eagle}
+        computed_blocks = self.single_type_manager.find_longest_cache_hit(
+            block_hashes=block_hashes,
+            max_length=max_cache_hit_length,
+            kv_cache_group_ids=[0],
+            block_pool=self.block_pool,
+            kv_cache_spec=self.single_type_manager.kv_cache_spec,
+            **eagle_kwarg,
+            alignment_tokens=self.block_size,
+        )
+        num_computed_tokens = len(computed_blocks[0]) * self.block_size
+        self.req_to_computed_blocks[request_id] = computed_blocks[0]
+        # We should touch these blocks in the concurrent scenarios.
+        self.block_pool.touch(computed_blocks)
+
+        # cup prefix cache status set and log
+        assert self.cpu_cache_stats is not None and self.cpu_cache_stats.prefix_cache_stats is not None
+        self.cpu_cache_stats.set_cache_stats(request.num_tokens, num_computed_tokens)
+        self.cpu_cache_stats.cpu_prefix_cache_metrics.observe(self.cpu_cache_stats.prefix_cache_stats)
+        self.cpu_cache_stats.log()
+
+        return num_computed_tokens, False
+
+    def _release_ahead_touch(self, request_id: str):
+        computed_blocks = self.req_to_computed_blocks[request_id]
+        if computed_blocks:
+            self.single_type_manager.block_pool.free_blocks(reversed(computed_blocks))
+            self.req_to_computed_blocks.pop(request_id, None)
+
+    def allocate_slots(self, req_to_num_tokens: dict[str, int], unallocated_req_ids: set[str]) -> dict[str, list[int]]:
+        for request_id in unallocated_req_ids:
+            self._free_slots(request_id)
+        req_to_new_blocks = {}
+        for request_id, num_tokens in req_to_num_tokens.items():
+            if self.req_failed_to_allocate[request_id]:
+                continue
+            new_computed_blocks = self.req_to_computed_blocks[request_id]
+            num_local_computed_tokens = len(new_computed_blocks) * self.block_size
+            num_blocks_to_allocate = self.single_type_manager.get_num_blocks_to_allocate(
+                request_id=request_id,
+                num_tokens=num_tokens,
+                new_computed_blocks=new_computed_blocks,
+                total_computed_tokens=num_local_computed_tokens,
+                num_tokens_main_model=num_tokens,
+            )
+            if num_blocks_to_allocate > self.block_pool.get_num_free_blocks():
+                self._release_ahead_touch(request_id)
+                self.req_failed_to_allocate[request_id] = True
+                continue
+            # Append the new computed blocks to the request blocks until now to
+            # avoid the case where the new blocks cannot be allocated.
+            self.single_type_manager.allocate_new_computed_blocks(
+                request_id,
+                new_computed_blocks,
+                num_local_computed_tokens=num_local_computed_tokens,
+                num_external_computed_tokens=0,
+            )
+            # Allocate new blocks but do not cache now.
+            new_blocks = self.single_type_manager.allocate_new_blocks(
+                request_id,
+                num_tokens,
+                num_tokens,
+            )
+            self.req_to_num_tokens[request_id] = num_tokens
+            # No need to release ref_cnt because we use officially.
+            self.req_to_computed_blocks.pop(request_id, None)
+            req_to_new_blocks[request_id] = [block.block_id for block in new_computed_blocks + new_blocks]
+        return req_to_new_blocks
+
+    def record_request_cache_and_free_slots(self, request: Request):
+        logger.debug("record_request_cache_and_free_slots for request %s in cpu_kv_cache_manager", request.request_id)
+        self.req_to_free[request.request_id] = request
+
+    def cache_and_free_slots(self, request_id: str):
+        logger.debug("Cache and free slots for request %s in cpu_kv_cache_manager", request_id)
+        if request_id not in self.req_to_free:
+            logger.error("request %s not in req_to_free, maybe bug!", request_id)
+            return
+        request = self.req_to_free[request_id]
+        if not self.req_failed_to_allocate[request_id]:
+            self.single_type_manager.cache_blocks(
+                request,
+                self.req_to_num_tokens[request_id],
+            )
+        self._free_slots(request_id)
+        logger.debug("delete request %s in cpu_kv_cache_manager req_to_free", request_id)
+        del self.req_to_free[request_id]
+
+    def _free_slots(self, request_id: str):
+        # This function is designed to be reentrant.
+        self._release_ahead_touch(request_id)
+        self.single_type_manager.free(request_id)
+        self.req_to_block_hashes.pop(request_id, None)
+        self.req_to_computed_blocks.pop(request_id, None)
+        self.req_failed_to_allocate.pop(request_id, None)
+        self.req_to_num_tokens.pop(request_id, None)

vllm_ascend/distributed/kv_transfer/kv_pool/cpu_offload/cpu_offload_connector.py

@@ -0,0 +1,448 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import copy
+import queue
+import threading
+import time
+from collections import defaultdict
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Optional
+
+import torch
+from vllm.config import VllmConfig, get_layers_from_vllm_config
+from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole
+from vllm.distributed.parallel_state import get_pp_group, get_tp_group
+from vllm.logger import logger
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.model_executor.layers.mamba.abstract import MambaBase
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheSpec
+
+from vllm_ascend.distributed.kv_transfer.kv_pool.cpu_offload.metadata import (
+    MetadataServer,
+    MetadataServerProc,
+    MLAConfig,
+)
+
+if TYPE_CHECKING:
+    from vllm.forward_context import ForwardContext
+    from vllm.v1.attention.backend import AttentionMetadata  # type: ignore
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.request import Request
+
+from vllm.model_executor.layers.attention import Attention, MLAAttention
+
+
+@dataclass
+class ReqMeta:
+    gpu_block_ids: list[int]
+    cpu_block_ids: list[int]
+    num_scheduled_tokens: int
+    num_computed_tokens: int
+    num_gpu_computed_tokens: int
+    num_cpu_computed_tokens: int
+
+    def update(self, other: "ReqMeta"):
+        self.gpu_block_ids.extend(other.gpu_block_ids)
+        self.cpu_block_ids.extend(other.cpu_block_ids)
+        self.num_scheduled_tokens = other.num_scheduled_tokens
+        self.num_computed_tokens = other.num_computed_tokens
+        self.num_gpu_computed_tokens = other.num_gpu_computed_tokens
+        self.num_cpu_computed_tokens = other.num_cpu_computed_tokens
+
+
+@dataclass
+class CPUOffloadingConnectorMetadata(KVConnectorMetadata):
+    requests: dict[str, ReqMeta]
+    finished_req_ids: set[str]
+
+
+class CPUOffloadingConnector(KVConnectorBase_V1):
+    def __init__(
+        self, vllm_config: VllmConfig, role: KVConnectorRole, kv_cache_config: Optional["KVCacheConfig"] = None
+    ):
+        self._connector_metadata = CPUOffloadingConnectorMetadata(requests={}, finished_req_ids=set())
+        if not vllm_config.cache_config.enable_prefix_caching:
+            self.connector_scheduler: CPUOffloadingConnectorScheduler | None = None
+            self.connector_worker: CPUOffloadingConnectorWorker | None = None
+        elif role == KVConnectorRole.SCHEDULER:
+            self.connector_scheduler = CPUOffloadingConnectorScheduler(vllm_config)
+            self.connector_worker = None
+        elif role == KVConnectorRole.WORKER:
+            self.connector_scheduler = None
+            self.connector_worker = CPUOffloadingConnectorWorker(vllm_config)
+
+    # ==============================
+    # Worker-side methods
+    # ==============================
+
+    def bind_connector_metadata(self, connector_metadata: KVConnectorMetadata) -> None:
+        if self.connector_worker is not None:
+            assert isinstance(connector_metadata, CPUOffloadingConnectorMetadata)
+            self.connector_worker.bind_connector_metadata(connector_metadata)
+
+    def clear_connector_metadata(self) -> None:
+        assert self.connector_worker is not None
+        self.connector_worker.clear_connector_metadata()
+
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+        if self.connector_worker is not None:
+            self.connector_worker.register_kv_caches(kv_caches)
+
+    def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
+        if self.connector_worker is not None:
+            self.connector_worker.start_load_kv()
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        if self.connector_worker is not None:
+            self.connector_worker.wait_for_layer_load()
+
+    def save_kv_layer(
+        self, layer_name: str, kv_layer: torch.Tensor, attn_metadata: "AttentionMetadata", **kwargs
+    ) -> None:
+        pass
+
+    def wait_for_save(self):
+        pass
+
+    def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str] | None, set[str] | None]:
+        assert self.connector_worker is not None
+        return self.connector_worker.get_finished(), None
+
+    # Scheduler-side methods
+    # ==============================
+
+    def get_num_new_matched_tokens(self, request: "Request", num_computed_tokens: int) -> tuple[int, bool]:
+        if self.connector_scheduler is not None:
+            return self.connector_scheduler.get_num_new_matched_tokens(request, num_computed_tokens)
+        return 0, False
+
+    def update_state_after_alloc(self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int):
+        if self.connector_scheduler is not None:
+            return self.connector_scheduler.update_state_after_alloc(request)
+
+    def build_connector_meta(self, scheduler_output: SchedulerOutput) -> KVConnectorMetadata:
+        if self.connector_scheduler is not None:
+            return self.connector_scheduler.build_connector_meta(scheduler_output)
+        return KVConnectorMetadata()
+
+    def request_finished(self, request: "Request", block_ids: list[int]) -> tuple[bool, dict[str, Any] | None]:
+        if self.connector_scheduler is not None:
+            self.connector_scheduler.request_finished(request)
+        return True, None
+
+
+class CPUOffloadingConnectorScheduler:
+    def __init__(self, vllm_config: VllmConfig):
+        logger.info("init CPUOffloadingConnectorScheduler")
+        self.vllm_config = vllm_config
+        self.block_size = vllm_config.cache_config.block_size
+        self.use_mla = vllm_config.model_config.use_mla
+        self.num_gpu_computed_tokens: dict[str, int] = {}
+        self.num_cpu_computed_tokens: dict[str, int] = {}
+        self.allocated_req_ids: set[str] = set()
+        self.finished_req_ids: list[str] = []
+        self.zmq_rpc_client = MetadataServer.ZMQRPCClient()
+        self.zmq_rpc_client.call("post_init")
+        if vllm_config.kv_transfer_config is not None:
+            self.swap_in_threshold = vllm_config.kv_transfer_config.get_from_extra_config("swap_in_threshold", 0)
+        else:
+            self.swap_in_threshold = 0
+        logger.info("swap_in_threshold: %s", self.swap_in_threshold)
+
+    def get_num_new_matched_tokens(self, ori_request: "Request", num_computed_tokens: int) -> tuple[int, bool]:
+        request = copy.deepcopy(ori_request)
+        request.get_hash_new_full_blocks = None
+        num_cpu_computed_tokens, load_async = self.zmq_rpc_client.call("get_matched_num_and_touch", request)
+        self.num_gpu_computed_tokens[request.request_id] = num_computed_tokens
+        self.num_cpu_computed_tokens[request.request_id] = num_cpu_computed_tokens
+        if num_cpu_computed_tokens - num_computed_tokens >= self.swap_in_threshold:
+            return num_cpu_computed_tokens - num_computed_tokens, load_async
+        else:
+            return 0, load_async
+
+    def update_state_after_alloc(self, request: "Request"):
+        self.allocated_req_ids.add(request.request_id)
+
+    def build_connector_meta(self, scheduler_output: SchedulerOutput) -> KVConnectorMetadata:
+        num_tokens = {}
+        # process scheduled_new_reqs
+        for req in scheduler_output.scheduled_new_reqs:
+            req_id = req.req_id
+            num_tokens[req_id] = req.num_computed_tokens + scheduler_output.num_scheduled_tokens[req_id]
+
+        # process scheduled_cached_reqs
+        cached_reqs = scheduler_output.scheduled_cached_reqs
+        for idx, req_id in enumerate(cached_reqs.req_ids):
+            num_tokens[req_id] = cached_reqs.num_computed_tokens[idx] + scheduler_output.num_scheduled_tokens[req_id]
+
+        unallocated_req_ids = set(
+            self.num_gpu_computed_tokens.keys() - self.allocated_req_ids - scheduler_output.num_scheduled_tokens.keys()
+        )
+        new_cpu_block_ids = self.zmq_rpc_client.call("allocate_slots", num_tokens, unallocated_req_ids)
+        metadata = CPUOffloadingConnectorMetadata(
+            requests={},
+            finished_req_ids=set(self.finished_req_ids),
+        )
+        for req in scheduler_output.scheduled_new_reqs:
+            req_id = req.req_id
+            gpu_block_ids = req.block_ids[0]
+            metadata.requests[req_id] = ReqMeta(
+                gpu_block_ids=[] if gpu_block_ids is None else gpu_block_ids,
+                cpu_block_ids=new_cpu_block_ids.get(req_id, []),
+                num_scheduled_tokens=scheduler_output.num_scheduled_tokens[req_id],
+                num_computed_tokens=req.num_computed_tokens,
+                num_gpu_computed_tokens=self.num_gpu_computed_tokens[req_id],
+                num_cpu_computed_tokens=self.num_cpu_computed_tokens[req_id],
+            )
+
+        for idx, req_id in enumerate(cached_reqs.req_ids):
+            gpu_block_ids = cached_reqs.new_block_ids[idx]
+            metadata.requests[req_id] = ReqMeta(
+                gpu_block_ids=[] if gpu_block_ids is None else gpu_block_ids,
+                cpu_block_ids=new_cpu_block_ids.get(req_id, []),
+                num_scheduled_tokens=scheduler_output.num_scheduled_tokens[req_id],
+                num_computed_tokens=cached_reqs.num_computed_tokens[idx],
+                num_gpu_computed_tokens=cached_reqs.num_computed_tokens[idx],
+                num_cpu_computed_tokens=cached_reqs.num_computed_tokens[idx],
+            )
+        self.num_gpu_computed_tokens.clear()
+        self.num_cpu_computed_tokens.clear()
+        self.allocated_req_ids.clear()
+        self.finished_req_ids.clear()
+        return metadata
+
+    def request_finished(self, ori_request: "Request"):
+        request = copy.deepcopy(ori_request)
+        request.get_hash_new_full_blocks = None
+        self.finished_req_ids.append(request.request_id)
+        # inform metadata server to record request, and free it after finish sending
+        self.zmq_rpc_client.call("record_request_cache_and_free_slots", request)
+
+
+class CPUOffloadingConnectorWorker:
+    def __init__(self, vllm_config: VllmConfig):
+        logger.info("init CPUOffloadingConnectorWorker")
+        self.vllm_config = vllm_config
+        self.block_size = vllm_config.cache_config.block_size
+        self.pp_rank = get_pp_group().rank_in_group
+        self.tp_group = get_tp_group()
+        self.tp_rank = self.tp_group.rank_in_group
+        self.tp_world_size = self.tp_group.world_size
+        self.use_mla = vllm_config.model_config.use_mla
+
+        self.requests: dict[str, ReqMeta] = {}
+        self.load_stream = torch.npu.Stream()
+        self.save_stream = torch.npu.Stream()
+        self.zmq_rpc_client = MetadataServer.ZMQRPCClient()
+        self.load_block_mapping: list[tuple[int, int]] = []
+        self.save_input_queue: queue.Queue[tuple[str, ReqMeta]] = queue.Queue()
+        self.save_output_queue: queue.Queue[str] = queue.Queue()
+        self.save_thread = threading.Thread(target=self._save_listener)
+        self.save_thread.start()
+        self.done_sending_count: defaultdict[str, int] = defaultdict(int)
+
+        # start metadata server to init cpu_kv_cache_manager and handle rpc requests
+        # all dp shared the same metadata server, only start the process on data_rank 0
+        if vllm_config.parallel_config.data_parallel_rank == 0 and self.tp_rank == 0 and self.pp_rank == 0:
+            config = VllmConfig()
+            config.cache_config = vllm_config.cache_config
+            config.parallel_config = vllm_config.parallel_config
+            config.kv_transfer_config = vllm_config.kv_transfer_config
+            self.init_metadata_server(config)
+        self._wait_for_metadata_process_start()
+
+    def init_metadata_server(self, vllm_config: VllmConfig):
+        self.metadata_thread = threading.Thread(
+            target=MetadataServerProc.run_metadata_server,
+            args=(vllm_config,),
+        )
+        self.metadata_thread.daemon = True
+        self.metadata_thread.start()
+
+    def _wait_for_metadata_process_start(self):
+        # TODO: wait for metadata server to start, add a rpc to check if ready
+        while True:
+            try:
+                if self.zmq_rpc_client.call("ready"):
+                    break
+            except Exception as e:
+                logger.info("wait for metadata server to start, error: %s", e)
+                time.sleep(1)
+
+    def bind_connector_metadata(self, connector_metadata: CPUOffloadingConnectorMetadata) -> None:
+        for req_id, req in connector_metadata.requests.items():
+            if req_id in self.requests:
+                self.requests[req_id].update(req)
+                req = self.requests[req_id]
+            else:
+                self.requests[req_id] = req
+            for i in range(req.num_gpu_computed_tokens // self.block_size, req.num_computed_tokens // self.block_size):
+                self.load_block_mapping.append((req.cpu_block_ids[i], req.gpu_block_ids[i]))
+        for req_id in connector_metadata.finished_req_ids:
+            if req_id in self.requests:
+                self.save_input_queue.put((req_id, self.requests[req_id]))
+
+    def clear_connector_metadata(self) -> None:
+        self.load_block_mapping.clear()
+
+    def register_kv_caches(self, kv_caches: dict[str, Sequence[torch.Tensor]]):
+        self.gpu_kv_caches = kv_caches
+        model_config = self.vllm_config.model_config
+        mla_config: MLAConfig | None = None
+        if model_config.use_mla:
+            mla_config = MLAConfig(
+                model_config.hf_text_config.kv_lora_rank, model_config.hf_text_config.qk_rope_head_dim
+            )
+        self.cpu_kv_caches = list(
+            self.zmq_rpc_client.call(
+                "init_cpu_kv_caches",
+                self.pp_rank,
+                self.tp_rank,
+                get_kv_cache_spec(self.vllm_config),
+                mla_config,
+            ).values()
+        )
+
+    def start_load_kv(self) -> None:
+        self.current_layer = 0
+        self.gpu_kv_caches_load_iter = iter(self.gpu_kv_caches.values())
+        self.load_kv_layer(0)
+
+    def wait_for_layer_load(self) -> None:
+        # TODO: Replace with `torch.npu.current_stream().wait_stream(self.load_stream)` after fixing the bug.
+        self.load_stream.synchronize()
+        self.current_layer += 1
+        self.load_kv_layer(self.current_layer)
+
+    def load_kv_layer(self, layer: int):
+        if layer == len(self.gpu_kv_caches):
+            return
+        gpu_kv_caches = next(self.gpu_kv_caches_load_iter)
+        cpu_kv_caches = self.cpu_kv_caches[layer]
+        with torch.npu.stream(self.load_stream):
+            for cpu_block_id, gpu_block_id in self.load_block_mapping:
+                for gpu_layer_part, cpu_layer_part in zip(gpu_kv_caches, cpu_kv_caches):
+                    gpu_layer_part[gpu_block_id].copy_(cpu_layer_part[cpu_block_id], non_blocking=True)
+
+    def get_finished(self) -> set[str]:
+        done_sending: set[str] = set()
+        while True:
+            try:
+                id = self.save_output_queue.get_nowait()
+            except queue.Empty:
+                break
+            done_sending.add(id)
+        for id in done_sending:
+            del self.requests[id]
+        if self.tp_world_size == 1:
+            return done_sending
+        if self.tp_rank == 0:
+            for req_id in done_sending:
+                self.done_sending_count[req_id] += 1
+            other_ranks_finished_ids: list[str] = []
+            for i in range(1, self.tp_world_size):
+                other_ranks_finished_ids.extend(self.tp_group.recv_object(src=i))
+            for req_id in other_ranks_finished_ids:
+                self.done_sending_count[req_id] += 1
+            all_done_sending: set[str] = set()
+            for req_id in list(self.done_sending_count.keys()):
+                if self.done_sending_count[req_id] == self.tp_world_size:
+                    del self.done_sending_count[req_id]
+                    all_done_sending.add(req_id)
+            # release cpu_kv_cache after request sending finished
+            # to avoid rpc blocking, use thread to call rpc asynchronously
+            sending_finished_thread = threading.Thread(target=self._sending_finished, args=(all_done_sending,))
+            sending_finished_thread.daemon = True
+            sending_finished_thread.start()
+
+            return all_done_sending
+        else:
+            self.tp_group.send_object(done_sending, dst=0)
+            return done_sending
+
+    def _sending_finished(self, all_done_sending):
+        for req_id in all_done_sending:
+            logger.debug("call cache_and_free_slots for req_id: %s", req_id)
+            self.zmq_rpc_client.call("cache_and_free_slots", req_id)
+
+    def _save_listener(self):
+        save_block_mapping = []
+        while True:
+            req_id, req = self.save_input_queue.get()
+            for i in range(
+                req.num_cpu_computed_tokens // self.block_size,
+                min((req.num_computed_tokens + req.num_scheduled_tokens) // self.block_size, len(req.cpu_block_ids)),
+            ):
+                save_block_mapping.append((req.gpu_block_ids[i], req.cpu_block_ids[i]))
+            with torch.npu.stream(self.save_stream):
+                # MLA: kv_layer is tuple[tensor, tensor] means (rope, nope).
+                # non-MLA: kv_layer is list[tensor], typically means [k, v].
+                if self.use_mla:
+                    start, step = self.tp_rank, self.tp_world_size
+                else:
+                    start, step = 0, 1
+                for i in range(start, len(save_block_mapping), step):
+                    gpu_block_id, cpu_block_id = save_block_mapping[i]
+                    for cpu_kv_caches, gpu_kv_caches in zip(self.cpu_kv_caches, self.gpu_kv_caches.values()):
+                        for cpu_layer_part, gpu_layer_part in zip(cpu_kv_caches, gpu_kv_caches):
+                            cpu_layer_part[cpu_block_id].copy_(gpu_layer_part[gpu_block_id], non_blocking=True)
+            self.save_stream.synchronize()
+            self.save_output_queue.put(req_id)
+            save_block_mapping.clear()
+
+
+# copied and modified from vllm_ascend/worker/model_runner_v1.py
+def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
+    """
+    Generates the KVCacheSpec by parsing the kv cache format from each
+    Attention module in the static forward context.
+    Returns:
+        KVCacheSpec: A dictionary mapping layer names to their KV cache
+        format. Layers that do not need KV cache are not included.
+    """
+    if has_ec_transfer() and get_ec_transfer().is_producer:
+        return {}
+
+    use_sparse = hasattr(vllm_config.model_config.hf_config, "index_topk")
+    if vllm_config.cache_config.cache_dtype == "auto":
+        kv_cache_dtype = vllm_config.model_config.dtype
+    else:
+        kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[vllm_config.cache_config.cache_dtype]
+
+    kv_cache_spec: dict[str, KVCacheSpec] = {}
+    attn_layers = get_layers_from_vllm_config(vllm_config, AttentionLayerBase)
+    # NOTE: Must process Attention/MLAAttention before MambaBase to maintain
+    # ordering expected by graph parameter update logic in attention backends.
+    mamba_layers: dict[str, MambaBase] = {}
+    for layer_name, attn_module in attn_layers.items():
+        if isinstance(attn_module, Attention):
+            if spec := attn_module.get_kv_cache_spec(vllm_config):
+                kv_cache_spec[layer_name] = spec
+
+        elif isinstance(attn_module, MLAAttention):
+            if use_sparse:
+                # TODO(cmq): This is a hack way to fix deepseek kvcache when
+                # using DSA. Fix the spec in vLLM is the final way.
+                block_size = vllm_config.cache_config.block_size
+                kv_cache_spec[layer_name] = FullAttentionSpec(
+                    block_size=block_size, num_kv_heads=1, head_size=attn_module.head_size, dtype=kv_cache_dtype
+                )
+            elif spec := attn_module.get_kv_cache_spec(vllm_config):
+                kv_cache_spec[layer_name] = spec
+
+        elif isinstance(attn_module, MambaBase):
+            mamba_layers[layer_name] = attn_module
+
+    if len(mamba_layers) > 0:
+        if vllm_config.cache_config.enable_prefix_caching:
+            raise NotImplementedError("Prefix caching is not supported for Mamba yet.")
+        for layer_name, mamba_module in mamba_layers.items():
+            if spec := mamba_module.get_kv_cache_spec(vllm_config):
+                kv_cache_spec[layer_name] = spec
+
+    return kv_cache_spec

vllm_ascend/distributed/kv_transfer/kv_pool/cpu_offload/metadata.py

@@ -0,0 +1,258 @@
+import math
+import os
+import pickle
+from collections.abc import Callable
+from dataclasses import dataclass
+from multiprocessing.shared_memory import SharedMemory
+from typing import Any
+
+import torch
+import vllm.envs as envs
+import zmq
+from vllm.config import KVTransferConfig, VllmConfig
+from vllm.logger import logger
+from vllm.utils.network_utils import make_zmq_socket
+from vllm.utils.torch_utils import get_dtype_size
+from vllm.v1.kv_cache_interface import AttentionSpec
+
+from vllm_ascend.core.kv_cache_interface import AscendMLAAttentionSpec
+from vllm_ascend.distributed.kv_transfer.kv_pool.cpu_offload.cpu_kv_cache_manager import CPUKVCacheManager
+
+
+@dataclass
+class MLAConfig:
+    nope_dim: int
+    rope_dim: int
+
+
+def get_cpu_offload_connector(vllm_config: VllmConfig) -> KVTransferConfig:
+    if vllm_config.kv_transfer_config is not None:
+        kv_transfer_config = vllm_config.kv_transfer_config
+        if kv_transfer_config.kv_connector == "CPUOffloadingConnector":
+            return kv_transfer_config
+        elif kv_transfer_config.kv_connector == "MultiConnector":
+            ktcs = kv_transfer_config.kv_connector_extra_config.get("connectors")
+            for ktc in ktcs:
+                kv_transfer_config = KVTransferConfig(**ktc)
+                if kv_transfer_config.kv_connector == "CPUOffloadingConnector":
+                    return kv_transfer_config
+    return None
+
+
+class MetadataServer:
+    METADATA_SERVER_ADDRESS = f"ipc://{envs.VLLM_RPC_BASE_PATH}/metadata.ipc"
+    DEFAULT_CPU_SWAP_SPACE_GB = 800
+
+    class ZMQRPCClient:
+        def __init__(self, identity=None):
+            if identity is None:
+                identity = f"worker-{os.getpid()}-{id(self)}"
+            logger.info("metadata client for worker %s started", identity)
+            self.ctx = zmq.Context()  # type: ignore
+            self.socket = make_zmq_socket(
+                self.ctx,
+                MetadataServer.METADATA_SERVER_ADDRESS,
+                zmq.DEALER,  # type: ignore
+                bind=False,
+                identity=identity.encode(),
+                linger=0,
+            )
+
+        def call(self, func_name: str, *args, **kwargs) -> Any:
+            request = (func_name, args, kwargs)
+            self.socket.send(b"", zmq.SNDMORE)  # type: ignore
+            self.socket.send(pickle.dumps(request))
+            _ = self.socket.recv()
+            response = pickle.loads(self.socket.recv())
+            result, error = response
+            if error:
+                logger.exception("call metadata sever error: %s", error)
+                raise error
+            if func_name == "init_cpu_kv_caches":
+                (memory_dict, layer_size, layer_dtype, mla_config) = result
+                # shared_memory_dict is recorded in self to close
+                self.shared_memory_dict = memory_dict
+                result = {}
+                for key, shm in memory_dict.items():
+                    tensor = torch.frombuffer(shm.buf, dtype=layer_dtype).reshape(layer_size)
+                    if mla_config is not None:
+                        tensor = tensor.split([mla_config.nope_dim, mla_config.rope_dim], dim=-1)
+                    result[key] = tensor
+            return result
+
+        def __del__(self):
+            # will be finalized by outer process
+            self.socket.close()
+            self.ctx.term()
+            if hasattr(self, "shared_memory_dict"):
+                for shm in self.shared_memory_dict.values():
+                    shm.close()
+
+    def __init__(self, vllm_config: VllmConfig):
+        self.world_size = vllm_config.parallel_config.world_size
+        self.pipeline_parallel_size = vllm_config.parallel_config.pipeline_parallel_size
+        kv_transfer_config = get_cpu_offload_connector(vllm_config)
+        assert kv_transfer_config is not None
+        available_memory_gb = kv_transfer_config.get_from_extra_config(
+            "cpu_swap_space_gb", MetadataServer.DEFAULT_CPU_SWAP_SPACE_GB
+        )
+        self.available_memory = available_memory_gb * 1024 * 1024 * 1024
+        logger.info("cpu swap space: %s bytes", self.available_memory)
+        self.ctx = zmq.Context()  # type: ignore
+        self.socket = make_zmq_socket(
+            self.ctx,
+            MetadataServer.METADATA_SERVER_ADDRESS,
+            zmq.ROUTER,  # type: ignore
+            bind=True,
+            linger=0,
+        )
+        self.functions: dict[str, Callable] = {
+            "init_cpu_kv_caches": self.init_cpu_kv_caches,
+            "post_init": self.post_init,
+            "ready": self.ready,
+        }
+        self.shared_memory = {}  # type: ignore
+        self.num_cpu_blocks = -1
+
+    @staticmethod
+    def _safe_create_shared_memory(name: str, size: int) -> SharedMemory:
+        try:
+            existing_shm = SharedMemory(name=name, create=False)
+            existing_shm.close()
+            existing_shm.unlink()
+        except FileNotFoundError:
+            pass
+        return SharedMemory(name=name, create=True, size=size)
+
+    def ready(self):
+        return True
+
+    def init_cpu_kv_caches(
+        self,
+        pp_rank: int,
+        tp_rank: int,
+        kv_cache_specs: dict[str, AttentionSpec],
+        mla_config: MLAConfig,
+    ) -> tuple[dict[str, SharedMemory], tuple[int, ...], torch.dtype, MLAConfig]:
+        logger.info("receive pp rank: %s, tp rank: %s", pp_rank, tp_rank)
+        # follow the assumption that each layer has the same spec
+        layer = next(iter(kv_cache_specs.values()))
+        assert all([layer.page_size_bytes == any.page_size_bytes for any in kv_cache_specs.values()])
+        use_mla = isinstance(layer, AscendMLAAttentionSpec)
+        # mla shares the same kv cache among different tp
+        if use_mla:
+            tp_rank = 0
+        if (pp_rank, tp_rank) in self.shared_memory:
+            return self.shared_memory[(pp_rank, tp_rank)]
+        available_memory = self.available_memory
+        shared_memory_dict = {}
+        if use_mla:
+            available_memory //= self.pipeline_parallel_size
+            available_memory //= len(kv_cache_specs)
+            num_blocks = available_memory // layer.page_size_bytes
+            layer_size = (num_blocks, layer.block_size, layer.num_kv_heads, layer.head_size)  # type: ignore
+        else:
+            available_memory //= self.world_size
+            available_memory //= len(kv_cache_specs)
+            num_blocks = available_memory // layer.page_size_bytes
+            layer_size = (2, num_blocks, layer.block_size, layer.num_kv_heads, layer.head_size)  # type: ignore
+        nbytes = math.prod(layer_size) * get_dtype_size(layer.dtype)
+        for layer_name in kv_cache_specs:
+            # only this format can share during ZeroMQ+pickle
+            shared_memory_dict[layer_name] = MetadataServer._safe_create_shared_memory(
+                f"cpu_kv_cache_{pp_rank}_{tp_rank}_{layer_name}", nbytes
+            )
+        if use_mla:
+            assert mla_config is not None
+            assert layer.head_size == mla_config.rope_dim + mla_config.nope_dim
+            self.shared_memory[(pp_rank, tp_rank)] = (shared_memory_dict, layer_size, layer.dtype, mla_config)
+        else:
+            self.shared_memory[(pp_rank, tp_rank)] = (shared_memory_dict, layer_size, layer.dtype, None)
+        if self.num_cpu_blocks == -1 or num_blocks < self.num_cpu_blocks:
+            self.num_cpu_blocks = num_blocks
+        self.layer = layer
+        return self.shared_memory[(pp_rank, tp_rank)]
+
+    def post_init(self):
+        # different processors in data parallel may call multiple times
+        if hasattr(self, "cpu_block_manager"):
+            return
+        # do shared_memory() at least once
+        logger.info("assign cpu num blocks: %s", self.num_cpu_blocks)
+        assert self.num_cpu_blocks >= 0
+        self.cpu_block_manager = CPUKVCacheManager(self.layer, self.num_cpu_blocks)
+        self.functions.update(
+            {
+                "get_matched_num_and_touch": self.cpu_block_manager.get_matched_num_and_touch,
+                "allocate_slots": self.cpu_block_manager.allocate_slots,
+                "record_request_cache_and_free_slots": self.cpu_block_manager.record_request_cache_and_free_slots,
+                "cache_and_free_slots": self.cpu_block_manager.cache_and_free_slots,
+            }
+        )
+
+    def serve_step(self):
+        client_id = self.socket.recv()
+        _ = self.socket.recv()
+        raw_msg = self.socket.recv()
+        try:
+            func_name, args, kwargs = pickle.loads(raw_msg)
+        except Exception as e:
+            response = (None, Exception(f"Invalid request: {str(e)}"))
+        else:
+            if func_name in self.functions:
+                try:
+                    result = self.functions[func_name](*args, **kwargs)
+                    response = (result, None)  # type: ignore
+                except Exception as e:
+                    logger.exception("metadata execute error: %s", e)
+                    response = (None, e)  # type: ignore
+            else:
+                response = (None, NameError(f"Function {func_name} not found"))
+        self.socket.send(client_id, zmq.SNDMORE)  # type: ignore
+        self.socket.send(b"", zmq.SNDMORE)  # type: ignore
+        self.socket.send(pickle.dumps(response))
+
+    def shutdown(self):
+        self.socket.close()
+        self.ctx.term()
+        socket_path = MetadataServer.METADATA_SERVER_ADDRESS.replace("ipc://", "")
+        if os.path.exists(socket_path):
+            os.remove(socket_path)
+        for cached in self.shared_memory.values():
+            for shm in cached[0].values():
+                shm.close()
+                shm.unlink()
+
+
+class MetadataServerProc:
+    @staticmethod
+    def run_metadata_server(vllm_config: VllmConfig):
+        if not vllm_config.cache_config.enable_prefix_caching or get_cpu_offload_connector(vllm_config) is None:
+            return
+
+        shutdown_requested = False
+
+        def _signal_handler(signum, frame):
+            nonlocal shutdown_requested
+            if not shutdown_requested:
+                shutdown_requested = True
+                raise SystemExit()
+
+        # Either SIGTERM or SIGINT will terminate the worker
+        # signal.signal(signal.SIGTERM, _signal_handler)
+        # signal.signal(signal.SIGINT, _signal_handler)
+        metadata_server: MetadataServer | None = None
+        try:
+            metadata_server = MetadataServer(vllm_config)
+            logger.info("Metadata server started.")
+            while True:
+                metadata_server.serve_step()
+        except SystemExit:
+            logger.info("Metadata server exiting.")
+            raise
+        except Exception as e:
+            logger.exception("Metadata server error: %s.", e)
+            raise e
+        finally:
+            if metadata_server is not None:
+                metadata_server.shutdown()

vllm_ascend/distributed/kv_transfer/kv_pool/lmcache_ascend_connector.py

@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+import lmcache_ascend  # noqa: F401
+from vllm.distributed.kv_transfer.kv_connector.v1.lmcache_connector import LMCacheConnectorV1
+
+__all__ = ["LMCacheConnectorV1"]

vllm_ascend/distributed/kv_transfer/kv_pool/recompute_cpu_offload/__init__.py

@@ -0,0 +1 @@
+

vllm_ascend/distributed/kv_transfer/kv_pool/recompute_cpu_offload/manager.py

@@ -0,0 +1,630 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Scheduler-side manager for recompute CPU offloading."""
+
+import contextlib
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from vllm.config import VllmConfig
+from vllm.distributed.kv_events import KVCacheEvent
+from vllm.logger import logger
+from vllm.utils.math_utils import cdiv
+from vllm.v1.core.block_pool import BlockPool
+from vllm.v1.core.kv_cache_coordinator import (
+    KVCacheCoordinator,
+    get_kv_cache_coordinator,
+)
+from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.kv_cache_interface import SlidingWindowSpec, UniformTypeKVCacheSpecs
+from vllm.v1.outputs import KVConnectorOutput
+
+from vllm_ascend.distributed.kv_transfer.kv_pool.recompute_cpu_offload.metadata import (
+    RecomputeCPUOffloadMetadata,
+    RecomputeCPUOffloadWorkerMetadata,
+)
+
+if TYPE_CHECKING:
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.core.kv_cache_utils import BlockHashWithGroupId, KVCacheBlock
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.request import Request
+
+
+@dataclass
+class TransferMeta:
+    gpu_block_ids: list[int]
+    cpu_block_ids: list[int]
+
+
+@dataclass
+class PreemptedRequestState:
+    req_id: str
+    cpu_block_ids: tuple[list[int], ...]
+    num_computed_tokens: int
+    store_transfer_meta: TransferMeta
+    store_event: int | None = None
+    load_event: int | None = None
+    load_transfer_meta: TransferMeta | None = None
+    load_start_tokens: int = 0
+    ready: bool = False
+    finished: bool = False
+
+
+class RecomputeCPUOffloadScheduler:
+    """Preserve preempted requests' KV blocks in CPU memory.
+
+    When offload prefix caching is enabled, full hashed blocks share CPU
+    blocks. Otherwise every offloaded block is private to its request.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        kv_cache_config: "KVCacheConfig | None",
+        cpu_capacity_bytes: int,
+        enable_offload_prefix_caching: bool = True,
+    ):
+        assert kv_cache_config is not None
+        self.vllm_config = vllm_config
+        self.enable_offload_prefix_caching = enable_offload_prefix_caching
+        self.cpu_kv_cache_config = self._derive_cpu_config(kv_cache_config, cpu_capacity_bytes)
+        self.num_cpu_blocks = self.cpu_kv_cache_config.num_blocks
+        self._group_is_sliding_window = self._get_group_is_sliding_window(kv_cache_config)
+        self.enable_kv_cache_events = (
+            vllm_config.kv_events_config is not None and vllm_config.kv_events_config.enable_kv_cache_events
+        )
+
+        logger.info(
+            "RecomputeCPUOffloadScheduler: allocating %d CPU blocks (%.2f GB) for recompute offload, prefix caching=%s",
+            self.num_cpu_blocks,
+            cpu_capacity_bytes / (1024**3),
+            self.enable_offload_prefix_caching,
+        )
+
+        dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
+        pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size
+        assert dcp_world_size == 1 and pcp_world_size == 1
+        scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes(kv_cache_config, vllm_config)
+        self.cpu_coordinator: KVCacheCoordinator = get_kv_cache_coordinator(
+            kv_cache_config=self.cpu_kv_cache_config,
+            max_model_len=vllm_config.model_config.max_model_len,
+            max_num_batched_tokens=(vllm_config.scheduler_config.max_num_batched_tokens),
+            use_eagle=False,
+            enable_caching=self.enable_offload_prefix_caching,
+            enable_kv_cache_events=self.enable_kv_cache_events,
+            dcp_world_size=dcp_world_size,
+            pcp_world_size=pcp_world_size,
+            scheduler_block_size=scheduler_block_size,
+            hash_block_size=hash_block_size,
+        )
+        self.cpu_block_pool: BlockPool = self.cpu_coordinator.block_pool
+        self._gpu_block_pool: BlockPool | None = None
+
+        self._preempted_req_states: dict[str, PreemptedRequestState] = {}
+        self._preempt_store_event_to_reqs: dict[int, list[str]] = {}
+        self._preempt_store_event_to_blocks: dict[int, TransferMeta] = {}
+        self._preempt_load_event_to_reqs: dict[int, list[str]] = {}
+
+        # Hash blocks created before build_connector_meta() are shared by all
+        # requests preempted in the same scheduling step.
+        self._pending_hash_blocks: dict[BlockHashWithGroupId, KVCacheBlock] = {}
+
+        self._load_event_counter = 0
+        self._store_event_counter = 0
+        self._expected_worker_count = vllm_config.parallel_config.world_size
+        self._store_event_pending_counts: dict[int, int] = {}
+
+    @staticmethod
+    def _get_group_is_sliding_window(kv_cache_config: "KVCacheConfig") -> list[bool]:
+        group_is_sliding_window: list[bool] = []
+        for group in kv_cache_config.kv_cache_groups:
+            if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs):
+                group_is_sliding_window.append(
+                    any(isinstance(spec, SlidingWindowSpec) for spec in group.kv_cache_spec.kv_cache_specs.values())
+                )
+            else:
+                group_is_sliding_window.append(isinstance(group.kv_cache_spec, SlidingWindowSpec))
+        return group_is_sliding_window
+
+    @staticmethod
+    def _derive_cpu_config(gpu_config: "KVCacheConfig", cpu_capacity_bytes: int) -> "KVCacheConfig":
+        from vllm.v1.kv_cache_interface import KVCacheConfig as KVCacheConfigCls
+        from vllm.v1.kv_cache_interface import KVCacheTensor
+
+        assert gpu_config.kv_cache_tensors
+        gpu_kv_cache_tensors = []
+        for t in gpu_config.kv_cache_tensors:
+            if t.shared_by:
+                gpu_kv_cache_tensors.append(t)
+        gpu_total_bytes = sum(t.size for t in gpu_kv_cache_tensors)
+        num_gpu_blocks = gpu_config.num_blocks
+        num_cpu_blocks = max(1, num_gpu_blocks * cpu_capacity_bytes // gpu_total_bytes)
+        cpu_tensors = [
+            KVCacheTensor(
+                size=t.size // num_gpu_blocks * num_cpu_blocks,
+                shared_by=list(t.shared_by),
+            )
+            for t in gpu_kv_cache_tensors
+        ]
+        return KVCacheConfigCls(
+            num_blocks=num_cpu_blocks,
+            kv_cache_tensors=cpu_tensors,
+            kv_cache_groups=gpu_config.kv_cache_groups,
+        )
+
+    def _align_group_block_ids(
+        self,
+        group_idx: int,
+        group_block_ids: list[int],
+        logical_num_blocks: int,
+    ) -> list[int]:
+        if logical_num_blocks <= 0:
+            return []
+        aligned_group_block_ids = list(group_block_ids)
+        if self._group_is_sliding_window[group_idx] and len(aligned_group_block_ids) < logical_num_blocks:
+            aligned_group_block_ids = [0] * (
+                logical_num_blocks - len(aligned_group_block_ids)
+            ) + aligned_group_block_ids
+        return aligned_group_block_ids[:logical_num_blocks]
+
+    def bind_gpu_block_pool(self, gpu_block_pool: BlockPool) -> None:
+        self._gpu_block_pool = gpu_block_pool
+
+    def has_preempted_request(self, req_id: str) -> bool:
+        return req_id in self._preempted_req_states
+
+    def get_num_new_matched_tokens(self, request: "Request", num_computed_tokens: int) -> tuple[int | None, bool]:
+        state = self._preempted_req_states.get(request.request_id)
+        if state is None:
+            return 0, False
+        if not state.ready:
+            return None, False
+
+        restorable_tokens = min(state.num_computed_tokens, request.num_tokens)
+        hit_length = max(0, restorable_tokens - num_computed_tokens)
+        if hit_length <= 0:
+            self._cleanup_preempt_cache_request(request.request_id)
+            return 0, False
+
+        state.load_start_tokens = num_computed_tokens
+        logger.debug(
+            "Recompute offload cache hit for request %s: load_start=%d, load_tokens=%d, stored_tokens=%d.",
+            request.request_id,
+            num_computed_tokens,
+            hit_length,
+            state.num_computed_tokens,
+        )
+        return hit_length, True
+
+    def update_state_after_alloc(
+        self,
+        request: "Request",
+        blocks: "KVCacheBlocks",
+        num_external_tokens: int,
+    ) -> None:
+        if num_external_tokens <= 0:
+            return
+        prepared = self._prepare_preempt_load_after_alloc(
+            request,
+            blocks.get_block_ids(),
+            num_external_tokens,
+        )
+        if not prepared:
+            raise RuntimeError(
+                "Failed to prepare recompute H2D load after KV block "
+                f"allocation: req_id={request.request_id}, "
+                f"num_external_tokens={num_external_tokens}"
+            )
+
+    def update_state_before_preempt(
+        self,
+        request: "Request",
+        block_ids: tuple[list[int], ...],
+        num_computed_tokens: int,
+    ) -> bool:
+        if request.request_id in self._preempted_req_states:
+            return True
+        return self._create_preempt_state(
+            request.request_id,
+            block_ids,
+            num_computed_tokens,
+        )
+
+    def _create_preempt_state(
+        self,
+        req_id: str,
+        block_ids_by_group: tuple[list[int], ...],
+        num_computed_tokens: int,
+    ) -> bool:
+        if num_computed_tokens <= 0 or self._gpu_block_pool is None:
+            return False
+
+        kv_cache_groups = self.cpu_kv_cache_config.kv_cache_groups
+        group_gpu_blocks: list[list[KVCacheBlock | None]] = []
+        group_gpu_hashes: list[list[BlockHashWithGroupId | None]] = []
+        missing_hashes: set[BlockHashWithGroupId] = set()
+        num_unhashed = 0
+
+        for g, group_gpu_ids in enumerate(block_ids_by_group):
+            group_block_size = kv_cache_groups[g].kv_cache_spec.block_size
+            logical_num_blocks = cdiv(num_computed_tokens, group_block_size)
+            aligned_group_gpu_ids = self._align_group_block_ids(g, group_gpu_ids, logical_num_blocks)
+            eviction_group_gpu_ids = self._align_group_block_ids(
+                g,
+                group_gpu_ids,
+                max(logical_num_blocks, len(group_gpu_ids)),
+            )
+            gpu_blocks: list[KVCacheBlock | None] = []
+            effective_hashes: list[BlockHashWithGroupId | None] = []
+
+            for block_idx, block_id in enumerate(eviction_group_gpu_ids):
+                if block_id <= 0:
+                    continue
+                gpu_block = self._gpu_block_pool.blocks[block_id]
+                block_is_computed = (block_idx + 1) * group_block_size <= num_computed_tokens
+                if not block_is_computed and gpu_block.block_hash is not None:
+                    # allocate_slots() may assign a hash using tokens planned
+                    # for this scheduling step. If the request is then
+                    # preempted before forward, that block does not contain the
+                    # hashed KV and must not remain in the GPU prefix cache.
+                    self._gpu_block_pool._maybe_evict_cached_block(gpu_block)
+
+            for block_idx, block_id in enumerate(aligned_group_gpu_ids):
+                if block_id <= 0:
+                    gpu_blocks.append(None)
+                    effective_hashes.append(None)
+                    continue
+
+                gpu_block = self._gpu_block_pool.blocks[block_id]
+                block_is_computed = (block_idx + 1) * group_block_size <= num_computed_tokens
+                block_hash = gpu_block.block_hash if block_is_computed and self.enable_offload_prefix_caching else None
+                gpu_blocks.append(gpu_block)
+                effective_hashes.append(block_hash)
+                if block_hash is None:
+                    num_unhashed += 1
+                elif (
+                    self.cpu_block_pool.cached_block_hash_to_block.get_one_block(block_hash) is None
+                    and block_hash not in self._pending_hash_blocks
+                ):
+                    missing_hashes.add(block_hash)
+            group_gpu_blocks.append(gpu_blocks)
+            group_gpu_hashes.append(effective_hashes)
+
+        num_needed = num_unhashed + len(missing_hashes)
+        if not any(any(gpu_block is not None for gpu_block in group) for group in group_gpu_blocks):
+            return False
+        if num_needed > self.cpu_block_pool.get_num_free_blocks():
+            logger.warning(
+                "Skip recompute offload for request %s: CPU cache has %d free blocks, but %d new blocks are required.",
+                req_id,
+                self.cpu_block_pool.get_num_free_blocks(),
+                num_needed,
+            )
+            return False
+
+        cpu_block_iter = iter(self.cpu_block_pool.get_new_blocks(num_needed))
+        cpu_block_ids_by_group: list[list[int]] = []
+        store_gpu_block_ids: list[int] = []
+        store_cpu_block_ids: list[int] = []
+        waiting_for_store = False
+
+        for gpu_blocks, effective_hashes in zip(group_gpu_blocks, group_gpu_hashes):
+            group_cpu_ids: list[int] = []
+            for gpu_block, block_hash in zip(gpu_blocks, effective_hashes):
+                if gpu_block is None:
+                    group_cpu_ids.append(0)
+                    continue
+
+                cpu_block = None
+
+                if block_hash is not None:
+                    cpu_block = self.cpu_block_pool.cached_block_hash_to_block.get_one_block(block_hash)
+                    if cpu_block is not None:
+                        self.cpu_block_pool.touch([cpu_block])
+                    else:
+                        cpu_block = self._pending_hash_blocks.get(block_hash)
+                        if cpu_block is not None:
+                            self.cpu_block_pool.touch([cpu_block])
+                            waiting_for_store = True
+                        else:
+                            cpu_block = next(cpu_block_iter)
+                            cpu_block._block_hash = block_hash
+                            self._pending_hash_blocks[block_hash] = cpu_block
+                            store_gpu_block_ids.append(gpu_block.block_id)
+                            store_cpu_block_ids.append(cpu_block.block_id)
+                            waiting_for_store = True
+                else:
+                    cpu_block = next(cpu_block_iter)
+                    store_gpu_block_ids.append(gpu_block.block_id)
+                    store_cpu_block_ids.append(cpu_block.block_id)
+                    waiting_for_store = True
+
+                group_cpu_ids.append(cpu_block.block_id)
+            cpu_block_ids_by_group.append(group_cpu_ids)
+
+        store_transfer = TransferMeta(store_gpu_block_ids, store_cpu_block_ids)
+        self._preempted_req_states[req_id] = PreemptedRequestState(
+            req_id=req_id,
+            cpu_block_ids=tuple(cpu_block_ids_by_group),
+            num_computed_tokens=num_computed_tokens,
+            store_transfer_meta=store_transfer,
+            ready=not waiting_for_store,
+        )
+        logger.info(
+            "Created recompute offload state for request %s: "
+            "computed_tokens=%d, cpu_blocks=%d, store_blocks=%d, "
+            "ready=%s.",
+            req_id,
+            num_computed_tokens,
+            sum(len(ids) for ids in cpu_block_ids_by_group),
+            len(store_cpu_block_ids),
+            not waiting_for_store,
+        )
+
+        return True
+
+    def _prepare_preempt_store_specs(
+        self,
+    ) -> tuple[list[int], list[int], list[str]]:
+        gpu_block_ids: list[int] = []
+        cpu_block_ids: list[int] = []
+        req_ids: list[str] = []
+
+        for req_id, state in self._preempted_req_states.items():
+            if state.store_event is not None or state.ready:
+                continue
+            gpu_block_ids.extend(state.store_transfer_meta.gpu_block_ids)
+            cpu_block_ids.extend(state.store_transfer_meta.cpu_block_ids)
+            req_ids.append(req_id)
+        return gpu_block_ids, cpu_block_ids, req_ids
+
+    def _prepare_preempt_load_after_alloc(
+        self,
+        request: "Request",
+        block_ids_by_group: tuple[list[int], ...],
+        num_external_tokens: int,
+    ) -> bool:
+        state = self._preempted_req_states.get(request.request_id)
+        if state is None or not state.ready:
+            return False
+
+        load_start_tokens = state.load_start_tokens
+        load_end_tokens = min(
+            load_start_tokens + num_external_tokens,
+            state.num_computed_tokens,
+        )
+        if load_end_tokens <= load_start_tokens:
+            return False
+
+        if len(block_ids_by_group) != len(state.cpu_block_ids):
+            raise RuntimeError(
+                "Recompute H2D KV group count mismatch: "
+                f"req_id={request.request_id}, "
+                f"gpu_groups={len(block_ids_by_group)}, "
+                f"cpu_groups={len(state.cpu_block_ids)}"
+            )
+
+        gpu_block_ids: list[int] = []
+        cpu_block_ids: list[int] = []
+        for g, group_cpu_ids in enumerate(state.cpu_block_ids):
+            group_block_size = self.cpu_kv_cache_config.kv_cache_groups[g].kv_cache_spec.block_size
+            start_block = load_start_tokens // group_block_size
+            end_block = min(
+                len(group_cpu_ids),
+                len(
+                    self._align_group_block_ids(
+                        g,
+                        block_ids_by_group[g],
+                        max(
+                            cdiv(load_end_tokens, group_block_size),
+                            len(block_ids_by_group[g]),
+                        ),
+                    )
+                ),
+                cdiv(load_end_tokens, group_block_size),
+            )
+            if end_block == start_block:
+                continue
+            if end_block < start_block:
+                raise RuntimeError(
+                    "Recompute H2D produced an empty block range: "
+                    f"req_id={request.request_id}, group={g}, "
+                    f"start_block={start_block}, end_block={end_block}, "
+                    f"gpu_blocks={len(block_ids_by_group[g])}, "
+                    f"cpu_blocks={len(group_cpu_ids)}"
+                )
+
+            aligned_group_gpu_ids = self._align_group_block_ids(
+                g,
+                block_ids_by_group[g],
+                end_block,
+            )
+            for block_idx in range(start_block, end_block):
+                cpu_block_id = group_cpu_ids[block_idx]
+                gpu_block_id = aligned_group_gpu_ids[block_idx]
+                if cpu_block_id <= 0 or gpu_block_id <= 0:
+                    continue
+                cpu_block_ids.append(cpu_block_id)
+                gpu_block_ids.append(gpu_block_id)
+
+        if not cpu_block_ids or len(cpu_block_ids) != len(gpu_block_ids):
+            raise RuntimeError(
+                "Recompute H2D block mapping is incomplete: "
+                f"req_id={request.request_id}, "
+                f"gpu_blocks={len(gpu_block_ids)}, "
+                f"cpu_blocks={len(cpu_block_ids)}"
+            )
+
+        assert self._gpu_block_pool is not None
+        self._gpu_block_pool.touch([self._gpu_block_pool.blocks[block_id] for block_id in gpu_block_ids])
+        state.load_transfer_meta = TransferMeta(gpu_block_ids, cpu_block_ids)
+        logger.info(
+            "Prepared recompute offload H2D load for request %s: tokens=[%d, %d), blocks=%d.",
+            request.request_id,
+            load_start_tokens,
+            load_end_tokens,
+            len(gpu_block_ids),
+        )
+        return True
+
+    def build_connector_meta(
+        self,
+        scheduler_output: SchedulerOutput,
+    ) -> RecomputeCPUOffloadMetadata:
+        store_event = -1
+        store_gpu, store_cpu, store_req_ids = self._prepare_preempt_store_specs()
+        if store_gpu:
+            store_event = self._store_event_counter
+            self._store_event_counter += 1
+            self._preempt_store_event_to_blocks[store_event] = TransferMeta(store_gpu, store_cpu)
+            self._preempt_store_event_to_reqs[store_event] = store_req_ids
+            for req_id in store_req_ids:
+                self._preempted_req_states[req_id].store_event = store_event
+            self._pending_hash_blocks.clear()
+
+        load_event = -1
+        load_gpu: list[int] = []
+        load_cpu: list[int] = []
+        load_req_ids: list[str] = []
+        for req_id, state in self._preempted_req_states.items():
+            if state.load_transfer_meta is None or state.load_event is not None:
+                continue
+            load_gpu.extend(state.load_transfer_meta.gpu_block_ids)
+            load_cpu.extend(state.load_transfer_meta.cpu_block_ids)
+            load_req_ids.append(req_id)
+
+        if load_req_ids:
+            load_event = self._load_event_counter
+            self._load_event_counter += 1
+            for req_id in load_req_ids:
+                self._preempted_req_states[req_id].load_event = load_event
+            self._preempt_load_event_to_reqs[load_event] = load_req_ids
+
+        return RecomputeCPUOffloadMetadata(
+            need_flush=bool(scheduler_output.preempted_req_ids),
+            preempt_store_event=store_event,
+            preempt_store_gpu_blocks=store_gpu,
+            preempt_store_cpu_blocks=store_cpu,
+            preempt_load_event=load_event,
+            preempt_load_gpu_blocks=load_gpu,
+            preempt_load_cpu_blocks=load_cpu,
+            preempt_load_event_to_reqs=self._preempt_load_event_to_reqs,
+        )
+
+    def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
+        for req_id in list(connector_output.finished_recving or []):
+            if req_id in self._preempted_req_states:
+                self._cleanup_preempt_load_request(req_id)
+
+        meta = connector_output.kv_connector_worker_meta
+        if not isinstance(meta, RecomputeCPUOffloadWorkerMetadata):
+            return
+        for event_idx, count in meta.completed_store_events.items():
+            total = self._store_event_pending_counts.get(event_idx, 0) + count
+            if total >= self._expected_worker_count:
+                self._store_event_pending_counts.pop(event_idx, None)
+                self._process_preempt_store_event(event_idx)
+            else:
+                self._store_event_pending_counts[event_idx] = total
+
+    def _process_preempt_store_event(self, event_idx: int) -> None:
+        transfer = self._preempt_store_event_to_blocks.pop(event_idx)
+        req_ids = self._preempt_store_event_to_reqs.pop(event_idx, [])
+
+        for cpu_block_id in transfer.cpu_block_ids:
+            cpu_block = self.cpu_block_pool.blocks[cpu_block_id]
+            block_hash = cpu_block.block_hash
+            if block_hash is None:
+                continue
+            cached_block = self.cpu_block_pool.cached_block_hash_to_block.get_one_block(block_hash)
+            if cached_block is None:
+                self.cpu_block_pool.cached_block_hash_to_block.insert(block_hash, cpu_block)
+            elif cached_block.block_id != cpu_block.block_id:
+                cpu_block.reset_hash()
+
+        for req_id in req_ids:
+            state = self._preempted_req_states.get(req_id)
+            if state is not None:
+                state.ready = True
+                if state.finished:
+                    self._cleanup_preempt_cache_request(req_id)
+
+    def has_pending_transfers(self) -> bool:
+        return bool(
+            self._store_event_pending_counts
+            or self._preempt_store_event_to_blocks
+            or any(
+                not state.ready or state.load_transfer_meta is not None for state in self._preempted_req_states.values()
+            )
+        )
+
+    def reset_cache(self) -> bool:
+        if self.has_pending_transfers():
+            logger.warning(
+                "Failed to reset recompute offload cache because transfers or request states are still pending."
+            )
+            return False
+        for req_id in list(self._preempted_req_states):
+            self._cleanup_preempt_cache_request(req_id)
+        self._preempt_store_event_to_reqs.clear()
+        self._preempt_store_event_to_blocks.clear()
+        self._preempt_load_event_to_reqs.clear()
+        self._pending_hash_blocks.clear()
+        return self.cpu_block_pool.reset_prefix_cache()
+
+    def request_finished(
+        self,
+        request: "Request",
+        block_ids: list[int],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        state = self._preempted_req_states.get(request.request_id)
+        if state is not None and state.load_event is None:
+            if state.ready:
+                self._cleanup_preempt_cache_request(request.request_id)
+            else:
+                state.finished = True
+        return False, None
+
+    def request_finished_all_groups(
+        self,
+        request: "Request",
+        block_ids: tuple[list[int], ...],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        return self.request_finished(request, block_ids=[])
+
+    def _cleanup_preempt_load_request(self, req_id: str) -> None:
+        state = self._preempted_req_states.get(req_id)
+        if state is None:
+            return
+
+        if state.load_event is not None:
+            reqs = self._preempt_load_event_to_reqs.get(state.load_event)
+            if reqs is not None:
+                with contextlib.suppress(ValueError):
+                    reqs.remove(req_id)
+                if not reqs:
+                    self._preempt_load_event_to_reqs.pop(state.load_event, None)
+
+        if state.load_transfer_meta is not None:
+            assert self._gpu_block_pool is not None
+            self._gpu_block_pool.free_blocks(
+                self._gpu_block_pool.blocks[block_id] for block_id in state.load_transfer_meta.gpu_block_ids
+            )
+        self._cleanup_preempt_cache_request(req_id)
+
+    def _cleanup_preempt_cache_request(self, req_id: str) -> None:
+        state = self._preempted_req_states.pop(req_id, None)
+        if state is None:
+            return
+        self.cpu_block_pool.free_blocks(
+            self.cpu_block_pool.blocks[block_id]
+            for group_cpu_ids in state.cpu_block_ids
+            for block_id in group_cpu_ids
+            if block_id > 0
+        )
+
+    def take_events(self) -> Iterable[KVCacheEvent]:
+        return self.cpu_block_pool.take_events()

vllm_ascend/distributed/kv_transfer/kv_pool/recompute_cpu_offload/metadata.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Metadata for RecomputeCPUOffloadConnector."""
+
+from dataclasses import dataclass, field
+
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorMetadata,
+    KVConnectorWorkerMetadata,
+)
+
+INVALID_JOB_ID = -1
+
+
+@dataclass
+class RecomputeCPUOffloadMetadata(KVConnectorMetadata):
+    """Recompute offload transfers passed from scheduler to worker."""
+
+    # Whether any requests were preempted this step and need flush pending transfers.
+    need_flush: bool = False
+
+    # Store blocks of newly preempted requests before their GPU blocks can
+    # be reused. The list may include a final partial block without a hash.
+    preempt_store_event: int = INVALID_JOB_ID
+    preempt_store_gpu_blocks: list[int] = field(default_factory=list)
+    preempt_store_cpu_blocks: list[int] = field(default_factory=list)
+
+    # Preemption load event. Used when a previously preempted request resumes.
+    preempt_load_event: int = INVALID_JOB_ID
+    preempt_load_gpu_blocks: list[int] = field(default_factory=list)
+    preempt_load_cpu_blocks: list[int] = field(default_factory=list)
+    preempt_load_event_to_reqs: dict[int, list[str]] = field(default_factory=dict)
+
+
+@dataclass
+class RecomputeCPUOffloadWorkerMetadata(KVConnectorWorkerMetadata):
+    """Worker -> Scheduler metadata for completed store events.
+
+    Each worker reports {event_idx: 1} for newly completed stores.
+    ``aggregate()`` sums counts across workers within a step.
+    The scheduler-side manager accumulates across steps and processes
+    a store completion only when count reaches ``world_size``.
+    """
+
+    completed_store_events: dict[int, int]
+
+    def aggregate(self, other: "KVConnectorWorkerMetadata") -> "KVConnectorWorkerMetadata":
+        assert isinstance(other, RecomputeCPUOffloadWorkerMetadata)
+        merged = dict(self.completed_store_events)
+        for k, v in other.completed_store_events.items():
+            merged[k] = merged.get(k, 0) + v
+        return RecomputeCPUOffloadWorkerMetadata(completed_store_events=merged)

vllm_ascend/distributed/kv_transfer/kv_pool/recompute_cpu_offload/recompute_cpu_offload_connector.py

@@ -0,0 +1,246 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""RecomputeCPUOffloadConnector: minimal CPU KV cache offloading."""
+
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Any
+
+import torch
+from vllm.config import VllmConfig
+from vllm.distributed.kv_events import KVCacheEvent
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorBase_V1,
+    KVConnectorMetadata,
+    KVConnectorRole,
+    SupportsHMA,
+)
+from vllm.logger import logger
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import KVConnectorOutput
+
+from vllm_ascend.distributed.kv_transfer.kv_pool.recompute_cpu_offload.manager import (
+    RecomputeCPUOffloadScheduler,
+)
+from vllm_ascend.distributed.kv_transfer.kv_pool.recompute_cpu_offload.metadata import (
+    RecomputeCPUOffloadMetadata,
+)
+from vllm_ascend.distributed.kv_transfer.kv_pool.recompute_cpu_offload.worker import (
+    RecomputeCPUOffloadWorker,
+)
+
+if TYPE_CHECKING:
+    from vllm.forward_context import ForwardContext
+    from vllm.v1.attention.backend import AttentionMetadata
+    from vllm.v1.core.block_pool import BlockPool
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.request import Request
+
+# Default CPU capacity: 8 GB
+DEFAULT_CPU_CAPACITY_BYTES = 8 * (1024**3)
+
+
+class RecomputeCPUOffloadConnectorV1(KVConnectorBase_V1, SupportsHMA):
+    """CPU KV cache preservation for recompute-preempted requests."""
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        role: KVConnectorRole,
+        kv_cache_config: "KVCacheConfig | None" = None,
+    ):
+        super().__init__(vllm_config, role, kv_cache_config)
+
+        extra_config = self._kv_transfer_config.kv_connector_extra_config or {}
+        cpu_capacity_bytes = int(extra_config.get("cpu_bytes_to_use", DEFAULT_CPU_CAPACITY_BYTES))
+        enable_offload_prefix_caching = extra_config.get("enable_offload_prefix_caching", False)
+        if not isinstance(enable_offload_prefix_caching, bool):
+            raise ValueError(f"enable_offload_prefix_caching must be a boolean, got {enable_offload_prefix_caching!r}")
+        world_size = vllm_config.parallel_config.world_size
+        cpu_capacity_per_rank = cpu_capacity_bytes // world_size
+        if "cpu_bytes_to_use_per_rank" in extra_config:
+            explicit = int(extra_config["cpu_bytes_to_use_per_rank"])
+            if explicit != cpu_capacity_per_rank:
+                logger.warning(
+                    "cpu_bytes_to_use_per_rank (%.2f GB) != "
+                    "cpu_bytes_to_use/world_size (%.2f GB). Using per-rank value.",
+                    explicit / (1024**3),
+                    cpu_capacity_per_rank / (1024**3),
+                )
+            cpu_capacity_per_rank = explicit
+
+        self.scheduler_manager: RecomputeCPUOffloadScheduler | None = None
+        self.worker_handler: RecomputeCPUOffloadWorker | None = None
+
+        logger.info(
+            "RecomputeCPUOffloadConnector: role=%s, per_rank=%.2f GB, world_size=%d, offload_prefix_caching=%s",
+            role.name,
+            cpu_capacity_per_rank / (1024**3),
+            world_size,
+            enable_offload_prefix_caching,
+        )
+
+        if role == KVConnectorRole.SCHEDULER:
+            self.scheduler_manager = RecomputeCPUOffloadScheduler(
+                vllm_config,
+                kv_cache_config,
+                cpu_capacity_per_rank,
+                enable_offload_prefix_caching,
+            )
+        elif role == KVConnectorRole.WORKER:
+            self.worker_handler = RecomputeCPUOffloadWorker(
+                vllm_config,
+                kv_cache_config,
+                cpu_capacity_per_rank,
+            )
+
+    # --- Worker-side methods ---
+
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None:
+        if self.worker_handler is not None:
+            self.worker_handler.register_kv_caches(kv_caches)
+
+    def bind_connector_metadata(
+        self,
+        connector_metadata: KVConnectorMetadata,
+    ) -> None:
+        super().bind_connector_metadata(connector_metadata)
+        if self.worker_handler is not None:
+            assert isinstance(connector_metadata, RecomputeCPUOffloadMetadata)
+            self.worker_handler.bind_connector_metadata(connector_metadata)
+
+    def clear_connector_metadata(self) -> None:
+        super().clear_connector_metadata()
+        if self.worker_handler is not None:
+            self.worker_handler.clear_connector_metadata()
+
+    def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata) -> None:
+        if self.worker_handler is not None:
+            assert isinstance(kv_connector_metadata, RecomputeCPUOffloadMetadata)
+            self.worker_handler.handle_preemptions(kv_connector_metadata)
+
+    def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:
+        if self.worker_handler is not None:
+            self.worker_handler.start_load_kv()
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        if self.worker_handler is not None:
+            self.worker_handler.wait_for_layer_load()
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        kv_layer: torch.Tensor,
+        attn_metadata: "AttentionMetadata",
+        **kwargs: Any,
+    ) -> None:
+        pass
+
+    def wait_for_save(self) -> None:
+        pass
+
+    def get_finished(
+        self,
+        finished_req_ids: set[str],
+    ) -> tuple[set[str] | None, set[str] | None]:
+        if self.worker_handler is not None:
+            return self.worker_handler.get_finished(finished_req_ids)
+        return None, None
+
+    def build_connector_worker_meta(self):
+        if self.worker_handler is not None:
+            return self.worker_handler.build_connector_worker_meta()
+        return None
+
+    # --- Scheduler-side methods ---
+
+    # NOTE: New API only for RecomputeCPUOffloadConnector.
+    def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None:
+        if self.scheduler_manager is not None:
+            self.scheduler_manager.bind_gpu_block_pool(gpu_block_pool)
+
+    def get_num_new_matched_tokens(
+        self,
+        request: "Request",
+        num_computed_tokens: int,
+    ) -> tuple[int | None, bool]:
+        if self.scheduler_manager is not None:
+            return self.scheduler_manager.get_num_new_matched_tokens(request, num_computed_tokens)
+        return 0, False
+
+    def update_state_after_alloc(
+        self,
+        request: "Request",
+        blocks: "KVCacheBlocks",
+        num_external_tokens: int,
+    ) -> None:
+        if self.scheduler_manager is not None:
+            self.scheduler_manager.update_state_after_alloc(request, blocks, num_external_tokens)
+
+    def update_state_before_preempt(
+        self,
+        request: "Request",
+        block_ids: tuple[list[int], ...],
+        num_computed_tokens: int,
+    ) -> bool:
+        if self.scheduler_manager is not None:
+            return self.scheduler_manager.update_state_before_preempt(
+                request,
+                block_ids,
+                num_computed_tokens,
+            )
+        return False
+
+    def build_connector_meta(
+        self,
+        scheduler_output: SchedulerOutput,
+    ) -> KVConnectorMetadata:
+        if self.scheduler_manager is not None:
+            return self.scheduler_manager.build_connector_meta(scheduler_output)
+        return RecomputeCPUOffloadMetadata()
+
+    def update_connector_output(
+        self,
+        connector_output: KVConnectorOutput,
+    ) -> None:
+        if self.scheduler_manager is not None:
+            self.scheduler_manager.update_connector_output(connector_output)
+
+    def request_finished(
+        self,
+        request: "Request",
+        block_ids: list[int],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        if self.scheduler_manager is not None:
+            return self.scheduler_manager.request_finished(request, block_ids)
+        return False, None
+
+    def request_finished_all_groups(
+        self,
+        request: "Request",
+        block_ids: tuple[list[int], ...],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        if self.scheduler_manager is not None:
+            return self.scheduler_manager.request_finished_all_groups(request, block_ids)
+        return False, None
+
+    # NOTE: New API only for RecomputeCPUOffloadConnector.
+    def has_pending_transfers(self) -> bool:
+        if self.scheduler_manager is not None:
+            return self.scheduler_manager.has_pending_transfers()
+        return False
+
+    def has_preempted_request(self, req_id: str) -> bool:
+        if self.scheduler_manager is not None:
+            return self.scheduler_manager.has_preempted_request(req_id)
+        return False
+
+    def take_events(self) -> Iterable[KVCacheEvent]:
+        if self.scheduler_manager is not None:
+            return self.scheduler_manager.take_events()
+        return []
+
+    def reset_cache(self) -> bool | None:
+        if self.scheduler_manager is not None:
+            return self.scheduler_manager.reset_cache()
+        return None

vllm_ascend/distributed/kv_transfer/kv_pool/recompute_cpu_offload/worker.py

@@ -0,0 +1,319 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Worker-side handler for Ascend RecomputeCPUOffloadConnector."""
+
+from typing import TYPE_CHECKING
+
+import torch
+from vllm.config import VllmConfig
+from vllm.logger import logger
+
+from vllm_ascend.distributed.kv_transfer.kv_pool.recompute_cpu_offload.metadata import (
+    RecomputeCPUOffloadMetadata,
+    RecomputeCPUOffloadWorkerMetadata,
+)
+
+if TYPE_CHECKING:
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+
+
+class RecomputeCPUOffloadWorker:
+    """Worker-side handler for recompute CPU/NPU KV cache transfers."""
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        kv_cache_config: "KVCacheConfig | None",
+        cpu_capacity_bytes: int,
+    ):
+        self.vllm_config = vllm_config
+        self.kv_cache_config = kv_cache_config
+        self.cpu_capacity_bytes = cpu_capacity_bytes
+
+        self.gpu_kv_caches: dict[str, torch.Tensor] | None = None
+        self.cpu_kv_caches: dict[str, torch.Tensor] | None = None
+        self.device: torch.device | None = None
+        self.num_cpu_blocks: int = 0
+
+        self.load_stream: torch.npu.Stream | None = None
+        self.store_stream: torch.npu.Stream | None = None
+
+        self._load_events: list[tuple[int, torch.npu.Event]] = []
+        self._load_hwm: int = -1
+
+        self._connector_metadata: RecomputeCPUOffloadMetadata | None = None
+        self._pending_load_event_indices: set[int] = set()
+        self._submitted_load_event_indices: set[int] = set()
+        self._completed_store_events: dict[int, int] = {}
+        self._load_stream_waited = False
+
+    def register_kv_caches(
+        self,
+        kv_caches: dict[str, torch.Tensor],
+    ) -> None:
+        """Register KV caches and initialize CPU/NPU transfer resources."""
+        if not kv_caches:
+            logger.warning("No KV caches to offload.")
+            return
+
+        any_tensor = next(iter(kv_caches.values()))
+        if isinstance(any_tensor, (tuple, list)):
+            any_tensor = any_tensor[0]
+        self.device = any_tensor.device
+
+        assert self.kv_cache_config is not None
+        self.num_gpu_blocks = self.kv_cache_config.num_blocks
+        self.block_size_scale = {}
+
+        scheduler_gpu_kv_cache_tensors = []
+        for t in self.kv_cache_config.kv_cache_tensors:
+            if t.shared_by:
+                scheduler_gpu_kv_cache_tensors.append(t)
+        scheduler_gpu_total_bytes = sum(t.size for t in scheduler_gpu_kv_cache_tensors)
+        scheduler_num_cpu_blocks = max(1, self.num_gpu_blocks * self.cpu_capacity_bytes // scheduler_gpu_total_bytes)
+
+        unique_gpu_caches: dict[str, torch.Tensor] = {}
+        register_cache_ptrs = []
+        for layer_name, layer_tensor in kv_caches.items():
+            if isinstance(layer_tensor, (tuple, list)):
+                for idx, single_tensor in enumerate(layer_tensor):
+                    if single_tensor.data_ptr() not in register_cache_ptrs:
+                        unique_gpu_caches[f"{layer_name}.{idx}"] = single_tensor.view(single_tensor.shape[0], -1)
+                        register_cache_ptrs.append(single_tensor.data_ptr())
+                        self.block_size_scale[f"{layer_name}.{idx}"] = single_tensor.shape[0] // self.num_gpu_blocks
+            else:
+                if layer_tensor.data_ptr() not in register_cache_ptrs:
+                    unique_gpu_caches[layer_name] = layer_tensor.view(layer_tensor.shape[0], -1)
+                    register_cache_ptrs.append(layer_tensor.data_ptr())
+                    self.block_size_scale[layer_name] = layer_tensor.shape[0] // self.num_gpu_blocks
+
+        per_tensor_bytes_per_block = [tensor.shape[-1] * tensor.element_size() for tensor in unique_gpu_caches.values()]
+        total_bytes_per_block = sum(per_tensor_bytes_per_block)
+        self.num_cpu_blocks = max(1, self.cpu_capacity_bytes // total_bytes_per_block)
+        if self.num_cpu_blocks != scheduler_num_cpu_blocks:
+            self.num_cpu_blocks = scheduler_num_cpu_blocks
+            logger.warning(
+                "RecomputeCPUOffloadScheduler has different num_blocks: %d,"
+                "worker-side num_block is set to %d to align with scheduler.",
+                scheduler_num_cpu_blocks,
+                scheduler_num_cpu_blocks,
+            )
+
+        self.gpu_kv_caches = unique_gpu_caches
+        self.cpu_kv_caches = {}
+        for name, gpu_tensor in unique_gpu_caches.items():
+            tensor_block_size_scale = self.block_size_scale[name]
+            cpu_shape = (self.num_cpu_blocks * tensor_block_size_scale,) + gpu_tensor.shape[1:]
+            self.cpu_kv_caches[name] = torch.zeros(
+                cpu_shape,
+                dtype=gpu_tensor.dtype,
+                pin_memory=True,
+                device="cpu",
+            )
+
+        self.load_stream = torch.npu.Stream()
+        self.store_stream = torch.npu.Stream()
+
+        logger.info(
+            "RecomputeCPUOffloadWorker scaffold registered %d unique KV tensors, allocating %d CPU blocks (%.2f GB).",
+            len(unique_gpu_caches),
+            self.num_cpu_blocks,
+            (self.num_cpu_blocks * total_bytes_per_block) / (1024**3),
+        )
+
+    def bind_connector_metadata(self, metadata: RecomputeCPUOffloadMetadata) -> None:
+        self._connector_metadata = metadata
+        self._load_stream_waited = False
+        if metadata.preempt_load_event >= 0:
+            self._pending_load_event_indices.add(metadata.preempt_load_event)
+
+    def clear_connector_metadata(self) -> None:
+        """Clear metadata after the model runner finishes the current step."""
+        self._connector_metadata = None
+
+    def handle_preemptions(
+        self,
+        kv_connector_metadata: RecomputeCPUOffloadMetadata,
+    ) -> None:
+        """Save preempted blocks before input preparation can overwrite them."""
+        if kv_connector_metadata.need_flush:
+            self._flush_and_sync_all()
+
+        # The scheduler may immediately reuse preempted block IDs in this same
+        # step. This blocking D2H must therefore run before _update_states()
+        # processes new_block_ids_to_zero and before model forward writes KV.
+        self._submit_transfer(
+            kv_connector_metadata.preempt_store_gpu_blocks,
+            kv_connector_metadata.preempt_store_cpu_blocks,
+            kv_connector_metadata.preempt_store_event,
+            is_store=True,
+            sync=True,
+        )
+
+    def start_load_kv(self) -> None:
+        """Submit pre-forward recompute H2D transfers."""
+        metadata = self._connector_metadata
+        if metadata is None:
+            return
+
+        self._submit_transfer(
+            metadata.preempt_load_cpu_blocks,
+            metadata.preempt_load_gpu_blocks,
+            metadata.preempt_load_event,
+            is_store=False,
+            sync=True,
+        )
+
+    def wait_for_layer_load(self) -> None:
+        """Make the current forward stream wait for the recompute H2D copy."""
+        if self._load_stream_waited or self.load_stream is None:
+            return
+        metadata = self._connector_metadata
+        if metadata is None or metadata.preempt_load_event < 0:
+            return
+        torch.npu.current_stream().wait_stream(self.load_stream)
+        self._load_stream_waited = True
+
+    def _flush_and_sync_all(self) -> None:
+        """Synchronize all in-flight transfer events."""
+        for event_idx, event in self._load_events:
+            event.synchronize()
+            self._load_hwm = event_idx
+        self._load_events.clear()
+        self._submitted_load_event_indices.clear()
+
+    def _poll_load_events(self) -> int:
+        """Return the highest completed H2D event index."""
+        events = self._load_events
+        hwm = self._load_hwm
+
+        while events:
+            event_idx, event = events[0]
+            if not event.query():
+                break
+            hwm = event_idx
+            events.pop(0)
+
+        self._load_hwm = hwm
+        return hwm
+
+    def _submit_transfer(
+        self,
+        src_block_ids: list[int],
+        dst_block_ids: list[int],
+        event_idx: int,
+        is_store: bool,
+        sync: bool = False,
+    ) -> None:
+        """Submit a CPU<->NPU block copy and record a completion event."""
+        if event_idx < 0:
+            return
+        if not is_store and event_idx in self._submitted_load_event_indices:
+            return
+        if not is_store:
+            self._submitted_load_event_indices.add(event_idx)
+
+        if not src_block_ids:
+            if is_store:
+                self._completed_store_events[event_idx] = 1
+            else:
+                self._load_hwm = max(self._load_hwm, event_idx)
+            return
+
+        assert len(src_block_ids) == len(dst_block_ids)
+        assert self.gpu_kv_caches is not None
+        assert self.cpu_kv_caches is not None
+
+        stream = self.store_stream if is_store else self.load_stream
+        assert stream is not None
+        torch.npu.synchronize()
+
+        with torch.npu.stream(stream):
+            for src_block_id, dst_block_id in zip(src_block_ids, dst_block_ids):
+                for name, gpu_tensor in self.gpu_kv_caches.items():
+                    cpu_tensor = self.cpu_kv_caches[name]
+                    tensor_block_size_scale = self.block_size_scale[name]
+                    if is_store:
+                        # TODO: Replace this D2H torch copy with the NPU copy
+                        # backend dedicated kernel.
+                        if tensor_block_size_scale > 1:
+                            cpu_tensor[
+                                dst_block_id * tensor_block_size_scale : (dst_block_id + 1) * tensor_block_size_scale
+                            ].copy_(
+                                gpu_tensor[
+                                    src_block_id * tensor_block_size_scale : (src_block_id + 1)
+                                    * tensor_block_size_scale
+                                ],
+                                non_blocking=True,
+                            )
+                        else:
+                            cpu_tensor[dst_block_id].copy_(
+                                gpu_tensor[src_block_id],
+                                non_blocking=True,
+                            )
+                    else:
+                        # TODO: Replace this H2D torch copy with the NPU copy
+                        # backend dedicated kernel.
+                        if tensor_block_size_scale > 1:
+                            gpu_tensor[
+                                dst_block_id * tensor_block_size_scale : (dst_block_id + 1) * tensor_block_size_scale
+                            ].copy_(
+                                cpu_tensor[
+                                    src_block_id * tensor_block_size_scale : (src_block_id + 1)
+                                    * tensor_block_size_scale
+                                ],
+                                non_blocking=True,
+                            )
+                        else:
+                            gpu_tensor[dst_block_id].copy_(
+                                cpu_tensor[src_block_id],
+                                non_blocking=True,
+                            )
+            event = torch.npu.Event()
+            event.record(stream)
+
+        if sync:
+            event.synchronize()
+            if is_store:
+                self._completed_store_events[event_idx] = 1
+            else:
+                self._load_hwm = max(self._load_hwm, event_idx)
+            return
+
+        assert not is_store
+        self._load_events.append((event_idx, event))
+
+    def get_finished(
+        self,
+        finished_req_ids: set[str],
+    ) -> tuple[set[str] | None, set[str] | None]:
+        """Poll recompute transfers and report completed request restores."""
+        metadata = self._connector_metadata
+        if metadata is None:
+            return None, None
+
+        finished_recving: set[str] = set()
+        if self._pending_load_event_indices:
+            load_hwm = self._poll_load_events()
+            completed_loads = [event_idx for event_idx in self._pending_load_event_indices if event_idx <= load_hwm]
+            for event_idx in completed_loads:
+                self._pending_load_event_indices.discard(event_idx)
+                self._submitted_load_event_indices.discard(event_idx)
+                finished_recving.update(metadata.preempt_load_event_to_reqs.get(event_idx, []))
+
+        return None, finished_recving or None
+
+    def build_connector_worker_meta(self) -> RecomputeCPUOffloadWorkerMetadata | None:
+        """Return completed store events since the previous call.
+
+        The scheduler aggregates this metadata across workers/ranks. A store
+        event becomes available to recompute requests only after all expected
+        workers have reported completion.
+        """
+        if not self._completed_store_events:
+            return None
+        meta = RecomputeCPUOffloadWorkerMetadata(
+            completed_store_events=self._completed_store_events,
+        )
+        self._completed_store_events = {}
+        return meta

vllm_ascend/distributed/kv_transfer/kv_pool/simple_cpu_offload/__init__.py

@@ -0,0 +1 @@
+

vllm_ascend/distributed/kv_transfer/kv_pool/simple_cpu_offload/simple_cpu_offload_connector.py

@@ -0,0 +1,60 @@
+"""Ascend NPU adaptation of vLLM's ``SimpleCPUOffloadConnector``.
+
+The scheduler-side ``SimpleCPUOffloadScheduler`` is platform-agnostic
+and reused as-is from upstream vLLM. The Ascend variant only swaps the
+worker-side handler with an NPU-native implementation that uses
+``aclrtMemcpyBatchAsync`` and ``torch.npu`` streams/events.
+"""
+
+from typing import TYPE_CHECKING
+
+from vllm.config import VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector import (  # noqa: E501
+    SimpleCPUOffloadConnector,
+)
+from vllm.logger import logger
+
+from vllm_ascend.simple_kv_offload.worker import SimpleCPUOffloadNPUWorker
+
+if TYPE_CHECKING:
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+
+
+class AscendSimpleCPUOffloadConnector(SimpleCPUOffloadConnector):
+    """NPU-flavored ``SimpleCPUOffloadConnector``.
+
+    Inherits the full scheduler/worker plumbing from upstream and only
+    replaces the CUDA worker handler with the NPU one. All other public
+    APIs (``register_kv_caches``, ``bind_connector_metadata``,
+    ``get_finished``, ``handle_preemptions``, every scheduler-side
+    method, etc.) are inherited verbatim — they all route through
+    ``self.worker_handler`` / ``self.scheduler_manager``.
+
+    Why post-init swap (instead of skipping ``super().__init__``):
+    ``SimpleCPUOffloadWorker.__init__`` and ``DmaCopyBackend.__init__``
+    only assign ``None``/empty-field defaults — no CUDA resource is
+    allocated until ``register_kv_caches`` runs. So letting the parent
+    construct a transient CUDA worker and then replacing it costs
+    nothing and keeps us free of duplicated configuration parsing.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        role: KVConnectorRole,
+        kv_cache_config: "KVCacheConfig | None" = None,
+    ) -> None:
+        super().__init__(vllm_config, role, kv_cache_config)
+
+        # If prefix caching is disabled, the parent leaves both handlers
+        # as None and the connector is a no-op — nothing to swap.
+        if role == KVConnectorRole.WORKER and self.worker_handler is not None:
+            cpu_capacity = self.worker_handler.cpu_capacity_bytes
+            self.worker_handler: SimpleCPUOffloadNPUWorker = SimpleCPUOffloadNPUWorker(
+                vllm_config, kv_cache_config, cpu_capacity
+            )
+            logger.info(
+                "AscendSimpleCPUOffloadConnector: swapped CUDA worker for NPU worker (per_rank=%.2f GB)",
+                cpu_capacity / (1024**3),
+            )

vllm_ascend/distributed/kv_transfer/kv_pool/ucm_connector.py

@@ -0,0 +1,324 @@
+# SPDX-License-Identifier: Apache-2.0
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Any, Optional
+
+import torch
+from ucm.integration.vllm.ucm_connector import UCMConnector
+from vllm.config import VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    CopyBlocksOp,
+    KVConnectorBase_V1,
+    KVConnectorHandshakeMetadata,
+    KVConnectorMetadata,
+    KVConnectorRole,
+    KVConnectorWorkerMetadata,
+    SupportsHMA,
+)
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import KVConnectorOutput
+
+# isort: off
+if TYPE_CHECKING:
+    from vllm.distributed.kv_events import KVCacheEvent, KVConnectorKVEvents
+    from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
+        KVConnectorPromMetrics,
+        KVConnectorStats,
+        PromMetric,
+        PromMetricT,
+    )
+    from vllm.forward_context import ForwardContext
+    from vllm.v1.attention.backend import AttentionMetadata
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.request import Request
+# isort: on
+
+
+class UCMConnectorV1(KVConnectorBase_V1, SupportsHMA):
+    def __init__(
+        self,
+        vllm_config: "VllmConfig",
+        role: KVConnectorRole,
+        kv_cache_config: "KVCacheConfig",
+    ):
+        super().__init__(vllm_config=vllm_config, role=role, kv_cache_config=kv_cache_config)
+        assert vllm_config.kv_transfer_config is not None
+
+        ImplCls = UCMConnector
+        self._ucm_engine = ImplCls(vllm_config, role, kv_cache_config)
+
+    def _get_ucm_delegate_for(self, name: str) -> Any:
+        """Return the UCM object that owns a reserved connector interface.
+
+        The imported UCMConnector is itself a thin dispatcher. Some reserved
+        KVConnectorBase_V1 hooks are not redeclared on that dispatcher yet, so
+        the base class no-op can otherwise hide the real inner connector hook.
+        Prefer an explicit method/property on the dispatcher; otherwise fall
+        through to the selected inner connector when present.
+        """
+        if name not in type(self._ucm_engine).__dict__:
+            inner_connector = getattr(self._ucm_engine, "connector", None)
+            if inner_connector is not None:
+                return inner_connector
+        return self._ucm_engine
+
+    def _call_ucm_reserved_hook(self, name: str, *args: Any, **kwargs: Any) -> Any:
+        hook = getattr(self._get_ucm_delegate_for(name), name, None)
+        if callable(hook):
+            return hook(*args, **kwargs)
+        return None
+
+    # ==============================
+    # Worker-side methods
+    # ==============================
+    def shutdown(self) -> None:
+        self._call_ucm_reserved_hook("shutdown")
+
+    def has_connector_metadata(self) -> bool:
+        """Check whether the connector metadata is currently set.
+
+        Returns:
+            bool: True if connector metadata exists, False otherwise.
+        """
+        return self._ucm_engine.has_connector_metadata()
+
+    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 set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp) -> None:
+        self._call_ucm_reserved_hook("set_host_xfer_buffer_ops", copy_operation)
+
+    def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata) -> None:
+        self._call_ucm_reserved_hook("handle_preemptions", kv_connector_metadata)
+
+    def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:
+        """
+        Start loading the KV cache from the connector to vLLM's paged
+        KV buffer. This is called from the forward context before the
+        forward pass to enable async loading during model execution.
+
+        Args:
+            forward_context (ForwardContext): the forward context.
+            **kwargs: additional arguments for the load operation
+
+        Note:
+            The number of elements in kv_caches and layer_names should be
+            the same.
+
+        """
+        self._ucm_engine.start_load_kv(forward_context, **kwargs)
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        """
+        Block until the KV for a specific layer is loaded into vLLM's
+        paged buffer. This is called from within attention layer to ensure
+        async copying from start_load_kv is complete.
+
+        This interface will be useful for layer-by-layer pipelining.
+
+        Args:
+            layer_name: the name of that layer
+        """
+        self._ucm_engine.wait_for_layer_load(layer_name)
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        kv_layer: torch.Tensor,
+        attn_metadata: "AttentionMetadata",
+        **kwargs: Any,
+    ) -> None:
+        """
+        Start saving the a layer of KV cache from vLLM's paged buffer
+        to the connector. This is called from within attention layer to
+        enable async copying during execution.
+
+        Args:
+            layer_name (str): the name of the layer.
+            kv_layer (torch.Tensor): the paged KV buffer of the current
+                layer in vLLM.
+            attn_metadata (AttentionMetadata): the attention metadata.
+            **kwargs: additional arguments for the save operation.
+        """
+        self._ucm_engine.save_kv_layer(layer_name, kv_layer, attn_metadata, **kwargs)
+
+    def wait_for_save(self) -> None:
+        """
+        Block until all the save operations is done. This is called
+        as the forward context exits to ensure that the async saving
+        from save_kv_layer is complete before finishing the forward.
+
+        This prevents overwrites of paged KV buffer before saving done.
+        """
+        self._ucm_engine.wait_for_save()
+
+    def clear_connector_metadata(self) -> None:
+        """Clear the connector metadata.
+
+        This function should be called by the model runner every time
+        after the model execution.
+        """
+        self._ucm_engine.clear_connector_metadata()
+
+    def bind_connector_metadata(self, connector_metadata: KVConnectorMetadata) -> None:
+        """Set the connector metadata from the scheduler.
+
+        This function should be called by the model runner every time
+        before the model execution. The metadata will be used for runtime
+        KV cache loading and saving.
+
+        Args:
+            connector_metadata (dict): the connector metadata.
+        """
+        self._ucm_engine.bind_connector_metadata(connector_metadata)
+
+    def get_block_ids_with_load_errors(self) -> set[int]:
+        """
+        Get the set of block IDs that failed to load.
+
+        Returns:
+            Set of block IDs that encountered load errors.
+            Empty set if no load errors occurred.
+        """
+        return self._ucm_engine.get_block_ids_with_load_errors()
+
+    # ==============================
+    # Scheduler-side methods
+    # ==============================
+    def get_num_new_matched_tokens(
+        self,
+        request: "Request",
+        num_computed_tokens: int,
+    ) -> tuple[int | None, bool]:
+        """
+        Get number of new tokens that can be loaded from the
+        external KV cache beyond the num_computed_tokens.
+
+        Args:
+            request (Request): the request object.
+            num_computed_tokens (int): the number of locally
+                computed tokens for this request
+
+        Returns:
+            the number of tokens that can be loaded from the
+            external KV cache beyond what is already computed.
+        """
+        return self._ucm_engine.get_num_new_matched_tokens(request, num_computed_tokens)
+
+    def update_state_after_alloc(self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int) -> None:
+        """
+        Update KVConnector state after block allocation.
+        """
+        self._ucm_engine.update_state_after_alloc(request, blocks, num_external_tokens)
+
+    def build_connector_meta(self, scheduler_output: SchedulerOutput) -> KVConnectorMetadata:
+        """
+        Build the connector metadata for this step.
+
+        This function should NOT modify fields in the scheduler_output.
+        Also, calling this function will reset the state of the connector.
+
+        Args:
+            scheduler_output (SchedulerOutput): the scheduler output object.
+        """
+        return self._ucm_engine.build_connector_meta(scheduler_output)
+
+    def request_finished(
+        self,
+        request: "Request",
+        block_ids: list[int],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        """
+        Called when a request has finished, before its blocks are freed.
+
+        Returns:
+            True if the request is being saved/sent asynchronously and blocks
+            should not be freed until the request_id is returned from
+            get_finished().
+            Optional KVTransferParams to be included in the request outputs
+            returned by the engine.
+        """
+        return self._ucm_engine.request_finished(request, block_ids)
+
+    def request_finished_all_groups(
+        self,
+        request: "Request",
+        block_ids: tuple[list[int], ...],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        return self._ucm_engine.request_finished_all_groups(request, block_ids)
+
+    def get_finished(
+        self,
+        finished_req_ids: set[str],
+    ) -> tuple[set[str] | None, set[str] | None]:
+        return self._ucm_engine.get_finished(finished_req_ids)
+
+    def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None:
+        return self._call_ucm_reserved_hook("build_connector_worker_meta")
+
+    def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
+        self._ucm_engine.update_connector_output(connector_output)
+
+    def take_events(self) -> Iterable["KVCacheEvent"]:
+        events = self._call_ucm_reserved_hook("take_events")
+        return () if events is None else events
+
+    def get_kv_connector_stats(self) -> Optional["KVConnectorStats"]:
+        return self._call_ucm_reserved_hook("get_kv_connector_stats")
+
+    def get_kv_connector_kv_cache_events(self) -> Optional["KVConnectorKVEvents"]:
+        return self._call_ucm_reserved_hook("get_kv_connector_kv_cache_events")
+
+    def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None:
+        return self._call_ucm_reserved_hook("get_handshake_metadata")
+
+    def set_xfer_handshake_metadata(self, metadata: dict[int, KVConnectorHandshakeMetadata]) -> None:
+        self._call_ucm_reserved_hook("set_xfer_handshake_metadata", metadata)
+
+    def get_finished_count(self) -> int | None:
+        return self._call_ucm_reserved_hook("get_finished_count")
+
+    def reset_cache(self) -> bool | None:
+        return self._call_ucm_reserved_hook("reset_cache")
+
+    # ==============================
+    # Metrics & Stats
+    # ==============================
+
+    @classmethod
+    def build_kv_connector_stats(cls, data: dict[str, Any] | None = None) -> Optional["KVConnectorStats"]:
+        """
+        KVConnectorStats resolution method. This method allows dynamically
+        registered connectors to return their own KVConnectorStats object,
+        which can implement custom aggregation logic on the data dict.
+        """
+        return UCMConnector.build_kv_connector_stats(data)
+
+    @classmethod
+    def build_prom_metrics(
+        cls,
+        vllm_config: "VllmConfig",
+        metric_types: dict[type["PromMetric"], type["PromMetricT"]],
+        labelnames: list[str],
+        per_engine_labelvalues: dict[int, list[object]],
+    ) -> Optional["KVConnectorPromMetrics"]:
+        """
+        Create a KVConnectorPromMetrics subclass which should register
+        per-connector Prometheus metrics and implement observe() to
+        expose connector transfer stats via Prometheus.
+
+        This implementation forwards the call to the underlying
+        UCMConnector engine.
+        """
+        return UCMConnector.build_prom_metrics(
+            vllm_config,
+            metric_types,
+            labelnames,
+            per_engine_labelvalues,
+        )

vllm_ascend/distributed/kv_transfer/utils/mooncake_transfer_engine.py

@@ -0,0 +1,43 @@
+import threading
+
+
+class GlobalTE:
+    def __init__(self):
+        self.transfer_engine = None
+        self.is_register_buffer: bool = False
+        self.transfer_engine_lock = threading.Lock()
+        self.register_buffer_lock = threading.Lock()
+
+    def get_transfer_engine(self, hostname: str, device_name: str | None):
+        if self.transfer_engine is None:
+            with self.transfer_engine_lock:
+                # Double-Checked Locking
+                if self.transfer_engine is None:
+                    try:
+                        from mooncake.engine import TransferEngine  # type: ignore
+                    except ImportError as e:
+                        raise ImportError(
+                            "Please install mooncake by following the instructions at "
+                            "https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md "  # noqa: E501
+                            "to run vLLM with MooncakeConnector."
+                        ) from e
+                    self.transfer_engine = TransferEngine()
+                    device_name = device_name if device_name is not None else ""
+                    ret_value = self.transfer_engine.initialize(hostname, "P2PHANDSHAKE", "ascend", device_name)
+                    if ret_value != 0:
+                        raise RuntimeError(f"TransferEngine initialization failed with ret_value: {ret_value}")
+        return self.transfer_engine
+
+    def register_buffer(self, ptrs: list[int], sizes: list[int]):
+        with self.register_buffer_lock:
+            assert self.transfer_engine is not None, "Transfer engine must be initialized"
+            if self.is_register_buffer:
+                return
+            for ptr, size in zip(ptrs, sizes):
+                ret_value = self.transfer_engine.register_memory(ptr, size)
+                if ret_value != 0:
+                    raise RuntimeError("Mooncake memory registration failed.")
+            self.is_register_buffer = True
+
+
+global_te = GlobalTE()

vllm_ascend/distributed/kv_transfer/utils/utils.py

@@ -0,0 +1,445 @@
+import math
+import os
+from collections import OrderedDict, defaultdict
+from collections.abc import Iterator
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.distributed as dist
+from vllm.logger import logger
+
+from vllm_ascend.distributed.parallel_state import get_p_tp_group
+
+MAX_HCCL_REGISTER_REGIONS = 256
+REGISTER_MERGE_GAP_BYTES = 4096
+
+
+def kv_alltoall_and_rearrange(pd_tp_ratio: int, key: torch.Tensor, value: torch.TensorType):
+    if pd_tp_ratio <= 1:
+        return None, None
+    elif key is None or value is None:
+        raise ValueError("key or value is None")
+    k_output = alltoall_and_rearrange(pd_tp_ratio, key)
+    v_output = alltoall_and_rearrange(pd_tp_ratio, value)
+    return k_output, v_output
+
+
+def alltoall_and_rearrange(tp_ratio: int, input_tensor: torch.Tensor):
+    num_kv_heads = input_tensor.size(1)
+    output_tensor = torch.zeros_like(input_tensor)
+    dist.all_to_all_single(output_tensor, input_tensor, group=get_p_tp_group().device_group)
+    input_tensor = 0
+    result = rearrange_output(output_tensor, tp_ratio, num_kv_heads)
+    output_tensor = 0
+    return result
+
+
+def rearrange_output(base_output: torch.Tensor, cut_num: int, num_kv_heads: int):
+    size_0 = base_output.size(0)
+    if size_0 % cut_num != 0:
+        raise ValueError(f"The size of dim 0 [{size_0}] must be divisible by the cut_num [{cut_num}]")
+    chunk_size = size_0 // cut_num
+    reshaped = base_output.view(cut_num, chunk_size, -1)
+    transposed = reshaped.transpose(0, 1)
+    return transposed.contiguous().view(size_0, num_kv_heads, -1)
+
+
+def align_memory(tensor: torch.Tensor, alignment: int) -> torch.Tensor:
+    data_ptr = tensor.data_ptr()
+    aligned_addr = (data_ptr + alignment - 1) // alignment * alignment
+    offset = (aligned_addr - data_ptr) // tensor.element_size()
+    return tensor[int(offset) :]
+
+
+def get_transfer_timeout_value():
+    ascend_transfer_timeout = os.getenv("ASCEND_TRANSFER_TIMEOUT", "")
+    if len(ascend_transfer_timeout) > 0:
+        return int(ascend_transfer_timeout)
+    hccl_rdma_timeout = int(os.getenv("HCCL_RDMA_TIMEOUT", "20"))  # type: ignore
+    hccl_rdma_retry_cnt = int(os.getenv("HCCL_RDMA_RETRY_CNT", "7"))  # type: ignore
+    return int((4.096 * (2**hccl_rdma_timeout)) * hccl_rdma_retry_cnt // 1000 + 3000)
+
+
+@dataclass
+class parallel_info:
+    tp_size: int
+    pcp_size: int
+    dcp_size: int
+    use_mla: bool
+    pd_head_ratio: int
+
+
+def get_cp_group(tp: int, heads: int, dcp: int):
+    # Partition the second dimension of [pcp][head_group][dcp] to obtain a complete head group
+    # head_group is all blocks for request in the same head
+    # tp8 dcp2 heads4 return[[0,1,2,3]]
+    # tp8 dcp1 heads4 return[[0,2,4,6],[1,3,5,7]]
+    step = tp // heads
+    if step == 0:
+        return [[i for i in range(tp // dcp)]]
+    else:
+        return [
+            set([k // dcp for h in range(heads) for k in range(h * step + i * dcp, h * step + (i + 1) * dcp)])
+            for i in range(step // dcp)
+        ]
+
+
+def context_parallel_parameters_check(
+    remote_pcp_size: int,
+    remote_dcp_size: int,
+    p_parallel_info: parallel_info,
+    d_parallel_info: parallel_info,
+    total_num_kv_heads: int,
+):
+    # Check whether the pcp–dcp ratio is supported
+    assert (p_parallel_info.pcp_size * p_parallel_info.dcp_size) % (remote_pcp_size * remote_dcp_size) == 0
+    if not p_parallel_info.use_mla:
+        p_node_heads_per_rank = math.ceil(total_num_kv_heads / p_parallel_info.tp_size)
+        d_node_heads_per_rank = math.ceil(total_num_kv_heads / d_parallel_info.dcp_size)
+        assert d_node_heads_per_rank % p_node_heads_per_rank == 0
+
+
+def get_tp_rank_head_mapping(num_key_value_heads: int, tp_size: int):
+    # Get the head_idx corresponding to the tp_rank, {tp_rank:[head_indx]}
+    mapping = {}
+    if tp_size <= num_key_value_heads:
+        if num_key_value_heads % tp_size != 0:
+            raise ValueError(f"Number of heads ({num_key_value_heads}) cannot be evenly divided by TP ({tp_size}).")
+
+        heads_per_rank = num_key_value_heads // tp_size
+
+        for rank in range(tp_size):
+            start_idx = rank * heads_per_rank
+            end_idx = start_idx + heads_per_rank
+            mapping[rank] = list(range(start_idx, end_idx))
+    else:
+        if tp_size % num_key_value_heads != 0:
+            raise ValueError(f"Number of heads ({num_key_value_heads}) cannot be evenly divided by TP ({tp_size}).")
+        ranks_per_head = tp_size // num_key_value_heads
+        for rank in range(tp_size):
+            head_idx = rank // ranks_per_head
+            mapping[rank] = [head_idx]
+    return mapping
+
+
+def get_head_group_mapping(num_key_value_heads: int, tp_size: int, num_groups: int, select_cp_group: list[int]):
+    # Get the mapping dictionary, where the key is head_group_rank and the value is head_idx
+    if tp_size % num_groups != 0:
+        raise ValueError(
+            f"Total number of devices ({tp_size}) cannot be divided by the number of groups ({num_groups})."
+        )
+    ranks_per_group = tp_size // num_groups
+    tp_mapping = get_tp_rank_head_mapping(num_key_value_heads, tp_size)
+    group_mapping = {}
+    for group_rank in range(num_groups):
+        if group_rank in select_cp_group:
+            start_rank = group_rank * ranks_per_group
+            end_rank = start_rank + ranks_per_group
+            heads_set = set()
+
+            for rank in range(start_rank, end_rank):
+                heads_set.update(tp_mapping[rank])
+            group_mapping[group_rank] = sorted(list(heads_set))
+    return group_mapping
+
+
+def get_local_remote_block_port_mappings(
+    to_trans_idx: int,
+    p_parallel_info: parallel_info,
+    d_parallel_info: parallel_info,
+    d_hosts: list[str],
+    d_port: int,
+    selected_p_cp_group: list[int],
+    selected_d_cp_group: list[int],
+    prompt_len: int,
+    block_size: int,
+    req_meta,
+    total_num_kv_heads: int,
+    req_id: str,
+):
+    p_head_group_size = p_parallel_info.tp_size // p_parallel_info.dcp_size
+    d_head_group_size = d_parallel_info.tp_size // d_parallel_info.dcp_size
+    world_size = d_parallel_info.pcp_size * d_head_group_size * d_parallel_info.dcp_size
+    # Compute which logic_block_idx corresponds to each tp_rank
+    p_rank_block_mapping: list[list[list[list[int]]]] = [
+        [[[] for _ in range(p_parallel_info.dcp_size)] for _ in range(p_head_group_size)]
+        for _ in range(p_parallel_info.pcp_size)
+    ]
+    for logic_block_idx in range(to_trans_idx):
+        pcp_rank = (logic_block_idx // p_parallel_info.dcp_size) % p_parallel_info.pcp_size
+        dcp_rank = logic_block_idx % p_parallel_info.dcp_size
+        for p_head_group_rank in range(p_head_group_size):
+            if p_head_group_rank in selected_p_cp_group:
+                p_rank_block_mapping[pcp_rank][p_head_group_rank][dcp_rank].append(logic_block_idx)
+
+    # Find the remote device that holds the logic_block_idx
+    d_block_rank_mapping: dict[int, dict[int, dict[str, Any]]] = defaultdict(lambda: defaultdict(dict))
+    for logic_block_idx in range(to_trans_idx):
+        pcp_rank = (logic_block_idx // d_parallel_info.dcp_size) % d_parallel_info.pcp_size
+        for d_head_group_rank in range(d_head_group_size):
+            if d_head_group_rank in selected_d_cp_group:
+                dcp_rank = logic_block_idx % d_parallel_info.dcp_size
+                world_rank = (
+                    pcp_rank * d_head_group_size * d_parallel_info.dcp_size
+                    + d_head_group_rank * d_parallel_info.dcp_size
+                    + dcp_rank
+                )
+                world_size = d_parallel_info.pcp_size * d_head_group_size * d_parallel_info.dcp_size
+                host = d_hosts[(len(d_hosts) * world_rank) // world_size]
+                port = d_port + world_rank
+                block_idx = (logic_block_idx - (pcp_rank * d_parallel_info.pcp_size + dcp_rank)) // (
+                    d_parallel_info.pcp_size * d_parallel_info.dcp_size
+                )
+                d_block_rank_mapping[logic_block_idx][d_head_group_rank] = {
+                    "pcp_rank": pcp_rank,
+                    "dcp_rank": dcp_rank,
+                    "host": host,
+                    "port": port,
+                    "block_idx": block_idx,
+                }
+    # Get how many times each device should receive done_single for this request
+    d_trans_count_mapping = {}
+    trans_block_size = math.ceil(prompt_len / block_size)  # Total number of blocks
+    transed_block_size = math.ceil(req_meta.remote_cache_tokens / block_size)  # Number of prefix cache hit blocks
+    d_cp_size = d_parallel_info.pcp_size * d_parallel_info.dcp_size
+    for d_pcp_rank in range(d_parallel_info.pcp_size):
+        for d_head_group_rank in range(d_head_group_size):
+            for d_dcp_rank in range(d_parallel_info.dcp_size):
+                if trans_block_size >= (p_parallel_info.pcp_size * p_parallel_info.dcp_size):
+                    trans_count = (p_parallel_info.pcp_size * p_parallel_info.dcp_size) // d_cp_size
+                else:
+                    current_rank_idx = d_pcp_rank * d_parallel_info.dcp_size + d_dcp_rank
+                    total_global_blocks = transed_block_size + trans_block_size
+
+                    target_total_count = total_global_blocks // d_cp_size
+                    if current_rank_idx < (total_global_blocks % d_cp_size):
+                        target_total_count += 1
+
+                    prev_processed_count = transed_block_size // d_cp_size
+                    if current_rank_idx < (transed_block_size % d_cp_size):
+                        prev_processed_count += 1
+
+                    trans_count = target_total_count - prev_processed_count
+                world_rank = (
+                    d_pcp_rank * d_head_group_size * d_parallel_info.dcp_size
+                    + d_head_group_rank * d_parallel_info.dcp_size
+                    + d_dcp_rank
+                )
+                host = d_hosts[(len(d_hosts) * world_rank) // world_size]
+                port = d_port + world_rank
+                d_trans_count_mapping[(host, port)] = trans_count * p_parallel_info.pd_head_ratio
+
+    # Compute the mapping between local and remote head_group_rank
+    p_tp_rank_head_mapping = get_head_group_mapping(
+        total_num_kv_heads, p_parallel_info.tp_size, p_head_group_size, selected_p_cp_group
+    )
+    d_tp_rank_head_mapping = get_head_group_mapping(
+        total_num_kv_heads, d_parallel_info.tp_size, d_head_group_size, selected_d_cp_group
+    )
+    head_to_d_groups = defaultdict(set)
+    for d_rank, heads in d_tp_rank_head_mapping.items():
+        for head in heads:
+            head_to_d_groups[head].add(d_rank)
+    pd_head_mapping = {}
+    for p_rank, p_heads in p_tp_rank_head_mapping.items():
+        target_d_ranks = set()
+        for head in p_heads:
+            if head in head_to_d_groups:
+                target_d_ranks.update(head_to_d_groups[head])
+            else:
+                logger.info("Warning: Head %s exists in P but not in D mapping.", head)
+        pd_head_mapping[p_rank] = sorted(list(target_d_ranks))
+    logger.debug(
+        "MooncakeLayerwiseConnector _get_kv_split_metadata req_id=%r "
+        "P-side logic_block to rank mapping: %s, "
+        "D-side logic_block to rank mapping: %s, "
+        "P&D head_group_rank mapping: %s",
+        req_id,
+        p_rank_block_mapping,
+        d_block_rank_mapping,
+        pd_head_mapping,
+    )
+    return p_rank_block_mapping, d_block_rank_mapping, pd_head_mapping, d_trans_count_mapping
+
+
+def get_transfer_mappings(
+    p_rank_block_mapping: list[list[list[list[int]]]],
+    d_block_rank_mapping: dict[int, dict[int, dict[str, Any]]],
+    pd_head_mapping: dict[int, set],
+    d_trans_count_mapping: dict[tuple[str, int], int],
+    req_meta,
+    block_group_idx: int,
+    p_parallel_info: parallel_info,
+    req_id: str,
+    transed_idx: int,
+    to_trans_idx: int,
+    tp_rank: int,
+    pcp_rank: int,
+    dcp_rank: int,
+):
+    transfer_mappings: dict[tuple[str, int], dict[str, Any]] = {}
+    p_head_group_rank = (tp_rank - dcp_rank) // p_parallel_info.dcp_size
+    p_block_idxs: list[int] = p_rank_block_mapping[pcp_rank][p_head_group_rank][dcp_rank]
+    p_block_ids = req_meta.local_block_ids[block_group_idx]
+    d_block_ids = req_meta.remote_block_ids[block_group_idx]
+    for p_block_idx, logic_block_idx in enumerate(p_block_idxs):
+        if logic_block_idx < transed_idx or logic_block_idx >= to_trans_idx:
+            continue
+        for d_head_group_rank in pd_head_mapping[p_head_group_rank]:
+            p_block_id = p_block_ids[p_block_idx]
+            remote_host = d_block_rank_mapping[logic_block_idx][d_head_group_rank]["host"]
+            remote_port = d_block_rank_mapping[logic_block_idx][d_head_group_rank]["port"]
+            d_block_idx = d_block_rank_mapping[logic_block_idx][d_head_group_rank]["block_idx"]
+            d_block_id = d_block_ids[d_block_idx]
+            if (remote_host, remote_port) not in transfer_mappings:
+                transfer_mappings[(remote_host, remote_port)] = {
+                    "local_block_ids": [],
+                    "remote_block_ids": [],
+                    "trans_count": 0,
+                }
+            transfer_mappings[(remote_host, remote_port)]["local_block_ids"].append(p_block_id)
+            transfer_mappings[(remote_host, remote_port)]["remote_block_ids"].append(d_block_id)
+    for (host, port), block_dict in transfer_mappings.items():
+        block_dict["trans_count"] = d_trans_count_mapping[(host, port)]
+    logger.debug("MooncakeLayerwiseConnector Request %s transfer tasks: %s", req_id, transfer_mappings)
+    return transfer_mappings
+
+
+@dataclass
+class RegisterRange:
+    start: int
+    end: int
+
+
+@dataclass
+class RegisterRegions:
+    ptrs: list[int]
+    lengths: list[int]
+    logical_tensor_count: int | None = None
+    logical_total_bytes: int | None = None
+
+    @property
+    def registered_bytes(self) -> int:
+        return sum(self.lengths)
+
+
+def iter_kv_cache_tensors(obj: Any) -> Iterator[torch.Tensor]:
+    """Flatten kv_caches into tensors without materializing new tensors."""
+    if obj is None:
+        return
+
+    if isinstance(obj, torch.Tensor):
+        yield obj
+        return
+
+    if isinstance(obj, (tuple, list)):
+        for item in obj:
+            yield from iter_kv_cache_tensors(item)
+        return
+
+    if isinstance(obj, dict):
+        for item in obj.values():
+            yield from iter_kv_cache_tensors(item)
+        return
+
+
+def tensor_storage_key(tensor: torch.Tensor) -> int:
+    """Return a stable grouping key for tensors sharing the same storage.
+
+    Do NOT use this key as the register address directly. For aligned KV cache
+    views, tensor.untyped_storage().data_ptr() may point to the original raw
+    allocation, whose address can be unaligned. We only use it to group views.
+    """
+    try:
+        return tensor.untyped_storage().data_ptr()
+    except Exception:
+        try:
+            return tensor.storage().data_ptr()
+        except Exception:
+            return tensor.data_ptr()
+
+
+def collect_storage_merged_register_regions(
+    kv_caches: dict[str, Any],
+) -> RegisterRegions:
+    """Collect HCCL/Mooncake register regions with storage-aware merging.
+
+    Metadata should still use each logical tensor's own data_ptr().
+    register_buffer should use the merged memory ranges returned here.
+    """
+    ranges_by_storage: OrderedDict[int, list[RegisterRange]] = OrderedDict()
+    logical_tensor_count = 0
+    logical_total_bytes = 0
+
+    for tensor in iter_kv_cache_tensors(kv_caches):
+        if tensor is None or tensor.numel() == 0:
+            continue
+
+        if not tensor.is_contiguous():
+            logger.warning(
+                "Mooncake register_buffer got a non-contiguous KV cache "
+                "tensor: shape=%s, dtype=%s, data_ptr=%s. "
+                "Registration will use logical numel * element_size.",
+                tuple(tensor.shape),
+                tensor.dtype,
+                hex(tensor.data_ptr()),
+            )
+
+        nbytes = tensor.nbytes
+        start = tensor.data_ptr()
+        end = start + nbytes
+        storage_key = tensor_storage_key(tensor)
+
+        logical_tensor_count += 1
+        logical_total_bytes += nbytes
+
+        ranges_by_storage.setdefault(storage_key, []).append(RegisterRange(start, end))
+
+    register_ptrs: list[int] = []
+    register_lengths: list[int] = []
+
+    for ranges in ranges_by_storage.values():
+        ranges.sort(key=lambda r: r.start)
+
+        merged_start = ranges[0].start
+        merged_end = ranges[0].end
+
+        for region in ranges[1:]:
+            if region.start <= merged_end + REGISTER_MERGE_GAP_BYTES:
+                merged_end = max(merged_end, region.end)
+            else:
+                register_ptrs.append(merged_start)
+                register_lengths.append(merged_end - merged_start)
+                merged_start = region.start
+                merged_end = region.end
+
+        register_ptrs.append(merged_start)
+        register_lengths.append(merged_end - merged_start)
+
+    return RegisterRegions(
+        ptrs=register_ptrs,
+        lengths=register_lengths,
+        logical_tensor_count=logical_tensor_count,
+        logical_total_bytes=logical_total_bytes,
+    )
+
+
+def validate_register_region_count(regions: RegisterRegions) -> None:
+    region_count = len(regions.ptrs)
+    if region_count <= MAX_HCCL_REGISTER_REGIONS:
+        return
+
+    detail = f"registered_bytes={regions.registered_bytes}"
+    if regions.logical_tensor_count is not None:
+        detail += f", logical_tensors={regions.logical_tensor_count}, logical_bytes={regions.logical_total_bytes}"
+
+    raise RuntimeError(
+        "Mooncake register_buffer region count "
+        f"{region_count} exceeds HCCL per-process limit "
+        f"{MAX_HCCL_REGISTER_REGIONS}. "
+        "KV cache registration would fail. "
+        f"{detail}. "
+        "Please reduce KV cache allocation fragmentation or merge "
+        "k/v/dsa/scale allocations further."
+    )

vllm_ascend/distributed/parallel_state.py

@@ -1,125 +1,287 @@
-from typing import Optional
-
 import torch
-from vllm.config import ParallelConfig
-from vllm.distributed.parallel_state import (GroupCoordinator, get_world_group,
-                                             init_model_parallel_group)
+from vllm.config import ParallelConfig, get_current_vllm_config
+from vllm.distributed.parallel_state import GroupCoordinator, get_tp_group, get_world_group, init_model_parallel_group
 
-import vllm_ascend.envs as envs_ascend
 from vllm_ascend.ascend_config import get_ascend_config
+from vllm_ascend.utils import enable_dsa_cp_with_layer_shard, flashcomm2_enable
 
 # Currently, mc2 op need their own group coordinator.
-_MC2: Optional[GroupCoordinator] = None
-_MLP_TP: Optional[GroupCoordinator] = None
-_OTP: Optional[GroupCoordinator] = None
-_LMTP: Optional[GroupCoordinator] = None
+_MC2: GroupCoordinator | None = None
+
+# Module specific tensor parallel groups
+_MLP_TP: GroupCoordinator | None = None
+_OTP: GroupCoordinator | None = None
+_LMTP: GroupCoordinator | None = None
+_EMBED_TP: GroupCoordinator | None = None
+
+# flashcomm specific groups
+_FLASHCOMM2_OTP: GroupCoordinator | None = None
+_FLASHCOMM2_ODP: GroupCoordinator | None = None
+_FC3_QUANT_X: GroupCoordinator | None = None
+
+# shard_weight across rank groups
+_SHARD_WEIGHT: GroupCoordinator | None = None
+
+_P_TP: GroupCoordinator | None = None
+
+_DYNAMIC_EPLB: GroupCoordinator | None = None
 
 
-def get_mc2_group() -> GroupCoordinator:
-    assert _MC2 is not None, ("mc2 group is not initialized")
-    return _MC2
-
-
-def get_otp_group() -> GroupCoordinator:
-    assert _OTP is not None, (
-        "output tensor parallel group is not initialized")
-    return _OTP
-
-
-def get_lmhead_tp_group() -> GroupCoordinator:
-    assert _LMTP is not None, (
-        "lm head tensor parallel group is not initialized")
-    return _LMTP
-
-
-def get_mlp_tp_group() -> GroupCoordinator:
-    assert _MLP_TP is not None, ("mlp group is not initialized")
-    return _MLP_TP
-
-
-def model_parallel_initialized():
-    return (_MC2 is not None)
-
-
-def init_ascend_model_parallel(parallel_config: ParallelConfig, ):
+def init_ascend_model_parallel(
+    parallel_config: ParallelConfig,
+):
     if model_parallel_initialized():
         return
     assert torch.distributed.is_initialized()
     world_size = torch.distributed.get_world_size()
     backend = torch.distributed.get_backend(get_world_group().device_group)
+    global_tp_size = parallel_config.tensor_parallel_size
+    global_dp_size = parallel_config.data_parallel_size
+    global_pp_size = parallel_config.pipeline_parallel_size
+    global_pcp_size = parallel_config.prefill_context_parallel_size
 
     # The layout of all ranks: ExternalDP * EP
     # ExternalDP is the data parallel group that is not part of the model,
     # every dp rank can generate independently (in verl integration).
     all_ranks = torch.arange(world_size).reshape(
-        -1, parallel_config.data_parallel_size *
-        parallel_config.tensor_parallel_size)
-    global _MC2
-    group_ranks = all_ranks.unbind(0)
+        -1,
+        global_dp_size,
+        global_pp_size,
+        global_pcp_size,
+        global_tp_size,
+    )
+
+    pd_tp_ratio = get_ascend_config().pd_tp_ratio
+    pd_head_ratio = get_ascend_config().pd_head_ratio
+    global _P_TP
+    assert _P_TP is None, "distributed prefill tensor parallel group is already initialized"
+    prefill_tensor_model_parallel_size = pd_tp_ratio
+    # divide alltoall groups
+    if pd_head_ratio > 1 and get_current_vllm_config().kv_transfer_config.is_kv_producer:
+        num_head_replica = get_ascend_config().num_head_replica
+        remote_tp_size = global_tp_size // pd_tp_ratio
+        if num_head_replica <= 1:
+            group_ranks = all_ranks.view(-1, prefill_tensor_model_parallel_size).unbind(0)
+        else:
+            group_ranks = all_ranks.clone().view(
+                global_dp_size * global_pp_size * global_pcp_size, -1, num_head_replica
+            )  # [DP_size, num_head, num_head_replica]
+            group_ranks = group_ranks.permute(0, 2, 1)
+            group_ranks = group_ranks.reshape(-1, group_ranks.size(-1))  # [DP_size * num_head_replica, num_head]
+            alltoall_group_size = group_ranks.size(-1) // remote_tp_size
+            group_ranks = group_ranks.unsqueeze(-1).view(
+                global_dp_size * global_pp_size * global_pcp_size,
+                num_head_replica,
+                -1,
+                alltoall_group_size,
+            )  # [DP_size, num_head_replica, num_alltoall_group, alltoall_group_size]
+            group_ranks = group_ranks.reshape(-1, alltoall_group_size).unbind(0)
+        group_ranks = [x.tolist() for x in group_ranks]
+        local_rank = get_world_group().local_rank
+        num = next((i for i, ranks in enumerate(group_ranks) if local_rank in ranks), None)
+        _P_TP = init_model_parallel_group(group_ranks, get_world_group().local_rank, backend, group_name=f"p_tp_{num}")
+
+    # EP like group ranks
+    group_ranks = (
+        all_ranks.transpose(1, 2)
+        .reshape(
+            -1,
+            global_dp_size * global_pcp_size * global_tp_size,
+        )
+        .unbind(0)
+    )
     group_ranks = [x.tolist() for x in group_ranks]
 
-    _MC2 = init_model_parallel_group(group_ranks,
-                                     get_world_group().local_rank,
-                                     backend,
-                                     group_name="mc2")
-    if envs_ascend.VLLM_ASCEND_ENABLE_MLP_OPTIMIZE:
-        global _MLP_TP
-        assert _MLP_TP is None, (
-            "mlp tensor model parallel group is already initialized")
+    global _MC2
+    _MC2 = init_model_parallel_group(group_ranks, get_world_group().local_rank, backend, group_name="mc2")
 
-        mlp_tp = parallel_config.data_parallel_size
+    if get_ascend_config().eplb_config.dynamic_eplb:
+        global _DYNAMIC_EPLB
+        _DYNAMIC_EPLB = init_model_parallel_group(
+            group_ranks, get_world_group().local_rank, backend, group_name="dynamic_eplb"
+        )
 
-        all_ranks_mlp_head = torch.arange(world_size).reshape(
-            -1, mlp_tp, parallel_config.pipeline_parallel_size, 1)  # noqa
-        group_ranks = all_ranks_mlp_head.view(-1, mlp_tp).unbind(0)
-        group_ranks = [x.tolist() for x in group_ranks]
+    if get_ascend_config().multistream_overlap_gate:
+        global _FC3_QUANT_X
+        _FC3_QUANT_X = init_model_parallel_group(
+            group_ranks, get_world_group().local_rank, backend, group_name="fc3_quant_x"
+        )
 
-        # message queue broadcaster is only used in tensor model parallel group
-        _MLP_TP = init_model_parallel_group(group_ranks,
-                                            get_world_group().local_rank,
-                                            backend,
-                                            group_name="mlp_tp")
+    # Initialize fine-grained TP process groups on Ascend for four components:
+    # 1. LM Head: output logits projection (`lmhead_tensor_parallel_size`)
+    # 2. O Proj: attention output projection (`oproj_tensor_parallel_size`)
+    # 3. Embedding: The token embedding table at the input of the model (`embedding_tensor_parallel_size`)
+    # 4. MLP: feed-forward network in transformer blocks (`mlp_tensor_parallel_size`)
+    _group_cache = {}
 
-    # If oproj tensor parallel size is set, we will create a group for it.
-    otp_size = get_ascend_config().oproj_tensor_parallel_size
-    if otp_size is not None:
+    def _create_or_get_group(group_size: int, group_name: str) -> GroupCoordinator:
+        if group_size is None:
+            return None
+        if group_size not in _group_cache:
+            rank_grid = torch.arange(world_size).reshape(global_pp_size, global_dp_size, global_tp_size)
+            num_chunks = global_dp_size // group_size
+            group_ranks = []
+            for pp_idx in range(global_pp_size):
+                stage_ranks = rank_grid[pp_idx]  # (dp, tp)
+                for chunk in range(num_chunks):
+                    for tp_idx in range(global_tp_size):
+                        group = stage_ranks[chunk * group_size : (chunk + 1) * group_size, tp_idx].tolist()
+                        group_ranks.append(group)
+            pg = init_model_parallel_group(group_ranks, get_world_group().local_rank, backend, group_name=group_name)
+            _group_cache[group_size] = pg
+
+        return _group_cache[group_size]
+
+    otp_size = get_ascend_config().finegrained_tp_config.oproj_tensor_parallel_size
+    lmhead_tp_size = get_ascend_config().finegrained_tp_config.lmhead_tensor_parallel_size
+    embedding_tp_size = get_ascend_config().finegrained_tp_config.embedding_tensor_parallel_size
+    mlp_tp_size = get_ascend_config().finegrained_tp_config.mlp_tensor_parallel_size
+
+    global _OTP, _LMTP, _EMBED_TP, _MLP_TP
+
+    if otp_size > 0:
+        _OTP = _create_or_get_group(otp_size, "otp")
+    if lmhead_tp_size > 0:
+        _LMTP = _create_or_get_group(lmhead_tp_size, "lmheadtp")
+    if embedding_tp_size > 0:
+        _EMBED_TP = _create_or_get_group(embedding_tp_size, "emtp")
+    if mlp_tp_size > 0:
+        _MLP_TP = _create_or_get_group(mlp_tp_size, "mlptp")
+
+    # TODO: Extract and unify the logic across different communication group.
+    flashcomm2_otp_group_ranks = []
+    if flashcomm2_enable():
+        flashcomm2_otp_size = get_ascend_config().flashcomm2_oproj_tensor_parallel_size
+        num_fc2_oproj_tensor_parallel_groups: int = global_tp_size // flashcomm2_otp_size
+        global _FLASHCOMM2_OTP
+        global _FLASHCOMM2_ODP
+
+        _FLASHCOMM2_OTP = None
+        _FLASHCOMM2_ODP = get_tp_group()
+
+        if flashcomm2_otp_size > 1:
+            odp_group_ranks: list[list[int]] = [
+                [] for _ in range(flashcomm2_otp_size * global_dp_size * global_pp_size)
+            ]
+            for dp_group_index in range(global_dp_size):
+                for pp_group_index in range(global_pp_size):
+                    dp_pp_serial_index = dp_group_index * global_pp_size + pp_group_index
+                    tp_base_rank = dp_pp_serial_index * global_tp_size
+                    odp_base_index = dp_pp_serial_index * flashcomm2_otp_size
+
+                    for i in range(num_fc2_oproj_tensor_parallel_groups):
+                        ranks = []
+                        for j in range(flashcomm2_otp_size):
+                            tp_local_rank = i + j * num_fc2_oproj_tensor_parallel_groups
+                            assert tp_local_rank < global_tp_size
+                            global_rank = tp_base_rank + tp_local_rank
+                            ranks.append(global_rank)
+
+                            odp_group_index = odp_base_index + j
+                            odp_group_ranks[odp_group_index].append(global_rank)
+                        flashcomm2_otp_group_ranks.append(ranks)
+
+            _FLASHCOMM2_OTP = init_model_parallel_group(
+                flashcomm2_otp_group_ranks, get_world_group().local_rank, backend, group_name="flashcomm2_otp"
+            )
+            _FLASHCOMM2_ODP = init_model_parallel_group(
+                odp_group_ranks, get_world_group().local_rank, backend, group_name="flashcomm2_odp"
+            )
+
+    def create_shard_weight_group(module_tp_group_ranks: None) -> GroupCoordinator:
+        # Argument module_tp_group_ranks: The module specific tensor parallel group.
+        # There are three situations.
+        # 1. If it is None, then the TP_size of the specific module is 1 and is replicated linear layer.
+        # 2. If it is not None, and the module tp_group is same as the global tp_group.
+        # 3. If it is not None, and the module tp_group is different from the global tp_group.(eg. flashcomm2_otp)
         group_ranks = []
-        global _OTP
-        num_oproj_tensor_parallel_groups: int = (world_size // otp_size)
-        for i in range(num_oproj_tensor_parallel_groups):
-            ranks = list(range(i * otp_size, (i + 1) * otp_size))
-            group_ranks.append(ranks)
-        _OTP = init_model_parallel_group(group_ranks,
-                                         get_world_group().local_rank,
-                                         backend,
-                                         group_name="otp")
+        pp_group_ranks = all_ranks.transpose(2, 4).reshape(-1, global_pp_size)
+        if module_tp_group_ranks is None:
+            # If it is None, then the TP_size of this shard weight is 1.
+            shard_weight_group_ranks = pp_group_ranks.transpose(0, 1).unbind(0)
+            group_ranks = [x.tolist() for x in shard_weight_group_ranks]
+        else:
+            # combine standard tp group and non-standard tp group to build  shard_weight comm_group
+            module_tp_tanspose_ranks = module_tp_group_ranks.transpose(0, 1)
+            G = world_size // (global_pp_size * module_tp_group_ranks.size(1))
+            shard_weight_group_ranks = torch.stack([t.view(global_pp_size, G) for t in module_tp_tanspose_ranks], dim=1)
+            group_ranks = shard_weight_group_ranks.view(-1, G).tolist()
+        return init_model_parallel_group(group_ranks, get_world_group().local_rank, backend, group_name="shard_weight")
 
-    lmhead_tensor_parallel_size = get_ascend_config(
-    ).lmhead_tensor_parallel_size
-    if lmhead_tensor_parallel_size is not None:
-        group_ranks = []
-        global _LMTP
-        num_lmhead_tensor_parallel_groups: int = (world_size //
-                                                  lmhead_tensor_parallel_size)
-        for i in range(num_lmhead_tensor_parallel_groups):
-            ranks = list(
-                range(i * lmhead_tensor_parallel_size,
-                      (i + 1) * lmhead_tensor_parallel_size))
-            group_ranks.append(ranks)
-        _LMTP = init_model_parallel_group(group_ranks,
-                                          get_world_group().local_rank,
-                                          backend,
-                                          group_name="lmheadtp")
+    # Create shard weight group if enabled
+    if get_ascend_config().layer_sharding is not None:
+        global _SHARD_WEIGHT
+        if flashcomm2_enable():
+            if len(flashcomm2_otp_group_ranks) == 0:
+                FC2_group_ranks = None
+            else:
+                FC2_group_ranks = torch.tensor(flashcomm2_otp_group_ranks).squeeze(0)
+            _SHARD_WEIGHT = create_shard_weight_group(FC2_group_ranks)
+        elif enable_dsa_cp_with_layer_shard():
+            # For dsa_cp, all shard layers are replicated.
+            _SHARD_WEIGHT = create_shard_weight_group(None)
+        else:
+            # For standard tp, use global tp group_ranks
+            tp_group_ranks = all_ranks.view(-1, global_tp_size)
+            _SHARD_WEIGHT = create_shard_weight_group(tp_group_ranks)
 
 
-def get_mlp_tensor_model_parallel_world_size():
-    """Return world size for the tensor model parallel group."""
-    return get_mlp_tp_group().world_size
+def model_parallel_initialized():
+    return _MC2 is not None
 
 
-def get_mlp_tensor_model_parallel_rank():
-    """Return world size for the tensor model parallel group."""
-    return get_mlp_tp_group().rank_in_group
+def get_mc2_group() -> GroupCoordinator:
+    assert _MC2 is not None, "mc2 group is not initialized"
+    return _MC2
+
+
+def get_mlp_tp_group() -> GroupCoordinator:
+    assert _MLP_TP is not None, "mlp group is not initialized"
+    return _MLP_TP
+
+
+def get_otp_group() -> GroupCoordinator:
+    assert _OTP is not None, "output tensor parallel group is not initialized"
+    return _OTP
+
+
+def get_lmhead_tp_group() -> GroupCoordinator:
+    assert _LMTP is not None, "lm head tensor parallel group is not initialized"
+    return _LMTP
+
+
+def get_embed_tp_group() -> GroupCoordinator:
+    assert _EMBED_TP is not None, "emtp group is not initialized"
+    return _EMBED_TP
+
+
+def get_flashcomm2_otp_group() -> GroupCoordinator:
+    return _FLASHCOMM2_OTP
+
+
+def get_flashcomm2_odp_group() -> GroupCoordinator:
+    assert _FLASHCOMM2_ODP is not None, "output data parallel group for flashcomm2 is not initialized"
+    return _FLASHCOMM2_ODP
+
+
+def get_shard_weight_group() -> GroupCoordinator:
+    assert _SHARD_WEIGHT is not None, "output shard weight parallel group for flashcomm2 is not initialized"
+    return _SHARD_WEIGHT
+
+
+def get_p_tp_group() -> GroupCoordinator:
+    assert _P_TP is not None, "distributed prefill tensor parallel group is not initialized"
+    return _P_TP
+
+
+def get_fc3_quant_x_group() -> GroupCoordinator:
+    assert _FC3_QUANT_X is not None, "fc3 quant x group is not initialized"
+    return _FC3_QUANT_X
+
+
+def get_dynamic_eplb_group() -> GroupCoordinator:
+    assert _DYNAMIC_EPLB is not None, "Dynamic eplb group is not initialized"
+    return _DYNAMIC_EPLB
 
 
 def destroy_ascend_model_parallel():
@@ -138,7 +300,75 @@ def destroy_ascend_model_parallel():
         _LMTP.destroy()
     _LMTP = None
 
+    global _EMBED_TP
+    if _EMBED_TP:
+        _EMBED_TP.destroy()
+    _EMBED_TP = None
+
     global _OTP
     if _OTP:
         _OTP.destroy()
     _OTP = None
+
+    global _P_TP
+    if _P_TP:
+        _P_TP.destroy()
+    _P_TP = None
+
+    global _FLASHCOMM2_OTP
+    if _FLASHCOMM2_OTP and get_ascend_config().flashcomm2_oproj_tensor_parallel_size != 1:
+        _FLASHCOMM2_OTP.destroy()
+        _FLASHCOMM2_OTP = None
+
+    global _FLASHCOMM2_ODP
+    if _FLASHCOMM2_ODP and get_ascend_config().flashcomm2_oproj_tensor_parallel_size != 1:
+        _FLASHCOMM2_ODP.destroy()
+        _FLASHCOMM2_ODP = None
+
+    global _SHARD_WEIGHT
+    if _SHARD_WEIGHT:
+        _SHARD_WEIGHT.destroy()
+    _SHARD_WEIGHT = None
+
+    global _FC3_QUANT_X
+    if _FC3_QUANT_X:
+        _FC3_QUANT_X.destroy()
+    _FC3_QUANT_X = None
+
+    global _DYNAMIC_EPLB
+    if _DYNAMIC_EPLB:
+        _DYNAMIC_EPLB.destroy()
+    _DYNAMIC_EPLB = None
+
+
+def get_global_rank(parallel_config: ParallelConfig | None = None) -> int:
+    """Return a globally unique rank for the current worker across all parallel
+     dimensions (TP/PP/CP/DP), compatible with both dense and MoE models.
+
+     vLLM does not expose a single ready-to-use cross-DP global rank:
+       - For dense models each DP rank is launched as an independent DP=1 engine,
+         so ``data_parallel_rank`` is reset to 0 and ``get_world_group()`` only
+         spans one replica (``rank_in_group`` is the local rank in the replica).
+       - For MoE DP / external_launcher the world group spans all DP ranks, so
+         ``rank_in_group`` already encodes the DP offset.
+
+     ``data_parallel_index`` always keeps the true DP rank (it is never reset),
+     and ``rank_in_group % replica_size`` yields the local rank within a replica
+     in both cases, so the formula below is correct everywhere. It mirrors vLLM's
+     own ``data_parallel_rank * world_size + rank`` (see
+     vllm/distributed/parallel_state.py).
+
+    Note: DCP (decode context parallel) reuses the TP NPUs and EP overlays
+    TP/DP, so neither adds new ranks and they are intentionally excluded from
+    ``replica_size``.
+    """
+    if parallel_config is None:
+        parallel_config = get_current_vllm_config().parallel_config
+    # Number of NPUs in a single DP replica (TP * PP * prefill-CP).
+    replica_size = (
+        parallel_config.tensor_parallel_size
+        * parallel_config.pipeline_parallel_size
+        * parallel_config.prefill_context_parallel_size
+    )
+    rank_in_replica = get_world_group().rank_in_group % replica_size
+    return parallel_config.data_parallel_index * replica_size + rank_in_replica

vllm_ascend/distributed/utils.py

@@ -0,0 +1,88 @@
+import torch
+import torch.distributed as dist
+from vllm.distributed import get_dcp_group
+from vllm.distributed.parallel_state import GroupCoordinator, get_dp_group
+from vllm.forward_context import get_forward_context
+
+from vllm_ascend.ascend_forward_context import _EXTRA_CTX
+from vllm_ascend.distributed.parallel_state import get_fc3_quant_x_group
+
+
+def get_decode_context_model_parallel_world_size() -> int:
+    """Return DCP world size (v0.21.0 helper removed on vLLM main)."""
+    return get_dcp_group().world_size
+
+
+def get_decode_context_model_parallel_rank() -> int:
+    """Return DCP rank within group (v0.21.0 helper removed on vLLM main)."""
+    return get_dcp_group().rank_in_group
+
+
+def fc3_all_gather_and_maybe_unpad_impl(
+    x: torch.Tensor,
+) -> torch.Tensor:
+    try:
+        forward_context = get_forward_context()
+    except AssertionError:
+        return x
+    x = get_fc3_quant_x_group().all_gather(x, 0)
+    dp_metadata = forward_context.dp_metadata
+    if dp_metadata is None:
+        pad_size = _EXTRA_CTX.pad_size
+        if pad_size > 0:
+            x = x[:-pad_size]
+    else:
+        # unpad
+        num_tokens_across_dp_cpu = dp_metadata.num_tokens_across_dp_cpu
+        result = torch.empty((num_tokens_across_dp_cpu.sum(), *x.shape[1:]), device=x.device, dtype=x.dtype)
+        dp_size = get_dp_group().world_size
+        x = x.view(dp_size, _EXTRA_CTX.padded_length, *x.shape[1:])
+        offset = 0
+        for idx in range(dp_size):
+            num_tokens_dp = num_tokens_across_dp_cpu[idx]
+            result[offset : offset + num_tokens_dp] = x[idx, :num_tokens_dp]
+            offset += num_tokens_dp
+        x = result
+
+    return x
+
+
+def all_gather_async(
+    input: torch.Tensor, group: GroupCoordinator, output: torch.Tensor | None = None, async_op: bool = True
+):
+    if group.world_size == 1:
+        return input, None
+    if output is None:
+        input_size = input.size()
+        output_size = (input_size[0] * group.world_size,) + input_size[1:]
+        output = torch.empty(output_size, dtype=input.dtype, device=input.device)
+    return output, dist.all_gather_into_tensor(output, input, group=group.device_group, async_op=async_op)
+
+
+def split_tensor_along_first_dim(
+    tensor: torch.Tensor,
+    num_partitions: int,
+    contiguous_split_chunks: bool = False,
+):
+    """Split a tensor along its first dimension.
+
+    Arguments:
+        tensor: input tensor.
+        num_partitions: number of partitions to split the tensor
+        contiguous_split_chunks: If True, make each chunk contiguous
+                                in memory.
+
+    Returns:
+        A list of Tensors
+    """
+    from vllm.distributed.utils import divide
+
+    # Get the size and dimension.
+    first_dim_size = divide(tensor.size()[0], num_partitions)
+    # Split.
+    tensor_list = torch.split(tensor, first_dim_size, dim=0)
+    # NOTE: torch.split does not create contiguous tensors by default.
+    if contiguous_split_chunks:
+        return tuple(chunk.contiguous() for chunk in tensor_list)
+
+    return tensor_list

vllm_ascend/distributed/weight_transfer/__init__.py

@@ -0,0 +1,32 @@
+#
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm-ascend project.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
+
+
+def register_engine():
+    """Register Ascend weight transfer engines as vLLM plugins."""
+    WeightTransferEngineFactory.register_engine(
+        "hccl",
+        "vllm_ascend.distributed.weight_transfer.hccl_engine",
+        "HCCLWeightTransferEngine",
+    )
+    WeightTransferEngineFactory.register_engine(
+        "npu_ipc",
+        "vllm_ascend.distributed.weight_transfer.npu_ipc_engine",
+        "NPUIPCWeightTransferEngine",
+    )

vllm_ascend/distributed/weight_transfer/hccl_engine.py

@@ -0,0 +1,336 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""HCCL-based weight transfer engine."""
+
+from collections.abc import Callable, Iterator
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+if TYPE_CHECKING:
+    from vllm_ascend.distributed.device_communicators.pyhccl import PyHcclCommunicator
+
+from vllm.config.parallel import ParallelConfig
+from vllm.config.weight_transfer import WeightTransferConfig
+from vllm.distributed.weight_transfer.base import (
+    WeightTransferEngine,
+    WeightTransferInitInfo,
+    WeightTransferUpdateInfo,
+)
+
+from vllm_ascend.distributed.weight_transfer.packed_tensor import (
+    DEFAULT_PACKED_BUFFER_SIZE_BYTES,
+    DEFAULT_PACKED_NUM_BUFFERS,
+    packed_broadcast_consumer,
+)
+
+
+@dataclass
+class HCCLWeightTransferInitInfo(WeightTransferInitInfo):
+    """Initialization info for HCCL weight transfer backend."""
+
+    master_address: str
+    """IP address of the trainer (rank 0) for HCCL process group setup."""
+    master_port: int
+    """Port on the trainer for HCCL process group setup."""
+    rank_offset: int
+    """Offset added to each vLLM worker's rank within the HCCL group.
+    Typically 1 (trainer is rank 0, workers start at rank 1)."""
+    world_size: int
+    """Total number of participants in the HCCL group (trainer + all workers)."""
+
+
+@dataclass
+class HCCLTrainerSendWeightsArgs:
+    """Arguments for HCCL trainer_send_weights method."""
+
+    group: Any
+    """Process group (PyHcclCommunicator) for HCCL communication."""
+    src: int = 0
+    """Source rank (default 0, trainer is typically rank 0)."""
+    post_iter_func: Callable[[tuple[str, torch.Tensor]], torch.Tensor] | None = None
+    """Optional function to apply to each (name, tensor) pair before broadcasting.
+    If None, extracts just the tensor."""
+    packed: bool = False
+    """Whether to use packed tensor broadcasting for efficiency.
+    When True, multiple tensors are batched together before broadcasting
+    to reduce HCCL communication overhead."""
+    stream: torch.npu.Stream | None = None
+    """ACL stream to use for broadcasting if packed is False.
+    If packed is True, new streams will be created for each buffer."""
+    packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES
+    """Size in bytes for each packed tensor buffer.
+    Must match the value used in HCCLWeightTransferUpdateInfo."""
+    packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS
+    """Number of buffers for double/triple buffering during packed transfer.
+    Must match the value used in HCCLWeightTransferUpdateInfo."""
+
+
+@dataclass
+class HCCLWeightTransferUpdateInfo(WeightTransferUpdateInfo):
+    """Update info for HCCL weight transfer backend."""
+
+    names: list[str]
+    """Names of the parameters to transfer (e.g. ``model.layers.0.weight``)."""
+    dtype_names: list[str]
+    """Torch dtype names (e.g. ``bfloat16``, ``float32``) for each parameter."""
+    shapes: list[list[int]]
+    """Shapes of each parameter as integer lists."""
+    packed: bool = False
+    """Whether to use packed tensor broadcasting for efficiency.
+    When True, multiple tensors are batched together before broadcasting
+    to reduce HCCL communication overhead."""
+    packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES
+    """Size in bytes for each packed tensor buffer.
+    Both producer and consumer must use the same value."""
+    packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS
+    """Number of buffers for double/triple buffering during packed transfer.
+    Both producer and consumer must use the same value."""
+
+    def __post_init__(self):
+        """Validate that all lists have the same length."""
+        num_params = len(self.names)
+        if len(self.dtype_names) != num_params:
+            raise ValueError(
+                f"`dtype_names` should be of the same size as `names`: "
+                f"got {len(self.dtype_names)} and {len(self.names)}"
+            )
+        if len(self.shapes) != num_params:
+            raise ValueError(
+                f"`shapes` should be of the same size as `names`: got {len(self.shapes)} and {len(self.names)}"
+            )
+
+
+class HCCLWeightTransferEngine(WeightTransferEngine[HCCLWeightTransferInitInfo, HCCLWeightTransferUpdateInfo]):
+    """
+    Weight transfer engine using HCCL for communication between trainer and workers.
+
+    This implementation uses HCCL broadcast operations to transfer weights from
+    the trainer (rank 0) to all inference workers in a process group.
+    """
+
+    # Define backend-specific dataclass types
+    init_info_cls = HCCLWeightTransferInitInfo
+    update_info_cls = HCCLWeightTransferUpdateInfo
+
+    def __init__(
+        self,
+        config: WeightTransferConfig,
+        parallel_config: ParallelConfig,
+        model: torch.nn.Module | None = None,
+    ) -> None:
+        """
+        Initialize the HCCL weight transfer engine.
+
+        Args:
+            config: The configuration for the weight transfer engine
+            parallel_config: The configuration for the parallel setup
+            model: The local model instance which will receive the weights.
+        """
+        super().__init__(config, parallel_config, model)
+        self.model_update_group: PyHcclCommunicator | None = None
+
+    def init_transfer_engine(self, init_info: HCCLWeightTransferInitInfo) -> None:
+        """
+        Initialize HCCL process group with the trainer.
+
+        Args:
+            init_info: HCCL initialization info containing master address, port,
+                      rank offset, and world size
+        """
+
+        # Calculate the global rank in the trainer-worker process group
+        # Must account for data parallel to get unique ranks across all workers
+        dp_rank = self.parallel_config.data_parallel_index
+        world_size_per_dp = self.parallel_config.world_size  # TP * PP
+        rank_within_dp = self.parallel_config.rank
+
+        # Unique rank across all DP groups
+        worker_rank = dp_rank * world_size_per_dp + rank_within_dp
+        rank = worker_rank + init_info.rank_offset
+        # Create stateless process group
+        device = torch.accelerator.current_device_index()
+        self.model_update_group = HCCLWeightTransferEngine._stateless_init_process_group(
+            init_info.master_address,
+            init_info.master_port,
+            rank,
+            init_info.world_size,
+            device=device,
+        )
+
+    def receive_weights(
+        self,
+        update_info: HCCLWeightTransferUpdateInfo,
+        load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
+    ) -> None:
+        """
+        Receive weights from trainer via HCCL broadcast and load them incrementally.
+
+        If update_info.packed is True, uses packed tensor broadcasting for
+        efficient transfer of multiple weights in batches. Otherwise, uses simple
+        one-by-one broadcasting.
+
+        Args:
+            update_info: HCCL update info containing parameter names, dtypes, shapes,
+                        and packed flag
+            load_weights: Callable that loads weights into the model. Called
+                         incrementally for each batch of weights to avoid OOM.
+        """
+        if self.model_update_group is None:
+            raise RuntimeError("HCCL weight transfer not initialized. Call init_transfer_engine() first.")
+
+        if update_info.packed:
+            # Build iterator of (name, (shape, dtype)) from update_info
+            def state_dict_info_iterator():
+                for name, dtype_name, shape in zip(update_info.names, update_info.dtype_names, update_info.shapes):
+                    dtype = getattr(torch, dtype_name)
+                    yield (name, (shape, dtype))
+
+            packed_broadcast_consumer(
+                iterator=state_dict_info_iterator(),
+                group=self.model_update_group,
+                src=0,
+                post_unpack_func=load_weights,
+                buffer_size_bytes=update_info.packed_buffer_size_bytes,
+                num_buffers=update_info.packed_num_buffers,
+            )
+        else:
+            # Use simple one-by-one broadcasting
+            for name, dtype_name, shape in zip(update_info.names, update_info.dtype_names, update_info.shapes):
+                dtype = getattr(torch, dtype_name)
+                weight = torch.empty(shape, dtype=dtype, device="npu")
+                self.model_update_group.broadcast(weight, src=0, stream=torch.npu.current_stream())
+                load_weights([(name, weight)])
+                del weight
+
+    def shutdown(self) -> None:
+        if self.model_update_group is not None:
+            # Clean up the communicator by removing the reference
+            self.model_update_group = None
+
+    @staticmethod
+    def trainer_send_weights(
+        iterator: Iterator[tuple[str, torch.Tensor]],
+        trainer_args: dict[str, Any] | HCCLTrainerSendWeightsArgs,
+    ) -> None:
+        """Broadcast weights from trainer to vLLM workers.
+
+        Args:
+            iterator: Iterator of model parameters. Returns (name, tensor) tuples
+            trainer_args: Dictionary or HCCLTrainerSendWeightsArgs instance containing
+                         HCCL-specific arguments. If a dict, should contain keys from
+                         HCCLTrainerSendWeightsArgs.
+
+        Example:
+            >>> from vllm.distributed.weight_transfer.hccl_engine import (
+            ...     HCCLWeightTransferEngine,
+            ...     HCCLTrainerSendWeightsArgs,
+            ... )
+            >>> param_iter = ((n, p) for n, p in model.named_parameters())
+            >>> args = HCCLTrainerSendWeightsArgs(group=group, packed=True)
+            >>> HCCLWeightTransferEngine.trainer_send_weights(param_iter, args)
+        """
+        # Parse trainer args - accept either dict or dataclass instance
+        if isinstance(trainer_args, dict):
+            args = HCCLTrainerSendWeightsArgs(**trainer_args)
+        else:
+            args = trainer_args
+
+        if args.post_iter_func is None:
+            # Default: extract just the tensor from (name, tensor) tuple
+            post_iter_func = lambda x: x[1]
+        else:
+            post_iter_func = args.post_iter_func
+
+        if args.packed:
+            # Use packed tensor broadcasting for efficiency
+            from vllm_ascend.distributed.weight_transfer.packed_tensor import (
+                packed_broadcast_producer,
+            )
+
+            packed_broadcast_producer(
+                iterator=iterator,
+                group=args.group,
+                src=args.src,
+                post_iter_func=post_iter_func,
+                buffer_size_bytes=args.packed_buffer_size_bytes,
+                num_buffers=args.packed_num_buffers,
+            )
+        else:
+            # Use simple one-by-one broadcasting
+            for item in iterator:
+                tensor = post_iter_func(item)
+                args.group.broadcast(
+                    tensor,
+                    src=args.src,
+                    stream=args.stream or torch.npu.current_stream(),
+                )
+
+    @staticmethod
+    def trainer_init(
+        init_info: HCCLWeightTransferInitInfo | dict,
+    ) -> "PyHcclCommunicator":
+        """
+        Initialize HCCL process group for trainer-side weight transfer.
+
+        The trainer is always rank 0 in the process group. Uses the current
+        Ascend device (torch.accelerator.current_device_index()).
+
+        Args:
+            init_info: Either an HCCLWeightTransferInitInfo object or a dict with keys:
+                - master_address: str
+                - master_port: int
+                - world_size: int
+
+        Returns:
+            PyHcclCommunicator for weight transfer.
+
+        Example:
+            >>> from vllm.distributed.weight_transfer.hccl_engine import (
+            ...     HCCLWeightTransferEngine,
+            ... )
+            >>> group = HCCLWeightTransferEngine.trainer_init(
+            ...     dict(
+            ...         master_address=master_address,
+            ...         master_port=master_port,
+            ...         world_size=world_size,
+            ...     ),
+            ... )
+        """
+        if isinstance(init_info, dict):
+            master_address = init_info["master_address"]
+            master_port = init_info["master_port"]
+            world_size = init_info["world_size"]
+        else:
+            # HCCLWeightTransferInitInfo object
+            master_address = init_info.master_address
+            master_port = init_info.master_port
+            world_size = init_info.world_size
+
+        # Trainer is always rank 0
+        device = torch.accelerator.current_device_index()
+        return HCCLWeightTransferEngine._stateless_init_process_group(
+            master_address,
+            master_port,
+            0,
+            world_size,
+            device,
+        )
+
+    @staticmethod
+    def _stateless_init_process_group(master_address, master_port, rank, world_size, device):
+        """
+        vLLM provides `StatelessProcessGroup` to create a process group
+        without considering the global process group in torch.distributed.
+        It is recommended to create `StatelessProcessGroup`, and then initialize
+        the data-plane communication (HCCL) between external (train processes)
+        and vLLM workers.
+        """
+        from vllm.distributed.utils import StatelessProcessGroup
+
+        from vllm_ascend.distributed.device_communicators.pyhccl import PyHcclCommunicator
+
+        pg = StatelessProcessGroup.create(host=master_address, port=master_port, rank=rank, world_size=world_size)
+        pyhccl = PyHcclCommunicator(pg, device=device)
+        return pyhccl

vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py

@@ -0,0 +1,405 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NPU IPC-based weight transfer engine using Ascend IPC for communication."""
+
+import os
+import pickle
+import socket
+from collections.abc import Callable, Iterator
+from dataclasses import asdict, dataclass
+from functools import lru_cache
+from typing import Any
+
+import pybase64 as base64
+import requests
+import torch
+from torch.multiprocessing.reductions import reduce_tensor
+from vllm import envs
+from vllm.config.parallel import ParallelConfig
+from vllm.config.weight_transfer import WeightTransferConfig
+from vllm.distributed.weight_transfer.base import (
+    WeightTransferEngine,
+    WeightTransferInitInfo,
+)
+from vllm.distributed.weight_transfer.ipc_engine import (
+    IPCTrainerSendWeightsArgs,
+    IPCWeightTransferUpdateInfo,
+)
+
+from vllm_ascend.distributed.weight_transfer.packed_tensor import (
+    packed_npu_ipc_consumer,
+    packed_npu_ipc_producer,
+)
+
+
+@dataclass
+class NPUIPCTrainerSendWeightsArgs(IPCTrainerSendWeightsArgs):
+    """NPU IPC variant — inherits all fields and validation from the CUDA IPC
+    base class.  Only the ``send_mode`` callable type is widened to accept the
+    NPU update-info type."""
+
+    send_mode: str | Callable[["NPUIPCWeightTransferUpdateInfo"], None]
+
+
+@dataclass
+class NPUIPCWeightTransferInitInfo(WeightTransferInitInfo):
+    """Initialization info for NPU IPC weight transfer backend.
+
+    No initialization needed for NPU IPC.
+    """
+
+    pass
+
+
+@dataclass
+class NPUIPCWeightTransferUpdateInfo(IPCWeightTransferUpdateInfo):
+    """NPU IPC variant — inherits all fields and validation from the CUDA IPC
+    base class.  No overrides needed; the field types and ``__post_init__`` are
+    identical."""
+
+
+@lru_cache(maxsize=1)
+def get_ip() -> str:
+    try:
+        # try to get ip from network interface
+        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
+            s.connect(("8.8.8.8", 80))
+            return s.getsockname()[0]
+    except Exception:  # noqa: BLE001
+        # fallback to get ip from hostname
+        return socket.gethostbyname(socket.gethostname())
+
+
+@lru_cache(maxsize=1)
+def npu_generate_uuid() -> str:
+    """Generate a unique identifier for the current process's physical NPU chip.
+
+    Returns ``{host_ip}-{physical_chip_id}`` where ``host_ip`` is the local
+    machine's IP address and ``physical_chip_id`` is derived from the current
+    logical device index mapped through ``ASCEND_RT_VISIBLE_DEVICES``.
+
+    On Ascend NPU, ``torch.accelerator.current_device_index()`` returns the
+    *logical* device index. When ``ASCEND_RT_VISIBLE_DEVICES`` is set, it
+    maps logical indices to physical chip IDs (e.g., ``ASCEND_RT_VISIBLE_DEVICES=2,3``
+    means logical device 0 → physical chip 2, logical device 1 → physical chip 3).
+    If the env var is not set, the logical index is used directly as the
+    physical chip ID (identity mapping).
+
+    The result is cached because it is constant for the lifetime of the
+    process. Both the trainer and inference worker processes co-located
+    on the same physical NPU chip will produce the same UUID, which is
+    required for NPU IPC handle matching.
+    """
+    logical_device = torch.accelerator.current_device_index()
+    visible_devices = os.environ.get("ASCEND_RT_VISIBLE_DEVICES", None)
+    if visible_devices:
+        physical_device = int(visible_devices.split(",")[logical_device].strip())
+    else:
+        physical_device = logical_device
+    return f"{get_ip()}-{physical_device}"
+
+
+class NPUIPCWeightTransferEngine(WeightTransferEngine[NPUIPCWeightTransferInitInfo, NPUIPCWeightTransferUpdateInfo]):
+    """
+    Weight transfer engine using NPU IPC for communication between
+    trainer and workers.
+
+    This implementation uses Ascend NPU IPC to transfer weights from the
+    trainer (rank 0) to all inference workers. IPC handles are used to
+    share memory between processes on the same node.
+
+    Requires ``torch_npu`` to be imported (which patches
+    ``torch.multiprocessing.reductions.reduce_tensor`` to support
+    NPU tensors via ``_share_npu_()`` / ``rebuild_npu_tensor``).
+    """
+
+    init_info_cls = NPUIPCWeightTransferInitInfo
+    update_info_cls = NPUIPCWeightTransferUpdateInfo
+
+    def __init__(
+        self,
+        config: WeightTransferConfig,
+        parallel_config: ParallelConfig,
+        model: torch.nn.Module | None = None,
+    ) -> None:
+        super().__init__(config, parallel_config, model)
+
+    def parse_update_info(self, update_dict: dict[str, Any]) -> NPUIPCWeightTransferUpdateInfo:
+        """Parse update dict, deserializing pickled IPC handles if present.
+
+        HTTP transport sends IPC handles as a base64-encoded pickle under the
+        key ``ipc_handles_pickled``. This method deserializes them back into
+        ``ipc_handles`` before constructing the typed dataclass, keeping
+        serialization concerns out of the dataclass itself.
+
+        Requires ``VLLM_ALLOW_INSECURE_SERIALIZATION=1`` because the
+        payload is deserialized via ``pickle.loads``.
+        """
+        if "ipc_handles_pickled" in update_dict:
+            if "ipc_handles" in update_dict:
+                raise ValueError("Cannot specify both `ipc_handles` and `ipc_handles_pickled`")
+
+            if not envs.VLLM_ALLOW_INSECURE_SERIALIZATION:
+                raise ValueError(
+                    "Refusing to deserialize `ipc_handles_pickled` without VLLM_ALLOW_INSECURE_SERIALIZATION=1"
+                )
+
+            pickled = update_dict.pop("ipc_handles_pickled")
+            update_dict["ipc_handles"] = pickle.loads(base64.b64decode(pickled))
+
+        return super().parse_update_info(update_dict)
+
+    def init_transfer_engine(self, init_info: NPUIPCWeightTransferInitInfo) -> None:
+        """No initialization needed for NPU IPC backend."""
+        pass
+
+    def receive_weights(
+        self,
+        update_info: NPUIPCWeightTransferUpdateInfo,
+        load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
+    ) -> None:
+        """Receive weights from the trainer via NPU IPC handles.
+
+        Args:
+            update_info: NPU IPC update info containing parameter names,
+                dtypes, shapes, and IPC handles.
+            load_weights: Callable that loads weights into the model.
+        """
+        device_index = torch.accelerator.current_device_index()
+        physical_npu_id = npu_generate_uuid()
+
+        if update_info.packed:
+            assert update_info.tensor_sizes is not None
+            assert isinstance(update_info.ipc_handles, dict)
+            weights = packed_npu_ipc_consumer(
+                ipc_handle=update_info.ipc_handles,
+                physical_npu_id=physical_npu_id,
+                names=update_info.names,
+                shapes=update_info.shapes,
+                dtype_names=update_info.dtype_names,
+                tensor_sizes=update_info.tensor_sizes,
+                device_index=device_index,
+            )
+            load_weights(weights)
+        else:
+            # Lazy import: ``rebuild_npu_tensor`` lives in ``torch_npu`` and
+            # must not be imported at module load time on non-NPU hosts.
+            from torch_npu.multiprocessing.reductions import rebuild_npu_tensor
+
+            assert isinstance(update_info.ipc_handles, list)
+            weights = []
+            for name, ipc_handle in zip(
+                update_info.names,
+                update_info.ipc_handles,
+            ):
+                if physical_npu_id not in ipc_handle:
+                    raise ValueError(
+                        f"IPC handle not found for NPU UUID {physical_npu_id}. "
+                        f"Available UUIDs: {list(ipc_handle.keys())}. "
+                        f"This may indicate that the trainer and worker are "
+                        f"not co-located on the same physical NPU (node)."
+                    )
+
+                args = ipc_handle[physical_npu_id]
+                list_args = list(args)
+                # Index 6 is the device_index parameter in torch's
+                # IPC handle tuple (rebuild_npu_tensor). Update it
+                # to the current device since the logical index can
+                # differ between sender and receiver.
+                list_args[6] = device_index
+                weight = rebuild_npu_tensor(*list_args)
+                weights.append((name, weight))
+
+            load_weights(weights)
+
+    def shutdown(self) -> None:
+        pass
+
+    @staticmethod
+    def trainer_send_weights(
+        iterator: Iterator[tuple[str, torch.Tensor]],
+        trainer_args: dict[str, Any] | NPUIPCTrainerSendWeightsArgs,
+    ) -> None:
+        """Send weights from trainer to inference workers via NPU IPC.
+
+        Supports two transport modes ('ray' and 'http') and two transfer
+        strategies:
+        - Non-packed (default): all weights in a single API call.
+        - Packed (packed=True): chunked transfer with bounded NPU memory.
+
+        For multi-NPU training, all ranks must call this method in
+        parallel. IPC handles are all-gathered across ranks and merged
+        so that each vLLM worker can find its own NPU UUID. Only rank 0
+        sends the payload to vLLM.
+
+        .. note::
+            This method calls ``update_weights`` internally. The caller must
+            handle ``pause`` / ``start_weight_update`` / ``finish_weight_update``
+            / ``resume`` before and after this method.
+
+        Args:
+            iterator: Iterator of (name, tensor) pairs.
+            trainer_args: NPUIPCTrainerSendWeightsArgs or equivalent dict.
+        """
+        args = NPUIPCTrainerSendWeightsArgs(**trainer_args) if isinstance(trainer_args, dict) else trainer_args
+        npu_uuid = npu_generate_uuid()
+        if args.packed:
+            NPUIPCWeightTransferEngine._send_packed(iterator, args, npu_uuid)
+        else:
+            NPUIPCWeightTransferEngine._send_unpacked(iterator, args, npu_uuid)
+
+    @staticmethod
+    def _is_rank_zero() -> bool:
+        """Return True if this is rank 0 or no distributed group exists."""
+        if not torch.distributed.is_initialized():
+            return True
+        return torch.distributed.get_rank() == 0
+
+    @staticmethod
+    def _all_gather_and_merge_handles(
+        handles: list[dict[str, tuple]],
+    ) -> list[dict[str, tuple]]:
+        """All-gather and merge IPC handle dicts across ranks.
+
+        Each rank contributes a list of ``{npu_uuid: ipc_args}`` dicts.
+        Rank 0 collects and merges per-index; other ranks receive a list
+        of empty dicts. No-op when no distributed group exists.
+        """
+        if not torch.distributed.is_initialized() or torch.distributed.get_world_size() == 1:
+            return handles
+
+        world_size = torch.distributed.get_world_size()
+        gathered: list[list[dict[str, tuple]] | None] = [None] * world_size
+        torch.distributed.all_gather_object(gathered, handles)
+        torch.distributed.barrier()
+        torch.npu.synchronize()
+
+        if torch.distributed.get_rank() == 0:
+            merged: list[dict[str, tuple]] = []
+            for param_idx in range(len(handles)):
+                m: dict[str, tuple] = {}
+                for rank_handles in gathered:
+                    if rank_handles is not None:
+                        m.update(rank_handles[param_idx])
+                merged.append(m)
+            return merged
+        return [{} for _ in handles]
+
+    @staticmethod
+    def _post_send_sync() -> None:
+        """Barrier + synchronize after a send; no-op if single-NPU."""
+        if torch.distributed.is_initialized() and torch.distributed.get_world_size() > 1:
+            torch.distributed.barrier()
+        torch.npu.synchronize()
+
+    @staticmethod
+    def _send_unpacked(
+        iterator: Iterator[tuple[str, torch.Tensor]],
+        args: NPUIPCTrainerSendWeightsArgs,
+        npu_uuid: str,
+    ) -> None:
+        """Send all weights in a single API call (non-packed mode)."""
+        names: list[str] = []
+        dtype_names: list[str] = []
+        shapes: list[list[int]] = []
+        ipc_handles: list[dict[str, tuple]] = []
+        # Hold strong refs to every contiguous copy until the send + post-send
+        # sync completes.  ``reduce_tensor``'s returned args do NOT keep
+        # storage alive.
+        weight_refs: list[torch.Tensor] = []
+
+        for name, tensor in iterator:
+            names.append(name)
+            dtype_names.append(str(tensor.dtype).split(".")[-1])
+            shapes.append(list(tensor.shape))
+
+            weight = tensor.detach().contiguous()
+            weight_refs.append(weight)
+            # Store only the rebuild args (drop the func); the consumer rebuilds
+            # with the well-known ``rebuild_npu_tensor``, mirroring upstream's
+            # CUDA IPC engine.
+            _, ipc_args = reduce_tensor(weight)
+            ipc_handles.append({npu_uuid: ipc_args})
+
+        ipc_handles = NPUIPCWeightTransferEngine._all_gather_and_merge_handles(ipc_handles)
+
+        if NPUIPCWeightTransferEngine._is_rank_zero():
+            NPUIPCWeightTransferEngine._do_send(
+                args=args,
+                names=names,
+                dtype_names=dtype_names,
+                shapes=shapes,
+                ipc_handles=ipc_handles,
+            )
+
+        NPUIPCWeightTransferEngine._post_send_sync()
+
+    @staticmethod
+    def _send_packed(
+        iterator: Iterator[tuple[str, torch.Tensor]],
+        args: NPUIPCTrainerSendWeightsArgs,
+        npu_uuid: str,
+    ) -> None:
+        """Send weights in bounded-memory chunks (packed mode)."""
+        post_iter_func: Callable = lambda item: item[1]
+
+        for chunk in packed_npu_ipc_producer(
+            iterator=iterator,
+            npu_uuid=npu_uuid,
+            post_iter_func=post_iter_func,
+            buffer_size_bytes=args.packed_buffer_size_bytes,
+        ):
+            ipc_handle = NPUIPCWeightTransferEngine._all_gather_and_merge_handles([chunk["ipc_handle"]])[0]
+
+            if NPUIPCWeightTransferEngine._is_rank_zero():
+                NPUIPCWeightTransferEngine._do_send(
+                    args=args,
+                    names=chunk["names"],
+                    dtype_names=chunk["dtype_names"],
+                    shapes=chunk["shapes"],
+                    ipc_handles=ipc_handle,
+                    tensor_sizes=chunk["tensor_sizes"],
+                    packed=True,
+                )
+
+            NPUIPCWeightTransferEngine._post_send_sync()
+
+    @staticmethod
+    def _do_send(
+        args: NPUIPCTrainerSendWeightsArgs,
+        names: list[str],
+        dtype_names: list[str],
+        shapes: list[list[int]],
+        ipc_handles: list[dict[str, tuple]] | dict[str, tuple],
+        tensor_sizes: list[int] | None = None,
+        packed: bool = False,
+    ) -> None:
+        """Send a single update payload via the configured transport."""
+        update_fields: dict[str, Any] = {
+            "names": names,
+            "dtype_names": dtype_names,
+            "shapes": shapes,
+            "packed": packed,
+        }
+        if tensor_sizes is not None:
+            update_fields["tensor_sizes"] = tensor_sizes
+
+        update_fields["ipc_handles"] = ipc_handles
+        update_info = NPUIPCWeightTransferUpdateInfo(**update_fields)
+
+        if callable(args.send_mode):
+            args.send_mode(update_info)
+        elif args.send_mode == "ray":
+            import ray
+
+            handles = args.llm_handle if isinstance(args.llm_handle, list) else [args.llm_handle]
+            ray.get([h.update_weights.remote(dict(update_info=asdict(update_info))) for h in handles])
+        elif args.send_mode == "http":
+            pickled_handles = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
+            http_fields = {k: v for k, v in update_fields.items() if k != "ipc_handles"}
+            http_fields["ipc_handles_pickled"] = pickled_handles
+
+            url = f"{args.url}/update_weights"
+            payload = {"update_info": http_fields}
+            response = requests.post(url, json=payload, timeout=300)
+            response.raise_for_status()

vllm_ascend/distributed/weight_transfer/packed_tensor.py

@@ -0,0 +1,325 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Packed tensor utilities for efficient weight transfer."""
+
+import math
+from collections.abc import Callable, Iterator
+from typing import Any
+
+import torch
+from torch.multiprocessing.reductions import reduce_tensor
+
+# Default values for packed tensor configuration.
+# These are imported by HCCLWeightTransferUpdateInfo and trainer_send_weights.
+DEFAULT_PACKED_BUFFER_SIZE_BYTES = 1024 * 1024 * 1024  # 1GB
+DEFAULT_PACKED_NUM_BUFFERS = 2
+
+
+def packed_broadcast_producer(
+    iterator: Iterator[tuple[str, torch.Tensor]],
+    group: Any,
+    src: int,
+    post_iter_func: Callable[[tuple[str, torch.Tensor]], torch.Tensor],
+    buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES,
+    num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS,
+) -> None:
+    """Broadcast tensors in a packed manner from trainer to workers.
+
+    Args:
+        iterator: Iterator of model parameters. Returns a tuple of (name, tensor)
+        group: Process group (PyHcclCommunicator)
+        src: Source rank (0 in current implementation)
+        post_iter_func: Function to apply to each (name, tensor) pair before
+                       packing, should return a tensor
+        buffer_size_bytes: Size in bytes for each packed tensor buffer.
+                          Both producer and consumer must use the same value.
+        num_buffers: Number of buffers for double/triple buffering.
+                    Both producer and consumer must use the same value.
+    """
+    target_packed_tensor_size = buffer_size_bytes
+
+    streams = [torch.npu.Stream() for _ in range(num_buffers)]
+    buffer_idx = 0
+
+    packing_tensor_list: list[list[torch.Tensor]] = [[] for _ in range(num_buffers)]
+    packing_tensor_sizes: list[int] = [0 for _ in range(num_buffers)]
+    packed_tensors: list[torch.Tensor] = [torch.empty(0, dtype=torch.uint8, device="npu") for _ in range(num_buffers)]
+
+    done = False
+    while not done:
+        # Synchronize the current stream (waits for previous
+        # iteration's work on this buffer to finish)
+        streams[buffer_idx].synchronize()
+        # Start tasks for the new buffer in a new stream
+        with torch.npu.stream(streams[buffer_idx]):
+            # Initialize the packing tensor list and sizes
+            packing_tensor_list[buffer_idx] = []
+            packing_tensor_sizes[buffer_idx] = 0
+            # Pack the tensors
+            while True:
+                try:
+                    item = next(iterator)
+                except StopIteration:
+                    done = True
+                    break
+                # Apply post processing and convert to linearized uint8 tensor
+                tensor = post_iter_func(item).contiguous().view(torch.uint8).view(-1)
+                packing_tensor_list[buffer_idx].append(tensor)
+                packing_tensor_sizes[buffer_idx] += tensor.numel()
+                if packing_tensor_sizes[buffer_idx] > target_packed_tensor_size:
+                    break
+            if len(packing_tensor_list[buffer_idx]) > 0:
+                # Pack the tensors
+                packed_tensors[buffer_idx] = torch.cat(packing_tensor_list[buffer_idx], dim=0)
+
+        if len(packing_tensor_list[buffer_idx]) == 0:
+            # No more tensors — nothing left to broadcast
+            break
+
+        # torch.cat runs on the custom stream.  Synchronize before
+        # broadcasting on the default stream so the packed data is ready.
+        streams[buffer_idx].synchronize()
+        group.broadcast(packed_tensors[buffer_idx], src=src)
+
+        # Move to the next buffer
+        buffer_idx = (buffer_idx + 1) % num_buffers
+
+    # Ensure the last broadcast on the default stream has completed
+    # before returning, so NPU tensor cleanup at exit doesn't hang.
+    torch.npu.current_stream().synchronize()
+
+
+def packed_broadcast_consumer(
+    iterator: Iterator[tuple[str, tuple[list[int], torch.dtype]]],
+    group: Any,
+    src: int,
+    post_unpack_func: Callable[[list[tuple[str, torch.Tensor]]], None],
+    buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES,
+    num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS,
+) -> None:
+    """Consume packed tensors and unpack them into a list of tensors.
+
+    Args:
+        iterator: Iterator of parameter metadata. Returns (name, (shape, dtype))
+        group: Process group (PyHcclCommunicator)
+        src: Source rank (0 in current implementation)
+        post_unpack_func: Function to apply to each list of (name, tensor) after
+                         unpacking
+        buffer_size_bytes: Size in bytes for each packed tensor buffer.
+                          Both producer and consumer must use the same value.
+        num_buffers: Number of buffers for double/triple buffering.
+                    Both producer and consumer must use the same value.
+    """
+
+    def unpack_tensor(
+        packed_tensor: torch.Tensor,
+        names: list[str],
+        shapes: list[list[int]],
+        dtypes: list[torch.dtype],
+        tensor_sizes: list[int],
+    ) -> list[tuple[str, torch.Tensor]]:
+        """Unpack a packed uint8 tensor into a list of typed tensors."""
+        unpacked_tensors = packed_tensor.split(tensor_sizes)
+        unpacked_list = [
+            (name, tensor.contiguous().view(dtype).view(*shape))
+            for name, shape, dtype, tensor in zip(names, shapes, dtypes, unpacked_tensors)
+        ]
+        return unpacked_list
+
+    target_packed_tensor_size = buffer_size_bytes
+
+    streams = [torch.npu.Stream() for _ in range(num_buffers)]
+    default_stream = torch.npu.current_stream()
+    buffer_idx = 0
+
+    packing_tensor_meta_data: list[list[tuple[str, list[int], torch.dtype, int]]] = [[] for _ in range(num_buffers)]
+    packing_tensor_sizes: list[int] = [0 for _ in range(num_buffers)]
+    packed_tensors: list[torch.Tensor] = [torch.empty(0, dtype=torch.uint8, device="npu") for _ in range(num_buffers)]
+
+    done = False
+    while not done:
+        # Synchronize the current stream (waits for previous
+        # iteration's load_weights on this buffer to finish)
+        streams[buffer_idx].synchronize()
+        with torch.npu.stream(streams[buffer_idx]):
+            # Collect parameter metadata for this buffer
+            packing_tensor_meta_data[buffer_idx] = []
+            packing_tensor_sizes[buffer_idx] = 0
+            while True:
+                try:
+                    name, (shape, dtype) = next(iterator)
+                except StopIteration:
+                    done = True
+                    break
+                tensor_size = math.prod(shape) * dtype.itemsize
+                packing_tensor_meta_data[buffer_idx].append((name, shape, dtype, tensor_size))
+                packing_tensor_sizes[buffer_idx] += tensor_size
+                if packing_tensor_sizes[buffer_idx] > target_packed_tensor_size:
+                    break
+            if len(packing_tensor_meta_data[buffer_idx]) > 0:
+                packed_tensors[buffer_idx] = torch.empty(
+                    packing_tensor_sizes[buffer_idx],
+                    dtype=torch.uint8,
+                    device="npu",
+                )
+
+        if len(packing_tensor_meta_data[buffer_idx]) == 0:
+            break
+
+        # Broadcast on the default stream.
+        group.broadcast(packed_tensors[buffer_idx], src=src)
+
+        # Synchronize the default stream so broadcast completes before
+        # load_weights (running on the custom stream) reads the data.
+        default_stream.synchronize()
+
+        # Unpack and load weights on the custom stream
+        with torch.npu.stream(streams[buffer_idx]):
+            names, shapes, dtypes, tensor_sizes = zip(*packing_tensor_meta_data[buffer_idx])
+            post_unpack_func(
+                unpack_tensor(
+                    packed_tensors[buffer_idx],
+                    list(names),
+                    list(shapes),
+                    list(dtypes),
+                    list(tensor_sizes),
+                )
+            )
+
+        # Move to the next buffer
+        buffer_idx = (buffer_idx + 1) % num_buffers
+
+    # Wait for all in-flight load_weights (on custom streams) to finish.
+    # Otherwise NPU tensor cleanup at exit may hang.
+    for s in streams:
+        s.synchronize()
+
+
+# ── NPU IPC packed transfer ────────────────────────────────────────────
+
+
+def packed_npu_ipc_producer(
+    iterator: Iterator[tuple[str, torch.Tensor]],
+    npu_uuid: str,
+    post_iter_func: Callable[[tuple[str, torch.Tensor]], torch.Tensor],
+    buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES,
+) -> Iterator[dict[str, Any]]:
+    """Pack tensors into a reusable NPU IPC buffer and yield chunks.
+
+    Allocates a single NPU buffer of ``buffer_size_bytes`` and registers
+    it for IPC once via ``reduce_tensor``.  Each chunk's packed data is
+    copied into this buffer before yielding, so only one IPC-shared
+    allocation exists for the lifetime of the transfer.
+
+    Args:
+        iterator: Iterator of (name, tensor) pairs.
+        npu_uuid: Physical NPU UUID string for this rank.
+        post_iter_func: Applied to each (name, tensor) before packing.
+        buffer_size_bytes: Exact capacity of the reusable IPC buffer.
+    """
+    ipc_buffer = torch.empty(buffer_size_bytes, dtype=torch.uint8, device="npu")
+    # Store only the rebuild args (drop the func); the consumer rebuilds with
+    # the well-known ``rebuild_npu_tensor``, mirroring upstream's CUDA IPC engine.
+    _, ipc_args = reduce_tensor(ipc_buffer)
+
+    names: list[str] = []
+    shapes: list[list[int]] = []
+    dtypes: list[torch.dtype] = []
+    tensor_sizes: list[int] = []
+    total_bytes = 0
+
+    for name, orig_tensor in iterator:
+        flat = post_iter_func((name, orig_tensor)).contiguous().view(torch.uint8).view(-1)
+
+        if flat.numel() > buffer_size_bytes:
+            raise ValueError(
+                f"Tensor '{name}' has size {flat.numel()} bytes, "
+                f"which exceeds buffer_size_bytes={buffer_size_bytes}. "
+                f"Increase buffer_size_bytes to at least {flat.numel()}."
+            )
+
+        if total_bytes and total_bytes + flat.numel() > buffer_size_bytes:
+            torch.npu.current_stream().synchronize()
+            yield {
+                "names": names,
+                "shapes": shapes,
+                "dtype_names": [str(d).split(".")[-1] for d in dtypes],
+                "tensor_sizes": tensor_sizes,
+                "ipc_handle": {npu_uuid: ipc_args},
+            }
+            names, shapes, dtypes, tensor_sizes = [], [], [], []
+            total_bytes = 0
+
+        ipc_buffer[total_bytes : total_bytes + flat.numel()].copy_(flat)
+        names.append(name)
+        shapes.append(list(orig_tensor.shape))
+        dtypes.append(orig_tensor.dtype)
+        tensor_sizes.append(flat.numel())
+        total_bytes += flat.numel()
+
+    if total_bytes:
+        torch.npu.current_stream().synchronize()
+        yield {
+            "names": names,
+            "shapes": shapes,
+            "dtype_names": [str(d).split(".")[-1] for d in dtypes],
+            "tensor_sizes": tensor_sizes,
+            "ipc_handle": {npu_uuid: ipc_args},
+        }
+
+
+def packed_npu_ipc_consumer(
+    ipc_handle: dict[str, tuple],
+    physical_npu_id: str,
+    names: list[str],
+    shapes: list[list[int]],
+    dtype_names: list[str],
+    tensor_sizes: list[int],
+    device_index: int,
+) -> list[tuple[str, torch.Tensor]]:
+    """Unpack a single packed IPC chunk into named tensors.
+
+    Reconstructs the packed buffer via the IPC handle, unpacks into
+    individual tensors, and clones each into independent storage before
+    returning.  The clone is required because the producer reuses one
+    IPC buffer across chunks.
+
+    Args:
+        ipc_handle: Mapping of NPU UUID to a ``rebuild_npu_tensor`` args tuple
+            from ``reduce_tensor``.
+        physical_npu_id: Physical NPU UUID string for the current process.
+        names: Parameter names in the packed buffer.
+        shapes: Parameter shapes.
+        dtype_names: Parameter dtype name strings (e.g. "float16").
+        tensor_sizes: Size in bytes of each parameter in the packed buffer.
+        device_index: Local NPU device index.
+    """
+    # Lazy import: ``rebuild_npu_tensor`` lives in ``torch_npu`` and must not be
+    # imported at module load time on non-NPU hosts.
+    from torch_npu.multiprocessing.reductions import rebuild_npu_tensor
+
+    if physical_npu_id not in ipc_handle:
+        raise ValueError(
+            f"IPC handle not found for NPU UUID {physical_npu_id}. Available UUIDs: {list(ipc_handle.keys())}"
+        )
+
+    args = ipc_handle[physical_npu_id]
+    list_args = list(args)
+    # Index 6 of the args from reduce_tensor is the device_index.
+    # Overwrite it with the receiver's device index.
+    list_args[6] = device_index
+    packed = rebuild_npu_tensor(*list_args)
+
+    content_size = sum(tensor_sizes)
+    packed = packed[:content_size]
+
+    dtypes = [getattr(torch, dn) for dn in dtype_names]
+    weights: list[tuple[str, torch.Tensor]] = []
+    offset = 0
+    for name, shape, dtype, size in zip(names, shapes, dtypes, tensor_sizes):
+        raw = packed[offset : offset + size]
+        tensor = raw.contiguous().view(dtype).view(*shape).clone()
+        weights.append((name, tensor))
+        offset += size
+
+    return weights