adapt store v1

Author: qyh111

ucm/integration/vllm/ucm_connector.py

@@ -4,7 +4,7 @@
 import pickle
 import time
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Callable, List, Optional
+from typing import TYPE_CHECKING, Callable, List, Optional, Tuple
 
 import torch
 from vllm.config import VllmConfig
@@ -20,8 +20,8 @@
 from ucm.logger import init_logger
 from ucm.shared.metrics import ucmmonitor
 from ucm.shared.metrics.observability import UCMStatsLogger
-from ucm.store.factory import UcmConnectorFactory
-from ucm.store.ucmstore import Task, UcmKVStoreBase
+from ucm.store.factory_v1 import UcmConnectorFactoryV1
+from ucm.store.ucmstore_v1 import Task, UcmKVStoreBaseV1
 from ucm.utils import Config
 
 if TYPE_CHECKING:
@@ -35,7 +35,7 @@
 
 @dataclass
 class RequestMeta:
-    ucm_block_ids: list[str] = field(default_factory=list)
+    ucm_block_ids: list[bytes] = field(default_factory=list)
     hbm_hit_block_num: int = 0
     # local_computed_block + external_computed_block
     total_hit_block_num: int = 0
@@ -47,9 +47,9 @@ class RequestMeta:
 @dataclass
 class RequestDispatchMeta:
     load_block_ids: tuple[
-        list[str], list[int]
+        list[bytes], list[int]
     ]  # [0] mean ucm_block_ids, [1] means vllm_block_ids
-    dump_block_ids: tuple[list[str], list[int]]
+    dump_block_ids: tuple[list[bytes], list[int]]
 
 
 @dataclass
@@ -69,14 +69,14 @@ def __init__(self, vllm_config, rank_id):
         if RequestHasher._SEED_HASH is None:
             RequestHasher._SEED_HASH = self("UCM_HASH_SEED")
 
-    def __call__(self, input_data) -> int:
-        if isinstance(input_data, str):
-            input_bytes = input_data.encode("utf-8")
+    def __call__(self, input_data) -> bytes:
+        if isinstance(input_data, bytes):
+            input_bytes = input_data
         else:
             input_bytes = pickle.dumps(input_data, protocol=pickle.HIGHEST_PROTOCOL)
 
         h = hashlib.md5(self.meta_bytes + input_bytes)
-        return int.from_bytes(h.digest(), byteorder="big")
+        return h.digest()
 
 
 class UCMDirectConnector(KVConnectorBase_V1):
@@ -110,9 +110,8 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
 
         if self.local_rank >= 0:
             self.device = torch_dev.device(f"{dev_name}:{self.local_rank}")
-            self._layer_offset_cache = {}
 
-        self.store: UcmKVStoreBase
+        self.store: UcmKVStoreBaseV1
 
         if role == KVConnectorRole.SCHEDULER:
             self.request_hasher = RequestHasher(vllm_config, 0)
@@ -160,7 +159,7 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
         config["io_size"] = block_size_per_layer * (
             1 if self.is_mla else num_head_per_tp
         )
-        self.store = UcmConnectorFactory.create_connector(name, config)
+        self.store = UcmConnectorFactoryV1.create_connector(name, config)
         self.block_data_size = config["kv_block_size"]
 
         logger.info("init UCConnectorImpl, connector: %s", name)
@@ -188,7 +187,7 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
         # invlalid block ids due to load errors
         self._invalid_block_ids: set[int] = set()
 
-    def generate_hash(self, block_size: int, request: "Request") -> list[str]:
+    def generate_hash(self, block_size: int, request: "Request") -> list[bytes]:
         token_ids = request.all_token_ids
 
         ret = []
@@ -205,7 +204,7 @@ def generate_hash(self, block_size: int, request: "Request") -> list[str]:
                 (parent_block_hash_value, block_token_ids_tuple)
             )
             parent_block_hash_value = hash_value
-            ret.append(str(hash_value))
+            ret.append(hash_value)
 
         return ret
 
