release v0.1.2

Author: mag1c-h

MANIFEST.in

@@ -2,6 +2,7 @@ include LICENSE
 include pyproject.toml
 include CMakeLists.txt
 include requirements.txt
+include setup.py
 
 recursive-include examples *
 recursive-include benchmarks *

docs/source/user-guide/pd-disaggregation/1p1d.md

@@ -13,14 +13,12 @@ For illustration purposes, let us take GPU as an example and assume the model us
 ### Run prefill server
 Prefiller Launch Command:
 ```bash
-export PYTHONHASHSEED=123456
 export CUDA_VISIBLE_DEVICES=0 
 vllm serve /home/models/Qwen2.5-7B-Instruct \
 --max-model-len 20000 \
 --tensor-parallel-size 1 \
 --gpu_memory_utilization 0.87 \
 --trust-remote-code \
---enforce-eager \
 --no-enable-prefix-caching \
 --port 7800 \
 --block-size 128 \
@@ -42,14 +40,12 @@ vllm serve /home/models/Qwen2.5-7B-Instruct \
 ### Run decode server
 Decoder Launch Command:
 ```bash
-export PYTHONHASHSEED=123456 
 export CUDA_VISIBLE_DEVICES=0 
 vllm serve /home/models/Qwen2.5-7B-Instruct \
 --max-model-len 20000 \
 --tensor-parallel-size 1 \
 --gpu_memory_utilization 0.87 \
 --trust-remote-code \
---enforce-eager \
 --no-enable-prefix-caching \
 --port 7801 \
 --block-size 128 \

docs/source/user-guide/pd-disaggregation/npgd.md

@@ -19,14 +19,12 @@ For illustration purposes, let us assume that the model used is Qwen2.5-7B-Instr
 ### Run prefill server
 Prefiller Launch Command:
 ```bash
-export PYTHONHASHSEED=123456
 export ASCEND_RT_VISIBLE_DEVICES=0
 vllm serve /home/models/Qwen2.5-7B-Instruct \
 --max-model-len 20000 \
 --tensor-parallel-size 1 \
 --gpu_memory_utilization 0.87 \
 --trust-remote-code \
---enforce-eager \
 --no-enable-prefix-caching \
 --port 7800 \
 --block-size 128 \
@@ -49,14 +47,12 @@ vllm serve /home/models/Qwen2.5-7B-Instruct \
 ### Run decode server
 Decoder Launch Command:
 ```bash
-export PYTHONHASHSEED=123456
 export CUDA_VISIBLE_DEVICES=0 
 vllm serve /home/models/Qwen2.5-7B-Instruct \
 --max-model-len 20000 \
 --tensor-parallel-size 1 \
 --gpu_memory_utilization 0.87 \
 --trust-remote-code \
---enforce-eager \
 --no-enable-prefix-caching \
 --port 7801 \
 --block-size 128 \

docs/source/user-guide/pd-disaggregation/xpyd.md

@@ -13,14 +13,12 @@ For illustration purposes, let us take GPU as an example and assume the model us
 ### Run prefill servers
 Prefiller1 Launch Command:
 ```bash
-export PYTHONHASHSEED=123456
 export CUDA_VISIBLE_DEVICES=0 
 vllm serve /home/models/Qwen2.5-7B-Instruct \
 --max-model-len 20000 \
 --tensor-parallel-size 1 \
 --gpu_memory_utilization 0.87 \
 --trust-remote-code \
---enforce-eager \
 --no-enable-prefix-caching \
 --port 7800 \
 --block-size 128 \
@@ -41,14 +39,12 @@ vllm serve /home/models/Qwen2.5-7B-Instruct \
 
 Prefiller2 Launch Command:
 ```bash
-export PYTHONHASHSEED=123456
 export CUDA_VISIBLE_DEVICES=1 
 vllm serve /home/models/Qwen2.5-7B-Instruct \
 --max-model-len 20000 \
 --tensor-parallel-size 1 \
 --gpu_memory_utilization 0.87 \
 --trust-remote-code \
---enforce-eager \
 --no-enable-prefix-caching \
 --port 7801 \
 --block-size 128 \

pyproject.toml

@@ -5,8 +5,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "uc-manager"
 authors = [{name = "UCM Team"}]
