Support Cache|Posix H2H pipeline for SGLang MLA models

ucm/integration/sglang/ucm_connector.py

@@ -21,6 +21,34 @@
 logger = logging.getLogger(__name__)
 
 
+def _page_first_kv_split_buffers(mem_pool_host: "HostKVCache") -> List[torch.Tensor]:
+    """Return the K/V Host buffers for a regular MLA split layout."""
+    buffers = [mem_pool_host.k_buffer, mem_pool_host.v_buffer]
+    for name, buffer in zip(("k", "v"), buffers):
+        if not buffer.is_contiguous():
+            raise ValueError(
+                f"page_first_kv_split {name}_buffer must be contiguous"
+            )
+    return buffers
+
+
+def _page_first_kv_split_tensor_sizes(
+    mem_pool_host: "HostKVCache",
+) -> List[int]:
+    page_num = int(mem_pool_host.page_num)
+    if page_num <= 0:
+        raise ValueError(f"invalid page_num for page_first_kv_split: {page_num}")
+    sizes = []
+    for buffer in _page_first_kv_split_buffers(mem_pool_host):
+        nbytes = int(buffer.numel()) * int(buffer.element_size())
+        if nbytes % page_num != 0:
+            raise ValueError(
+                f"buffer bytes {nbytes} are not divisible by page_num {page_num}"
+            )
+        sizes.append(nbytes // page_num)
+    return sizes
+
+
 def _load_extra_config_from_yaml_env() -> Optional[Dict[str, Any]]:
     cfg_path = os.environ.get("UNIFIEDCACHE_CONFIG_FILE")
     if not cfg_path:
@@ -71,6 +99,8 @@ def load_from_config(
         page_bytes = page_size * mem_pool_host.get_size_per_token()
         tensor_size = page_bytes if storage_config.is_mla_model else page_bytes // 2
         block_size = tensor_size * (1 if storage_config.is_mla_model else 2)
+        is_kv_split = mem_pool_host.layout == "page_first_kv_split"
+        use_cache_pipeline = storage_config.is_mla_model
 
         ucm_cfg = kvc.get("ucm_connector_config")
         name = kvc.get("ucm_connector_name")
@@ -85,12 +115,51 @@ def load_from_config(
             )
 
         cfg = dict(ucm_cfg)
-        cfg["store_pipeline"] = "Posix"
-        cfg["storage_backends"] = [
-            path for path in cfg["storage_backends"].split(":") if path
-        ]
+        if use_cache_pipeline:
+            tensor_sizes = (
+                _page_first_kv_split_tensor_sizes(mem_pool_host)
+                if is_kv_split
+                else [page_bytes]
+            )
+            payload_size = sum(tensor_sizes)
+            io_direct = bool(cfg.get("io_direct", False))
+            block_size = (
+                (payload_size + 4095) // 4096 * 4096
+                if io_direct
+                else payload_size
+            )
+            cfg["store_pipeline"] = "Cache|Posix"
+            cfg.pop("tensor_size", None)
+            cfg["tensor_size_list"] = tensor_sizes
+            cfg["cache_use_host_buffer"] = True
+            # Host pointers must never enter device SDMA or GDR paths.
+            cfg["cache_sdma_direct"] = False
+            cfg["use_gdr"] = False
+            cfg.pop("gpu_kv_buffer_addrs", None)
+            cfg.pop("gpu_kv_buffer_sizes", None)
+            safe_model_name = (
+                "-".join(storage_config.model_name.split("/"))
+                if storage_config.model_name
+                else "unknown-model"
+            )
+            if not cfg.get("unique_id"):
+                tensor_fingerprint = "-".join(str(size) for size in tensor_sizes)
+                cfg["unique_id"] = (
+                    f"sglang-{safe_model_name}-{mem_pool_host.layout}-"
+                    f"tp{storage_config.tp_size}-{tensor_fingerprint}"
+                )
+        else:
+            cfg["store_pipeline"] = "Posix"
+            cfg["tensor_size"] = tensor_size
+
+        storage_backends = cfg["storage_backends"]
+        if isinstance(storage_backends, str):
+            cfg["storage_backends"] = [
+                path for path in storage_backends.split(":") if path
+            ]
+        else:
+            cfg["storage_backends"] = list(storage_backends)
         cfg["device_id"] = get_world_group().local_rank
-        cfg["tensor_size"] = tensor_size
         cfg["shard_size"] = block_size
         cfg["block_size"] = block_size
         cfg["stream_number"] = 8
@@ -117,6 +186,26 @@ def __init__(
         self.cache_nums = 1 if self.is_mla else 2
         self.tp_rank = storage_config.tp_rank
         self.tp_size = storage_config.tp_size
+        self.is_kv_split = mem_pool_host.layout == "page_first_kv_split"
+        self.split_buffers = (
+            _page_first_kv_split_buffers(mem_pool_host)
+            if self.is_kv_split
+            else []
+        )
+        self.split_tensor_sizes = (
+            _page_first_kv_split_tensor_sizes(mem_pool_host)
+            if self.is_kv_split
+            else []
+        )
+        self.cache_tensor_sizes = (
+            self.split_tensor_sizes
+            if self.is_kv_split
+            else (
+                [self.page_size * mem_pool_host.get_size_per_token()]
+                if self.is_mla
+                else []
+            )
+        )
 
         self.config_suffix = self._build_config_suffix()
 
@@ -150,7 +239,13 @@ def _encode_keys(self, keys: List[str]) -> List[bytes]:
     def _build_config_suffix(self) -> str:
         model_name = "-".join(self.model.split("/")) if self.model else ""
         if self.is_mla:
-            return f"_{model_name}"
+            tensor_fingerprint = "-".join(
+                str(size) for size in self.cache_tensor_sizes
+            )
+            return (
+                f"_{model_name}_{self.mem_pool_host.layout}_p{self.page_size}_"
+                f"tp{self.tp_size}_{tensor_fingerprint}"
+            )
         return f"_{model_name}_{self.tp_rank}_{self.tp_size}"
 
     def _get_physical_key(self, logical_key: str) -> str:
@@ -168,6 +263,35 @@ def _generate_task(
             return [], [], []
 
         shard_index_list = [0] * len(encoded_keys)
+        if self.is_kv_split:
+            indices = host_indices.tolist()
+            if len(indices) % self.page_size != 0:
+                raise ValueError(
+                    "host_indices length must be a multiple of page_size"
+                )
+            if len(indices) // self.page_size != len(encoded_keys):
+                raise ValueError(
+                    "page count mismatch between keys and host_indices: "
+                    f"{len(encoded_keys)} != {len(indices) // self.page_size}"
+                )
+            ptr_list = []
+            for offset in range(0, len(indices), self.page_size):
+                first_token_index = int(indices[offset])
+                if first_token_index % self.page_size != 0:
+                    raise ValueError(
+                        "page_first_kv_split host index must start at a page boundary"
+                    )
+                page_index = first_token_index // self.page_size
+                ptr_list.append(
+                    [
+                        buffer.data_ptr() + page_index * tensor_size
+                        for buffer, tensor_size in zip(
+                            self.split_buffers, self.split_tensor_sizes
+                        )
+                    ]
+                )
+            return encoded_keys, shard_index_list, ptr_list
+
         ptr_list, _ = self.mem_pool_host.get_page_buffer_meta(host_indices)
 
         if not self.is_mla:

ucm/integration/sglang/unifiedcache_store.py

@@ -43,9 +43,11 @@ def _ensure_initialized(self) -> SglangUcmConnector:
 
     def register_mem_pool_host(self, mem_pool_host: HostKVCache):
         super().register_mem_pool_host(mem_pool_host)
-        if mem_pool_host.layout != "page_first":
+        supported_layouts = {"page_first", "page_first_kv_split", "page_first_direct"}
+        if mem_pool_host.layout not in supported_layouts:
             raise ValueError(
-                "UnifiedCacheStore currently requires --hicache-mem-layout page_first, "
+                "UnifiedCacheStore currently requires one of "
+                f"{sorted(supported_layouts)}, "
                 f"got {mem_pool_host.layout!r}."
             )
 

ucm/store/cache/cc/cache_store.cc

@@ -49,7 +49,8 @@ class CacheStore : public StoreV1 {
             UC_ERROR("Failed to check config params: {}.", s);
             return s;
         }
-        if (config.deviceId >= 0 && !config.gpuKvBufferAddrs.empty()) {
+        if (!config.cacheUseHostBuffer && config.deviceId >= 0 &&
+            !config.gpuKvBufferAddrs.empty()) {
             gpuKvBufferRegistrations_ = std::make_unique<Trans::GdrKVBufferConfig>();
             s = gpuKvBufferRegistrations_->Register(config.gpuKvBufferAddrs,
                                                     config.gpuKvBufferSizes);
@@ -161,6 +162,7 @@ class CacheStore : public StoreV1 {
         config.GetNumbers("gpu_kv_buffer_sizes", param.gpuKvBufferSizes);
         config.Get("use_gdr", param.useGdr);
         config.Get("cache_sdma_direct", param.cacheSdmaDirect);
+        config.Get("cache_use_host_buffer", param.cacheUseHostBuffer);
         config.GetNumber("local_rank_size", param.localRankSize);
         return param;
     }
@@ -184,6 +186,10 @@ class CacheStore : public StoreV1 {
         if (config.deviceId < -1) {
             return Status::InvalidParam("invalid device({})", config.deviceId);
         }
+        if (config.cacheUseHostBuffer && config.deviceId < 0) {
+            return Status::InvalidParam(
+                "cache_use_host_buffer requires a non-negative device_id as cache owner id");
+        }
         if (config.uniqueId.empty()) { return Status::InvalidParam("invalid unique id"); }
         auto s =
             Trans::GdrKVBufferConfig::Validate(config.gpuKvBufferAddrs, config.gpuKvBufferSizes);
@@ -196,6 +202,12 @@ class CacheStore : public StoreV1 {
         if (config.deviceId == -1) { return Status::OK(); }
         s = CheckSizeConfig(config);
         if (s.Failure()) { return s; }
+        if (config.cacheUseHostBuffer &&
+            (config.cacheSdmaDirect || config.useGdr || !config.gpuKvBufferAddrs.empty())) {
+            return Status::InvalidParam(
+                "cache_use_host_buffer is incompatible with cache_sdma_direct, use_gdr, "
+                "and gpu_kv_buffer_addrs");
+        }
 #if !UCM_RUNTIME_ASCEND_SDMA_DIRECT
         if (config.cacheSdmaDirect) {
             return Status::InvalidParam("Cache SDMA Direct requires RUNTIME_ENVIRONMENT=ascend-a3");
@@ -253,6 +265,7 @@ class CacheStore : public StoreV1 {
             UC_INFO("Set {}::StreamNumber to {}.", ns, config.EffectiveStreamNumber());
         }
         UC_INFO("Set {}::CacheSdmaDirect to {}.", ns, config.cacheSdmaDirect);
+        UC_INFO("Set {}::CacheUseHostBuffer to {}.", ns, config.cacheUseHostBuffer);
         UC_INFO("Set {}::LoadExclusiveBufferNumber to {}.", ns, config.loadExclusiveBufferNumber);
         UC_INFO("Set {}::GpuKvBufferNumber to {}.", ns, config.gpuKvBufferAddrs.size());
         UC_INFO("Set {}::UseGdr to {}.", ns, config.useGdr);

ucm/store/cache/cc/dump_queue.cc

@@ -24,6 +24,8 @@
 #include "dump_queue.h"
 #include <algorithm>
 #include <atomic>
+#include <cstddef>
+#include <cstring>
 #include <memory>
 #include "logger/logger.h"
 #include "metrics_api.h"
@@ -48,6 +50,7 @@ Status DumpQueue::Setup(const Config& config, TaskIdSet* failureSet, TransBuffer
     streamNumber_ = config.EffectiveStreamNumber();
     useGdr_ = config.useGdr;
     cacheSdmaDirect_ = config.cacheSdmaDirect;
+    useHostBuffer_ = config.cacheUseHostBuffer;
     cpuAffinityCores_ = config.cpuAffinityCores;
     waiting_.Setup(config.waitingQueueDepth);
     dumping_.Setup(config.runningQueueDepth);
@@ -76,8 +79,9 @@ void DumpQueue::DispatchStage(std::promise<Status>& started)
         UC_WARN("Failed({}) to set UCM dump dispatcher name.", nameStatus);
     }
     CopyStream stream;
-    auto s = cacheSdmaDirect_ ? stream.SetupSdmaDirect(deviceId_, useGdr_)
-                              : stream.Setup(deviceId_, streamNumber_, useGdr_);
+    auto s = useHostBuffer_ ? Status::OK()
+                            : (cacheSdmaDirect_ ? stream.SetupSdmaDirect(deviceId_, useGdr_)
+                                                : stream.Setup(deviceId_, streamNumber_, useGdr_));
     started.set_value(s);
     if (s.Failure()) [[unlikely]] { return; }
     if (!cpuAffinityCores_.empty()) {
@@ -115,6 +119,10 @@ Status DumpQueue::DumpOneTask(CopyStream& stream, TaskPtr task)
     dumpCtx.taskHandle = task->id;
     std::shared_ptr<std::atomic<double>> eventReadyTp;
     if (task->desc.prerequisiteHandle != 0) {
+        if (useHostBuffer_) {
+            return Status::InvalidParam(
+                "host-to-host dump does not accept a device prerequisite event");
+        }
         auto s = stream.WaitEvent(reinterpret_cast<void*>(task->desc.prerequisiteHandle));
         if (s.Failure()) [[unlikely]] {
             UC_ERROR("Failed({}) to wait prerequisite event for dump task({}).", s, task->id);
@@ -134,9 +142,11 @@ Status DumpQueue::DumpOneTask(CopyStream& stream, TaskPtr task)
         if (!handle.Owner()) { continue; }
         if (!handle.Ready()) {
             auto* host = cacheSdmaDirect_ ? handle.DeviceData() : handle.Data();
-            auto s = DeviceToHostAsync(stream, shard.addrs.data(), host);
+            auto s = useHostBuffer_ ? HostToHostGather(shard, host)
+                                    : DeviceToHostAsync(stream, shard.addrs.data(), host);
             if (s.Failure()) [[unlikely]] {
-                UC_ERROR("Failed({}) to do D2H for task({}).", s, task->id);
+                UC_ERROR("Failed({}) to do {} gather for task({}).", s,
+                         useHostBuffer_ ? "H2H" : "D2H", task->id);
                 UC::Metrics::UpdateStats(NAME_TO_METRIC_ID("cache_d2h_errors_total"), 1.0);
                 return s;
             }
@@ -152,7 +162,7 @@ Status DumpQueue::DumpOneTask(CopyStream& stream, TaskPtr task)
                              static_cast<double>(backendTaskDesc.size()));
     if (backendTaskDesc.empty()) { return Status::OK(); }
     auto tpSyncStart = NowTime::Now();
-    auto s = stream.Synchronize();
+    auto s = useHostBuffer_ ? Status::OK() : stream.Synchronize();
     if (s.Failure()) [[unlikely]] {
         UC_ERROR("Failed({}) to sync on stream for task({}).", s, task->id);
         UC::Metrics::UpdateStats(NAME_TO_METRIC_ID("cache_d2h_errors_total"), 1.0);
@@ -201,6 +211,25 @@ Status DumpQueue::DeviceToHostAsync(CopyStream& stream, void** device, void* hos
     return stream.DeviceToHostAsync(device, host, tensorSizes_);
 }
 
+Status DumpQueue::HostToHostGather(const Detail::Shard& shard, void* destination)
+{
+    if (shard.addrs.size() != tensorSizes_.size()) {
+        return Status::InvalidParam("invalid host addr number({}, expect {})", shard.addrs.size(),
+                                    tensorSizes_.size());
+    }
+    if (destination == nullptr) { return Status::InvalidParam("invalid null host destination"); }
+    auto* dst = static_cast<std::byte*>(destination);
+    size_t offset = 0;
+    for (size_t i = 0; i < tensorSizes_.size(); ++i) {
+        if (shard.addrs[i] == nullptr) {
+            return Status::InvalidParam("invalid null host source({})", i);
+        }
+        std::memcpy(dst + offset, shard.addrs[i], tensorSizes_[i]);
+        offset += tensorSizes_[i];
+    }
+    return Status::OK();
+}
+
 void DumpQueue::BackendDumpStage()
 {
     auto nameStatus = CpuAffinity::SetCurrentThreadName("ucm_dump_back");

ucm/store/cache/cc/dump_queue.h

@@ -59,6 +59,7 @@ class DumpQueue {
     size_t streamNumber_{1};
     bool useGdr_{false};
     bool cacheSdmaDirect_{false};
+    bool useHostBuffer_{false};
     std::vector<ssize_t> cpuAffinityCores_{};
     SpscRingQueue<TaskPair> waiting_;
     SpscRingQueue<DumpCtx> dumping_;
@@ -75,6 +76,7 @@ class DumpQueue {
     void DispatchOneTask(CopyStream& stream, TaskPair&& pair);
     Status DumpOneTask(CopyStream& stream, TaskPtr task);
     Status DeviceToHostAsync(CopyStream& stream, void** device, void* host);
+    Status HostToHostGather(const Detail::Shard& shard, void* destination);
     void BackendDumpStage();
 };
 

ucm/store/cache/cc/global_config.h

@@ -58,6 +58,7 @@ struct Config {
     std::vector<size_t> gpuKvBufferSizes{};
     bool useGdr{false};
     bool cacheSdmaDirect{UCM_RUNTIME_ASCEND_SDMA_DIRECT};
+    bool cacheUseHostBuffer{false};
     size_t localRankSize{8};
 
     size_t EffectiveStreamNumber() const noexcept { return cacheSdmaDirect ? 1 : streamNumber; }