@@ -395,70 +394,37 @@ def _extract_layer_index(layer_name: str) -> Optional[int]:
                 return int(chunk)
         return None
 
-    def _precompute_layer_offsets(self):
-        if not self.kv_caches:
-            return
-
-        sample_kv_layer = next(iter(self.kv_caches.values()))
-        elem_size = sample_kv_layer[0].element_size()
-        block_data_size = (
-            sample_kv_layer[0].numel() if self.is_mla else sample_kv_layer[0][0].numel()
-        ) * elem_size
-        layer_data_size = block_data_size if self.is_mla else block_data_size * 2
-
-        # precompute all layers offset
-        for layer_name, _ in self.kv_caches.items():
-            layer_id = self._extract_layer_index(layer_name)
-            assert layer_id is not None
-            k_offset = layer_data_size * layer_id
-            v_offset = k_offset + block_data_size if not self.is_mla else 0
-            self._layer_offset_cache[layer_name] = (k_offset, v_offset)
-
-    def _get_tensor_and_offset(
-        self, vllm_block_ids: list[int], kv_layer: torch.Tensor, layer_name: str
-    ) -> tuple[list[torch.Tensor], list[int]]:
+    def _get_tensors(
+        self, vllm_block_id: int
+    ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
         """
         GQA/MHA: one layer shape is (2, num_blocks, block_size, num_kv_heads, head_size)
         MLA: one layer shape is (num_blocks, block_size, head_size)
         """
-        k_tensors, k_offsets = [], []
-        v_tensors, v_offsets = [], []
-        k_offset, v_offset = self._layer_offset_cache[layer_name]
-
-        for vllm_block_id in vllm_block_ids:
+        k_tensors, v_tensors = [], []
+        for _, kv_layer in self.kv_caches.items():
             k_tensors.append(
                 kv_layer[vllm_block_id] if self.is_mla else kv_layer[0][vllm_block_id]
             )
-            k_offsets.append(k_offset)
             if not self.is_mla:
                 v_tensors.append(kv_layer[1][vllm_block_id])
-                v_offsets.append(v_offset)
-        return k_tensors + v_tensors, k_offsets + v_offsets
-
-    def _generate_task(self, vllm_block_ids: List[int], ucm_block_ids: List[str]):
-        if not self._layer_offset_cache:
-            self._precompute_layer_offsets()
-
-        num_layers = len(self.kv_caches)
-        num_blocks_per_layer = len(vllm_block_ids)
-        num_tensors_per_layer = num_blocks_per_layer * (1 if self.is_mla else 2)
-        dst_tensor_addr = [None] * (num_layers * num_tensors_per_layer)
-        ucm_offsets = [0] * (num_layers * num_tensors_per_layer)
-
-        idx = 0
-        for layer_name, one_layer_kv_cache in self.kv_caches.items():
-            tensors, offsets = self._get_tensor_and_offset(
-                vllm_block_ids, one_layer_kv_cache, layer_name
-            )
-            dst_tensor_addr[idx : idx + len(tensors)] = tensors
-            ucm_offsets[idx : idx + len(offsets)] = offsets
-            idx += len(tensors)
+        return k_tensors, v_tensors
 
-        repeat_times = len(self.kv_caches) * (1 if self.is_mla else 2)
-        ucm_total_block_ids = ucm_block_ids * repeat_times
+    def _generate_task(
+        self, vllm_block_ids: List[int], ucm_block_ids: List[bytes]
+    ) -> Tuple[List[bytes], List[int], List[List[torch.Tensor]]]:
+        """
+        GQA/MHA: one layer shape is (2, num_blocks, block_size, num_kv_heads, head_size)
+        MLA: one layer shape is (num_blocks, block_size, head_size)
+        """
+        block_ids, shard_indexs, tensors = [], [], []
+        for i, vllm_block_id in enumerate(vllm_block_ids):
+            k_tensors, v_tensors = self._get_tensors(vllm_block_id)
+            block_ids.append(ucm_block_ids[i])
+            tensors.append(k_tensors + v_tensors)
+            shard_indexs.append(0)
 
-        assert len(ucm_total_block_ids) == len(ucm_offsets) == len(dst_tensor_addr)
-        return ucm_total_block_ids, ucm_offsets, dst_tensor_addr
+        return block_ids, shard_indexs, tensors
 
     def _broadcast(self, dst_tensor_addr: list[torch.Tensor]):
         rec_tensor: torch.Tensor = None
@@ -501,26 +467,28 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
             ucm_block_ids, vllm_block_ids = request.load_block_ids
             if self.global_rank != 0 and not self.is_mla and not self.is_dsa:
                 for i, ucm_block_id in enumerate(ucm_block_ids):
-                    ucm_block_ids[i] = str(self.request_hasher(ucm_block_id))
-            ucm_total_block_ids, ucm_offsets, dst_tensor_addr = self._generate_task(
+                    ucm_block_ids[i] = self.request_hasher(ucm_block_id)
+            block_ids, shard_indexs, tensors = self._generate_task(
                 vllm_block_ids, ucm_block_ids
             )
             if self.global_rank == 0 or not self.load_only_first_rank:
                 request_to_task[request_id] = self.store.load(
-                    ucm_total_block_ids, ucm_offsets, dst_tensor_addr
+                    block_ids, shard_indexs, tensors
                 )
             else:
                 request_to_task[request_id] = None
-            req_broadcast_addr[request_id] = dst_tensor_addr
+            req_broadcast_addr[request_id] = [t for row in tensors for t in row]
 
         for request_id, task in request_to_task.items():
             # TODO error handling
             if self.global_rank == 0 or not self.load_only_first_rank:
-                if self.store.wait(task) != 0:
+                try:
+                    self.store.wait(task)
+                except RuntimeError as e:
+                    logger.error("request {request_id} load kv cache failed.:", e)
                     self._invalid_block_ids.update(
                         metadata.request_meta[request_id].load_block_ids[1]
                     )
-                    logger.error(f"request {request_id} load kv cache failed.")
             if self.load_only_first_rank:
                 self._broadcast(req_broadcast_addr[request_id])
         load_end_time = time.perf_counter() * 1000
@@ -568,7 +536,6 @@ def wait_for_save(self) -> None:
         assert isinstance(metadata, UCMConnectorMetadata)
 
         request_to_task: dict[str, Task] = {}
-        request_to_blocks: dict[str, list[str]] = {}
         is_save = False
         num_saved_block = 0
         num_saved_request = 0
@@ -583,36 +550,19 @@ def wait_for_save(self) -> None:
             ucm_block_ids, vllm_block_ids = request.dump_block_ids
             if self.global_rank != 0:
                 for i, ucm_block_id in enumerate(ucm_block_ids):
-                    ucm_block_ids[i] = str(self.request_hasher(ucm_block_id))
-            rets = self.store.create(ucm_block_ids)
-            end = 0
-            for i, ret in enumerate(rets):
-                if ret != 0:
-                    logger.error(
-                        f"create blocks for {request_id} failed, block index: {i}, ret code: {ret}"
-                    )
-                    break
-                end += 1
-
-            if end == 0:
-                continue
-            ucm_block_ids = ucm_block_ids[:end]
-            vllm_block_ids = vllm_block_ids[:end]
-            ucm_total_block_ids, ucm_offsets, dst_tensor_addr = self._generate_task(
+                    ucm_block_ids[i] = self.request_hasher(ucm_block_id)
+            block_ids, shard_indexs, tensors = self._generate_task(
                 vllm_block_ids, ucm_block_ids
             )
             request_to_task[request_id] = self.store.dump(
-                ucm_total_block_ids, ucm_offsets, dst_tensor_addr
+                block_ids, shard_indexs, tensors
             )
-            request_to_blocks[request_id] = ucm_block_ids
 
         for request_id, task in request_to_task.items():
-            ucm_block_ids = request_to_blocks[request_id]
-            if self.store.wait(task) == 0:
-                self.store.commit(ucm_block_ids, True)
-            else:
-                logger.error(f"request {request_id} dump kv cache failed.")
-                self.store.commit(ucm_block_ids, False)
+            try:
+                self.store.wait(task)
+            except RuntimeError as e:
+                logger.error("request {request_id} dump kv cache failed.:", e)
         save_end_time = time.perf_counter() * 1000
         save_speed = (
             num_saved_block

ucm/store/factory.py

@@ -63,9 +63,6 @@ def create_connector(cls, connector_name: str, config: dict) -> UcmKVStoreBase:
 UcmConnectorFactory.register_connector(
     "UcmNfsStore", "ucm.store.nfsstore.nfsstore_connector", "UcmNfsStore"
 )
-UcmConnectorFactory.register_connector(
-    "UcmPcStore", "ucm.store.pcstore.pcstore_connector", "UcmPcStore"
-)
 UcmConnectorFactory.register_connector(
     "UcmMooncakeStore",
     "ucm.store.mooncakestore.mooncake_connector",

ucm/store/factory_v1.py

@@ -0,0 +1,61 @@
+#
+# MIT License
+#
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+
+import importlib
+from typing import Callable
+
+from ucm.logger import init_logger
+from ucm.store.ucmstore_v1 import UcmKVStoreBaseV1
+
+logger = init_logger(__name__)
+
+
+class UcmConnectorFactoryV1:
+    _registry: dict[str, Callable[[], type[UcmKVStoreBaseV1]]] = {}
+
+    @classmethod
+    def register_connector(cls, name: str, module_path: str, class_name: str) -> None:
+        """Register a connector with a lazy-loading module and class name."""
+        if name in cls._registry:
+            raise ValueError(f"Connector '{name}' is already registered.")
+
+        def loader() -> type[UcmKVStoreBaseV1]:
+            module = importlib.import_module(module_path)
+            return getattr(module, class_name)
+
+        cls._registry[name] = loader
+
+    @classmethod
+    def create_connector(cls, connector_name: str, config: dict) -> UcmKVStoreBaseV1:
+        if connector_name in cls._registry:
+            connector_cls = cls._registry[connector_name]()
+        else:
+            raise ValueError(f"Unsupported connector type: {connector_name}")
+        assert issubclass(connector_cls, UcmKVStoreBaseV1)
+        logger.info("Creating connector with name: %s", connector_name)
+        return connector_cls(config)
+
+UcmConnectorFactoryV1.register_connector(
+    "UcmPcStore", "ucm.store.pcstore.pcstore_connector", "UcmPcStore"
+)

ucm/store/pcstore/cc/domain/trans/trans_task.h

@@ -58,10 +58,10 @@ class TransTask {
         : id{NextId()}, type{std::move(type)}, startTp{NowTp()}, brief_{std::move(brief)}
     {
     }
-    void Append(const std::string& block, const uintptr_t address)
+    void Append(const std::string& block, const std::vector<uintptr_t>& addresses)
     {
-        grouped_[block].push_back(address);
-        number_++;
+        grouped_[block] = addresses;
+        number_ += addresses.size();
     }
     auto Str() const noexcept { return fmt::format("{},{},{}", id, brief_, number_); }
     size_t GroupNumber() const { return grouped_.size(); }

ucm/store/pcstore/cpy/pcstore.py.cc

@@ -71,7 +71,9 @@ class PcStorePy : public PcStore {
         auto blockId = blockIds.begin();
         auto address = addresses.begin();
         while ((blockId != blockIds.end()) && (address != addresses.end())) {
-            task.Append(blockId->cast<std::string>(), address->cast<uintptr_t>());
+            std::string id = blockId->cast<py::bytes>();
+            std::vector<uintptr_t> addrs = address->cast<std::vector<uintptr_t>>();
+            task.Append(id, addrs);
             blockId++;
             address++;
         }