Mooncake|Posix

Author: UESTC-AHao

examples/ucm_mooncake_config.yaml

@@ -0,0 +1,38 @@
+# UCM + Mooncake Configuration Example
+#
+# Prerequisites:
+#   1. mooncake master_service running: mooncake_master --port 50088
+#   2. libmooncakestore.so built and installed
+#   3. NPU device available
+#
+# Usage:
+#   kv_connector_extra_config={"UCM_CONFIG_FILE": "/path/to/ucm_mooncake_config.yaml"}
+
+ucm_connectors:
+  - ucm_connector_name: "UcmPipelineStore"
+    ucm_connector_config:
+      store_pipeline: "Mooncake|Posix"
+      local_hostname: "127.0.0.1"
+      master_server_address: "127.0.0.1:50088"
+      metadata_server: "P2PHANDSHAKE"
+      protocol: "ascend"
+      global_segment_size: 32212254720
+      replica_num: 1
+      local_buffer_size: 1073741824
+      storage_backends: "/mnt/test"
+      io_direct: false
+      # device_name: ""                 # RDMA device, leave empty for ascend
+
+# prerequisiteHandle requires event sync to ensure NPU compute finishes
+# before Mooncake reads device memory for dump.
+enable_event_sync: true
+
+# Enable UCM metrics (optional)
+# metrics_config_path: "/workspace/unified-cache-management/examples/metrics/metrics_configs.yaml"
+
+# Sparse attention (optional)
+# ucm_sparse_config:
+#   GSAOnDevice: {}
+
+# Whether to use layerwise loading/saving (optional, default: True for UCMConnector)
+use_layerwise: true

ucm/integration/vllm/ucm_connector.py

@@ -382,6 +382,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         )
 
         self.store = self._create_store(self.kv_cache_layout, store_cores)
+        self._register_kv_cache_memory()
 
         if worker_cores:
             try:
@@ -393,6 +394,31 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         if self.device is None:
             raise RuntimeError(f"Unsupported device platform for UCMDirectConnector.")
 
+    def _register_kv_cache_memory(self):
+        for layer_name, kv_layer in self.kv_caches.items():
+            if isinstance(kv_layer, torch.Tensor):
+                if kv_layer.dim() == 5:
+                    num_blocks = kv_layer.shape[1]
+                    block_size = kv_layer[0].shape[1:].numel() * kv_layer.element_size()
+                    total_size = num_blocks * block_size
+                    self.store.register_memory(kv_layer[0].data_ptr(), total_size)
+                    self.store.register_memory(kv_layer[1].data_ptr(), total_size)
+                elif kv_layer.dim() == 3:
+                    num_blocks = kv_layer.shape[0]
+                    total_size = kv_layer.numel() * kv_layer.element_size()
+                    self.store.register_memory(kv_layer.data_ptr(), total_size)
+                else:
+                    raise ValueError(
+                        f"Unsupported kv cache tensor shape: {kv_layer.shape}"
+                    )
+            elif isinstance(kv_layer, Tuple):
+                for tensor in kv_layer:
+                    total_size = tensor.numel() * tensor.element_size()
+                    self.store.register_memory(tensor.data_ptr(), total_size)
+            else:
+                raise TypeError(f"Unsupported kv cache type: {type(kv_layer)}")
+        logger.info(f"Registered {len(self.kv_caches)} layers' kv cache memory")
+
     def get_num_new_matched_tokens(
         self,
         request: "Request",

ucm/shared/vendor/helper.cmake

@@ -15,7 +15,7 @@ function(find_reachable_git_url OUT_REACHABLE_URL IN_URL_LIST)
             COMMAND ${GIT_EXECUTABLE} ls-remote --heads "${GIT_URL}"
             RESULT_VARIABLE GIT_RESULT
             OUTPUT_QUIET ERROR_QUIET
-            TIMEOUT 5
+            TIMEOUT 20
         )
         if(GIT_RESULT EQUAL 0)
             set(${OUT_REACHABLE_URL} ${GIT_URL} PARENT_SCOPE)

ucm/store/cache/cc/cache_store.cc

@@ -99,6 +99,7 @@ class CacheStore : public StoreV1 {
         if (s.Failure()) [[unlikely]] { UC_ERROR("Failed({}) to wait task({}).", s, taskId); }
         return s;
     }
+    Status RegisterMemory(void* base_addr, size_t total_size) override { return Status::OK(); }
 
 private:
     Config ParseConfig(const Detail::Dictionary& config)

ucm/store/empty/cc/empty_store.cc