-license = "MIT"
-license-files = ["LICENSE"]
+license = { file="LICENSE" }
 readme = "README.md"
 description = "Persist and reuse KV Cache to speedup your LLM."
 requires-python = ">=3.10"

setup.py

@@ -23,9 +23,11 @@
 #
 
 import os
+import shutil
 import subprocess
 import sys
 import sysconfig
+import warnings
 from glob import glob
 
 import pybind11
@@ -34,6 +36,12 @@
 from setuptools import Extension, find_packages, setup
 from setuptools.command.build_ext import build_ext
 
+# Suppress warnings about packages absent from packages configuration
+# These are expected for C++ source directories, test directories, etc.
+warnings.filterwarnings(
+    "ignore", message=".*Package.*is absent from the `packages` configuration.*"
+)
+
 ROOT_DIR = os.path.abspath(os.path.dirname(__file__))
 PLATFORM = os.getenv("PLATFORM")
 
@@ -45,15 +53,11 @@ def _enable_sparse() -> bool:
 
 
 def _is_cuda() -> bool:
-    return PLATFORM == "cuda"
-
+    return PLATFORM == "cuda" or (hasattr(torch, "cuda") and torch.cuda.is_available())
 
-def _is_npu() -> bool:
-    return PLATFORM == "ascend"
 
-
-def _is_musa() -> bool:
-    return PLATFORM == "musa"
+def _is_maca() -> bool:
+    return PLATFORM == "maca"
 
 
 class CMakeExtension(Extension):
@@ -67,6 +71,8 @@ def run(self):
         for ext in self.extensions:
             self.build_cmake(ext)
 
+        self._copy_so_files_to_build_lib()
+
     def build_cmake(self, ext: CMakeExtension):
         build_dir = self.build_temp
         os.makedirs(build_dir, exist_ok=True)
@@ -93,15 +99,8 @@ def build_cmake(self, ext: CMakeExtension):
 
         if _is_cuda():
             cmake_args.append("-DRUNTIME_ENVIRONMENT=cuda")
-        elif _is_npu():
-            cmake_args.append("-DRUNTIME_ENVIRONMENT=ascend")
-        elif _is_musa():
-            cmake_args.append("-DRUNTIME_ENVIRONMENT=musa")
         else:
-            raise RuntimeError(
-                "No supported accelerator found. "
-                "Please ensure either CUDA/MUSA or NPU is available."
-            )
+            cmake_args.append("-DRUNTIME_ENVIRONMENT=ascend")
 
         if _enable_sparse():
             cmake_args.append("-DBUILD_UCM_SPARSE=ON")
@@ -119,33 +118,58 @@ def build_cmake(self, ext: CMakeExtension):
             cwd=build_dir,
         )
 
+    def _copy_so_files_to_build_lib(self):
+        """Copy .so files from source directories to build_lib for installation."""
+        if not hasattr(self, "build_lib") or not self.build_lib:
+            return
 
-def _get_packages():
-    """Discover Python packages, optionally filtering out sparse-related ones."""
-    packages = find_packages()
-    if not _enable_sparse():
-        packages = [pkg for pkg in packages if not pkg.startswith("ucm.sparse")]
-    return packages
+        packages = _get_packages()
+        copied_count = 0
 
+        for package in packages:
+            # Source directory where CMake outputs .so files
+            source_package_dir = os.path.join(ROOT_DIR, package.replace(".", os.sep))
 
-def _get_package_data_with_so(packages=None):
-    """Automatically discover all packages and include .so files."""
-    if packages is None:
-        packages = _get_packages()
-    package_data = {}
+            # Destination in build_lib
+            build_package_dir = os.path.join(
+                self.build_lib, package.replace(".", os.sep)
+            )
+
+            # Find all .so files in the source package directory
+            so_files = glob(os.path.join(source_package_dir, "*.so"))
+
+            if so_files:
+                # Ensure destination directory exists
+                os.makedirs(build_package_dir, exist_ok=True)
+
+                # Copy each .so file
+                for so_file in so_files:
+                    dest_file = os.path.join(
+                        build_package_dir, os.path.basename(so_file)
+                    )
+                    shutil.copy2(so_file, dest_file)
+                    copied_count += 1
+                    print(
+                        f"[INFO] Copied {os.path.basename(so_file)} to {build_package_dir}"
+                    )
+
+        if copied_count > 0:
+            print(f"[INFO] Successfully copied {copied_count} .so file(s) to build_lib")
+        else:
+            print(
+                "[WARNING] No .so files found to copy. Extensions may not have been built."
+            )
 