@@ -42,6 +42,7 @@ class EmptyStore : public StoreV1 {
     Expected<Detail::TaskHandle> Dump(Detail::TaskDesc task) { return NextId(); }
     Expected<bool> Check(Detail::TaskHandle taskId) { return true; }
     Status Wait(Detail::TaskHandle taskId) { return Status::OK(); }
+    Status RegisterMemory(void* base_addr, size_t total_size) { return Status::OK(); }
 
 private:
     static Detail::TaskHandle NextId() noexcept

ucm/store/fake/cc/fake_store.cc

@@ -82,6 +82,7 @@ class FakeStore : public StoreV1 {
     }
     Expected<bool> Check(Detail::TaskHandle taskId) override { return true; }
     Status Wait(Detail::TaskHandle taskId) override { return Status::OK(); }
+    Status RegisterMemory(void* base_addr, size_t total_size) override { return Status::OK(); }
 
 private:
     static Detail::TaskHandle NextId() noexcept

ucm/store/mooncakestore/CMakeLists.txt

@@ -0,0 +1,50 @@
+find_path(ASCEND_ACL_INCLUDE_DIR NAMES acl/acl.h PATHS /usr/local/Ascend/ascend-toolkit/latest/include NO_DEFAULT_PATH)
+find_library(ASCEND_ACL_LIBRARY NAMES ascendcl PATHS /usr/local/Ascend/ascend-toolkit/latest/lib64 NO_DEFAULT_PATH)
+find_library(MOONCAKE_STORE_LIBRARY NAMES mooncake_store PATHS /vllm-workspace/Mooncake/mooncake-store/lib NO_DEFAULT_PATH)
+
+if(ASCEND_ACL_INCLUDE_DIR AND ASCEND_ACL_LIBRARY AND MOONCAKE_STORE_LIBRARY)
+    file(GLOB_RECURSE SOURCES "./cc/*.cc")
+    add_library(mooncakestore SHARED ${SOURCES})
+
+    target_compile_features(mooncakestore PUBLIC cxx_std_20)
+    if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
+        target_compile_options(mooncakestore PRIVATE -fcoroutines)
+    endif()
+
+    target_compile_definitions(mooncakestore PRIVATE YLT_ENABLE_IBV SPDLOG_FMT_EXTERNAL)
+
+    if(NOT MOONCAKE_STORE_INCLUDE_DIR)
+        set(MOONCAKE_STORE_INCLUDE_DIR "/vllm-workspace/Mooncake/mooncake-store/include")
+    endif()
+
+    target_include_directories(mooncakestore PUBLIC
+        ${CMAKE_CURRENT_SOURCE_DIR}/cc
+        ${MOONCAKE_STORE_INCLUDE_DIR}
+        ${MOONCAKE_STORE_INCLUDE_DIR}/cachelib_memory_allocator
+        ${MOONCAKE_STORE_INCLUDE_DIR}/cachelib_memory_allocator/fake_include
+    )
+    target_include_directories(mooncakestore SYSTEM PUBLIC
+        ${MOONCAKE_STORE_INCLUDE_DIR}/cachelib_memory_allocator/include
+        ${ASCEND_ACL_INCLUDE_DIR}
+    )
+
+    target_link_libraries(mooncakestore PUBLIC
+        storeintf infra_logger fmt ${MOONCAKE_STORE_LIBRARY}
+        ${ASCEND_ACL_LIBRARY}
+    )
+
+    file(RELATIVE_PATH INSTALL_REL_PATH ${UCM_ROOT_DIR} ${CMAKE_CURRENT_SOURCE_DIR})
+    install(TARGETS mooncakestore LIBRARY DESTINATION ${INSTALL_REL_PATH} COMPONENT ucm)
+else()
+    message(STATUS "mooncakestore: Skipping build - required Ascend/Mooncake dependencies not found")
+    if(NOT ASCEND_ACL_INCLUDE_DIR)
+        message(STATUS "  * Missing: Ascend ACL headers (acl/acl.h)")
+    endif()
+    if(NOT ASCEND_ACL_LIBRARY)
+        message(STATUS "  * Missing: libascendcl.so library")
+    endif()
+    if(NOT MOONCAKE_STORE_LIBRARY)
+        message(STATUS "  * Missing: libmooncake_store.so library")
+    endif()
+    message(STATUS "  Please ensure Ascend and Mooncake dependencies are installed or build paths are correct")
+endif()

ucm/store/mooncakestore/cc/global_config.h

@@ -0,0 +1,67 @@
+/**
+ * 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.
+ * */
+#ifndef UNIFIEDCACHE_MOONCAKE_STORE_CC_GLOBAL_CONFIG_H
+#define UNIFIEDCACHE_MOONCAKE_STORE_CC_GLOBAL_CONFIG_H
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+namespace UC {
+class StoreV1;
+}
+
+namespace UC::MooncakeStore {
+
+struct Config {
+    std::string localHostname{};
+    std::string metadataServer{"P2PHANDSHAKE"};
+    std::string masterServerAddress{"127.0.0.1:50088"};
+    std::string protocol{"ascend"};
+    std::string deviceName{};
+
+    uint64_t globalSegmentSize{32212254720};
+    uint64_t localBufferSize{1073741824};
+    uint32_t replicaNum{1};
+    bool withSoftPin{false};
+
+    std::vector<uint64_t> tensorSizeList{};
+
+    int32_t deviceId{-1};
+    uint32_t loadWorkerNum{32};
+    uint32_t dumpWorkerNum{32};
+    uint32_t dumpGetBufferNum{4};
+    uint32_t dumpSubmitNum{4};
+    uint32_t dumpWaitNum{4};
+    uint32_t loadWaitNum{4};
+    uint32_t loadPutNum{4};
+    uint32_t loadGetNum{4};
+    uint32_t hostBufPoolSize{8192};
+
+    StoreV1* storeBackend{nullptr};
+};
+
+}  // namespace UC::MooncakeStore
+
+#endif

ucm/store/mooncakestore/cc/host_buffer_pool.h

@@ -0,0 +1,104 @@
+/**
+ * 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.
+ * */
+#ifndef UNIFIEDCACHE_MOONCAKE_STORE_CC_HOST_BUFFER_POOL_H
+#define UNIFIEDCACHE_MOONCAKE_STORE_CC_HOST_BUFFER_POOL_H
+
+#include <chrono>
+#include <condition_variable>
+#include <cstddef>
+#include <memory>
+#include <mutex>
+#include "thread/index_pool.h"
+
+namespace UC::MooncakeStore {
+
+class HostBufferPool {
+public:
+    HostBufferPool() = default;
+    ~HostBufferPool() = default;
+
+    void Setup(uint32_t count, size_t unitSize)
+    {
+        if (count == 0 || unitSize == 0) { return; }
+        size_t totalSize = static_cast<size_t>(count) * unitSize;
+        if (totalSize / unitSize != count) { return; }
+        unitSize_ = unitSize;
+        count_ = count;
+        pool_ = std::make_unique<char[]>(totalSize);
+        index_.Setup(count);
+    }
+
+    void* Acquire()
+    {
+        if (!pool_ || unitSize_ == 0) { return nullptr; }
+        auto idx = index_.Acquire();
+        if (idx == IndexPool::npos) { return nullptr; }
+        return pool_.get() + static_cast<size_t>(idx) * unitSize_;
+    }
+
+    void* AcquireWithTimeout(std::chrono::milliseconds timeout)
+    {
+        if (!pool_ || unitSize_ == 0) { return nullptr; }
+        auto idx = index_.Acquire();
+        if (idx != IndexPool::npos) { return pool_.get() + static_cast<size_t>(idx) * unitSize_; }
+        std::unique_lock<std::mutex> lk(cvMtx_);
+        auto deadline = std::chrono::steady_clock::now() + timeout;
+        while (true) {
+            cv_.wait_until(lk, deadline);
+            idx = index_.Acquire();
+            if (idx != IndexPool::npos) {
+                return pool_.get() + static_cast<size_t>(idx) * unitSize_;
+            }
+            if (std::chrono::steady_clock::now() >= deadline) { return nullptr; }
+        }
+    }
+
+    void Release(void* buf)
+    {
+        if (!buf || !pool_ || unitSize_ == 0) { return; }
+        auto offset = static_cast<char*>(buf) - pool_.get();
+        if (offset < 0 || static_cast<size_t>(offset) >= static_cast<size_t>(count_) * unitSize_) {
+            return;
+        }
+        if (static_cast<size_t>(offset) % unitSize_ != 0) { return; }
+        auto idx = static_cast<IndexPool::Index>(static_cast<size_t>(offset) / unitSize_);
+        index_.Release(idx);
+        cv_.notify_one();
+    }
+
+    size_t UnitSize() const { return unitSize_; }
+    uint32_t Count() const { return count_; }
+
+private:
+    std::unique_ptr<char[]> pool_;
+    size_t unitSize_{0};
+    uint32_t count_{0};
+    IndexPool index_;
+    std::mutex cvMtx_;
+    std::condition_variable cv_;
+};
+
+}  // namespace UC::MooncakeStore
+
+#endif