-    for package in packages:
-        # Convert package name to directory path
-        package_dir = os.path.join(ROOT_DIR, package.replace(".", os.sep))
 
-        # Check if this package directory contains .so files
-        so_files = glob(os.path.join(package_dir, "*.so"))
-        if so_files:
-            package_data[package] = ["*.so"]
-            print(f"[INFO] Including .so files for package: {package}")
+def _get_packages():
+    """Discover Python packages, optionally filtering out sparse-related ones."""
+    sparse_enabled = _enable_sparse()
+    exclude_patterns = []
+    if not sparse_enabled:
+        exclude_patterns.append("ucm.sparse*")
 
-    print(f"[INFO] Package data: {package_data}")
-    return package_data
+    packages = find_packages(exclude=exclude_patterns)
+    return packages
 
 
 ext_modules = []
@@ -155,13 +179,12 @@ def _get_package_data_with_so(packages=None):
 
 setup(
     name="uc-manager",
-    version="0.1.1",
+    version="0.1.2",
     description="Unified Cache Management",
     author="Unified Cache Team",
     packages=packages,
     python_requires=">=3.10",
     ext_modules=ext_modules,
     cmdclass={"build_ext": CMakeBuild},
-    package_data=_get_package_data_with_so(packages),
     zip_safe=False,
 )

ucm/integration/vllm/ucm_connector.py

@@ -16,7 +16,6 @@
 from vllm.distributed.parallel_state import get_tp_group, get_world_group
 from vllm.platforms import current_platform
 from vllm.v1.core.sched.output import SchedulerOutput
-from vllm.v1.request import Request
 
 from ucm.logger import init_logger
 from ucm.shared.metrics import ucmmonitor
@@ -29,6 +28,7 @@
     from vllm.attention.backends.abstract import AttentionMetadata
     from vllm.forward_context import ForwardContext
     from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.request import Request
 
 logger = init_logger(__name__)
 
@@ -178,11 +178,15 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
                 self.metrics_config,
             )
             self.monitor = ucmmonitor.StatsMonitor.get_instance()
-            self.synchronize = (
-                torch.cuda.synchronize
-                if current_platform.is_cuda_alike()
-                else torch.npu.synchronize
-            )
+
+        self.synchronize = (
+            torch.cuda.synchronize
+            if current_platform.is_cuda_alike()
+            else torch.npu.synchronize
+        )
+
+        # 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]:
         token_ids = request.all_token_ids
@@ -513,6 +517,9 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
             # TODO error handling
             if self.global_rank == 0 or not self.load_only_first_rank:
                 if self.store.wait(task) != 0:
+                    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])
@@ -552,7 +559,9 @@ def wait_for_save(self) -> None:
         # TODO support PP
         if (self.is_mla or self.is_dsa) and self.global_rank != 0:
             return
-        if self.metrics_config:
+        if self.metrics_config or current_platform.device_type == "npu":
+            # When use vllm_ascend, we should add synchronize here, otherwise accuracy problem will raise
+            # This has already been fixed in the latest main branch of vllm_ascend, so synchronize will no longer be needed in future versions.
             self.synchronize()
 
         metadata = self._get_connector_metadata()
@@ -626,6 +635,18 @@ def wait_for_save(self) -> None:
     def clear_connector_metadata(self) -> None:
         super().clear_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.
+        """
+        res = self._invalid_block_ids
+        self._invalid_block_ids = set()
+        return res
+
 
 class UCMLayerWiseConnector(UCMDirectConnector):
     """
@@ -866,3 +887,13 @@ def clear_connector_metadata(self) -> None:
         after the model execution.
         """
         self.connector.clear_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.connector.get_block_ids_with_load_errors()

ucm/shared/trans/CMakeLists.txt

@@ -1,6 +1,9 @@
 if(RUNTIME_ENVIRONMENT STREQUAL "ascend")
     add_subdirectory(ascend)
 endif()
+if(RUNTIME_ENVIRONMENT STREQUAL "maca")
+    add_subdirectory(maca)
+endif()
 if(RUNTIME_ENVIRONMENT STREQUAL "cuda")
     add_subdirectory(cuda)
 endif()