执行摘要

• 一句话：新增 Mooncake 后端外部链接器，支持跨节点 KV 加载
• 推荐动作：值得精读。本 PR 展示了如何将外部分布式存储接入 unified cache：异步后台线程 + Future 计数器的 layer-wise 加载协议、跨 rank 的 restorable 集合求交、以及 buffer 注册与键后缀对齐等设计，都是可复用的模式。建议重点看 MooncakeDirectLinker.load_layer_wise 与 MooncakeStore.batch_exists_v2 的改动，理解 TRAILING_PAGES 空洞语义为什么必须返回集合而不是单一前缀长度。

功能与动机

PR body 明确说明这是从主 PR #35687 拆出的第 4 个拆分项，目标是为 external linker 增加 Mooncake 后端。外部 linker 允许 KV 缓存通过外部存储（如 Mooncake 的分布式共享内存）跨节点共享，从而支撑多机推理、prefill/decode 分离等场景。已有 split 分别完成 linker 框架（#37091）、基础存储（#37098）与混合池组装（#37151），本 PR 补齐 Mooncake 这一具体后端实现。

实现拆解

1. 新增 MooncakeDirectLinker 类（python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_direct_linker.py，+414 行）：继承 UnifiedCacheLinker，实现统一接口 lookup/load/cancel_queued_load/start_layer_wise_loading。构造时通过 resolve_hybrid_device_pool_group 组装混合设备池组，计算 tp/pp/cp 拓扑与 storage suffix，把 storage 的 mem_pool_host、registered_pools、mla_suffix、mha_suffix 注入进去；随后 register_buffers 把各 pool 的 GPU buffer 注册进 Mooncake，最后启动 load/offload 两个 daemon 线程。
2. 实现 LayerWiseLoadCounter 异步加载协议：该类兼容 KV pool 的 layer wait hook。update_producer 为每个批次分配 num_layers 个 Future，load 线程在每层加载完成后调用 complete(layer) 置位，消费端 wait_until(threshold) 等待到指定层，最后层完成时清理批次；失败时 fail() 给所有未完成的 Future 设异常。配合 req_to_token_pool.register_layer_transfer_counter 注册（仅当存在 MAMBA 池时）。
3. 异步加载流程：load() 只把请求放入 pending_loads（允许 partial/missing KV），start_layer_wise_loading() 冻结 GC、record CUDA event 后把批次放入 load_queue；load_thread_func() 先同步 event，再逐层调用 batch_get_session_start、batch_get_into_multi_buffer_ranges 写缓冲，失败则 fail 计数器，完成则把请求列表放入 completed_loads 供调用方消费。
4. MooncakeStore 查询语义改造（python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py，+23/-8）：_get_hybrid_page_component_keys 新增 PoolName.KV 分支，生成 _{mla_suffix}_k 后缀；batch_exists_v2 不再只返回单一 final_pages，而是返回完整的 restorable 前缀集合（ALL_PAGES 为连续前缀，TRAILING_PAGES 为所有满足尾窗完整的前缀点），调用方需要跨 rank 求交；同时对不支持的 hit_policy 显式抛 ValueError。
5. 通用化缓冲 meta 打包：_batch_io_v2 中判定是否调用 _pack_multi_buffer_meta 的条件从 transfer.name == PoolName.DEEPSEEK_V4_C4 放宽为 len(ptr_list) != len(key_strs)，使多缓冲打包逻辑适用于任何 key/指针数量不匹配的 pool。
6. 测试与 CI 配套：本 PR 未直接新增测试文件，依赖已有 test_mooncake_direct_linker.py、test_unified_cache_linker.py 以及 hicache/radix_cache 系列测试；作者多次通过 /rerun-test、/rerun-group 重跑这些测试组，最终 PR Test (Extra) 通过。

关键文件：

• python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_direct_linker.py（模块 缓存后端；类别 source；类型 core-logic；符号 _storage_suffix, LayerWiseLoadCounter, MooncakeDirectLinker.init, MooncakeDirectLinker.register_buffers）: 新增 414 行的核心后端实现：MooncakeDirectLinker 继承 UnifiedCacheLinker，实现 buffer 注册、layer-wise 异步加载、后台双线程调度与 LayerWiseLoadCounter 计数协议，是本 PR 的主体。
• python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py（模块 KV 存储；类别 source；类型 core-logic；符号 MooncakeStore._get_hybrid_page_component_keys, MooncakeStore.batch_exists_v2, MooncakeStore._batch_io_v2）: 修改 MooncakeStore 的混合池键生成与存在性查询语义：支持 KV 池后缀、返回 restorable 前缀集合、通用化 buffer meta 打包，是 direct linker 正确工作的底层支撑。

关键符号：MooncakeDirectLinker.lookup, MooncakeDirectLinker.load, MooncakeDirectLinker.start_layer_wise_loading, MooncakeDirectLinker.load_thread_func, MooncakeDirectLinker.load_layer_wise, MooncakeDirectLinker.register_buffers, LayerWiseLoadCounter.update_producer, LayerWiseLoadCounter.complete, LayerWiseLoadCounter.fail, LayerWiseLoadCounter.wait_until, MooncakeStore.batch_exists_v2, MooncakeStore._batch_io_v2, MooncakeStore._get_hybrid_page_component_keys

关键源码片段

python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_direct_linker.py

新增 414 行的核心后端实现：MooncakeDirectLinker 继承 UnifiedCacheLinker，实现 buffer 注册、layer-wise 异步加载、后台双线程调度与 LayerWiseLoadCounter 计数协议，是本 PR 的主体。

class LayerWiseLoadCounter:
    """CPU 侧完成计数器，兼容 KV pool 的 layer wait hook。
​
    加载线程按批次调用 update_producer() 分配一组 Future（每层一个），
    每完成一层加载就 complete(layer) 置位；消费侧通过 wait_until(threshold)
    等待某层的数据就绪，从而实现解码逐层消费、加载提前推进的流水线。
    """
​
    def __init__(self, num_layers: int):
        self.num_layers = num_layers
        self.producer_index = -1
        self.consumer_index = -1
        self.futures: dict[int, list[Future]] = {}
​
    def update_producer(self) -> int:
        # 新批次：为 num_layers 层各分配一个 Future，返回本批次索引
        self.producer_index += 1
        self.futures[self.producer_index] = [Future() for _ in range(self.num_layers)]
        return self.producer_index
​
    def set_consumer(self, index: int) -> None:
        # 消费侧把当前批次索引告诉计数器，wait_until 据此取 Future
        self.consumer_index = index
​
    def complete(self, index: int, layer: int) -> None:
        # 加载线程在每层数据写入 GPU buffer 后调用，唤醒等待该层的解码
        self.futures[index][layer].set_result(None)
​
    def fail(self, index: int, error: BaseException) -> None:
        # 某层加载失败时，把该批次所有未完成 Future 置为异常，避免解码侧永久卡死
        for future in self.futures.get(index, ()):
            if not future.done():
                future.set_exception(error)
​
    def wait_until(self, threshold: int) -> None:
        index = self.consumer_index
        futures = self.futures.get(index)
        if futures is None:
            return
        try:
            futures[threshold].result()
        except BaseException as error:
            raise RuntimeError("Mooncake layer-wise KV load failed.") from error
        finally:
            # 最后层完成后清理整个批次的 Future，避免内存累积
            if threshold == self.num_layers - 1:
                self.futures.pop(index, None)
​
    def reset(self) -> None:
        self.producer_index = -1
        self.consumer_index = -1
        self.futures.clear()

    def load_thread_func(self) -> None:
        while True:
            task = self.load_queue.get()
            try:
                if task is None:
                    return
                counter_index, pending, ready_event = task
                try:
                    # 等待 GPU 侧事件同步后再开始 H2D 加载，保证引用 buffer 的前序 kernel 已完成
                    ready_event.synchronize()
                    self.load_layer_wise(counter_index, list(pending.values()))
                except BaseException as error:
                    # 批量失败：让所有等该批次的层都拿到异常，而不是挂起
                    self.layer_done_counter.fail(counter_index, error)
                    logger.exception("Mooncake layer-wise load batch failed")
                finally:
                    # 无论成败都把请求列表交回主线程，由调用方决定重试或清理
                    self.completed_loads.put(list(pending))
            finally:
                self.load_queue.task_done()

python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py

修改 MooncakeStore 的混合池键生成与存在性查询语义：支持 KV 池后缀、返回 restorable 前缀集合、通用化 buffer meta 打包，是 direct linker 正确工作的底层支撑。

    def batch_exists_v2(
        self,
        keys: List[str],
        pool_transfers: Optional[List[PoolTransfer]] = None,
        extra_info: Optional[HiCacheStorageExtraInfo] = None,
    ) -> PoolTransferResult:
        if self.mem_pool_host.kv_buffer is None:
            # 逻辑锚点：物理 KV 对象不存在时，可用前缀完全由 sidecar 对象决定
            kv_pages = len(keys)
        else:
            kv_pages = self.batch_exists(keys, extra_info)
​
        hit_count: dict = {PoolName.KV: kv_pages} if kv_pages else {}
        # 从所有可能的 KV 前缀出发，逐个 pool 剔除其无法服务的停止点。
        # 必须收集完整集合而非最大值：TRAILING_PAGES 池会留下“空洞”，
        # 调用方需要跨 rank 求交（见 PoolTransferResult）。
        restorable = list(range(1, kv_pages + 1))
​
        for transfer in pool_transfers or []:
            if not restorable:
                break
            component_keys, key_multiplier = self._get_hybrid_page_component_keys(
                keys, transfer
            )
            component_keys = self._tag_keys(component_keys)
            ex = self._batch_exist(component_keys)
            if key_multiplier > 0:
                # 每个逻辑页展开为 key_multiplier 个组件对象，全部存在才算该页可恢复
                page_exists = [
                    all(
                        r == 1
                        for r in ex[i * key_multiplier : (i + 1) * key_multiplier]
                    )
                    for i in range(kv_pages)
                ]
            else:
                page_exists = [False] * kv_pages
            boundary = 0
            pool_restorable = []
            if transfer.hit_policy == PoolHitPolicy.ALL_PAGES:
                try:
                    boundary = page_exists.index(False)
                except ValueError:
                    boundary = kv_pages
                pool_restorable = list(range(1, boundary + 1))
            elif transfer.hit_policy == PoolHitPolicy.TRAILING_PAGES:
                # 尾窗完整的前缀点都可作为停止点，扫描全部点而不是只取最长
                trailing = max(1, len(transfer.keys) if transfer.keys else 1)
                for prefix_len in range(kv_pages, 0, -1):
                    if all(
                        page_exists[i]
                        for i in range(max(0, prefix_len - trailing), prefix_len)
                    ):
                        pool_restorable.append(prefix_len)
                        if boundary == 0:
                            boundary = prefix_len
            else:
                raise ValueError(f"Unsupported pool hit policy: {transfer.hit_policy}")
            if boundary:
                hit_count[transfer.name] = boundary
            # 与当前已累积的可恢复集合求交，保留双方都支持的停止点
            pool_restorable_set = set(pool_restorable)
            restorable = [p for p in restorable if p in pool_restorable_set]
​
        final_pages = restorable[-1] if restorable else 0
        return PoolTransferResult(final_pages, hit_count, restorable)

评论区精华

该 PR 没有正式的 review 评论，有价值的讨论发生在代码内注释与 CI 重跑记录中：

• 代码注释明确记录了一个关键设计约束：“Suffix order must match get_page_buffer_meta() for one page, because Mooncake zips object keys with registered buffer pointers”——键后缀顺序必须与 buffer 元数据顺序严格一致，否则 Mooncake 会把对象键和缓冲指针错误对应。
• batch_exists_v2 的注释解释了为什么从单一 boundary 改为返回 restorable 集合：“a TRAILING_PAGES pool leaves holes (see PoolTransferResult), and the caller has to intersect these sets across ranks”——TRAILING_PAGES 池会留下可恢复前缀的空洞，调用方必须跨 rank 求交而不是取最小值。

• CI 中多次重跑 hicache、radix_cache/unified_radix_tree、unit/mem_cache 组测试，说明这些测试对运行环境（GPU 型号、Mooncake 依赖）比较敏感，存在偶发失败，最终靠 rerun 通过；AMD ROCm 7.2 测试仍未通过。

• CI 重跑与 Mooncake 后端测试稳定性 (testing): 依赖已有 hicache/radix_cache 测试覆盖新后端；Extra CI 通过，但环境相关的偶发失败需要关注。

• 本 PR 缺少新增测试文件 (testing): 测试覆盖依赖前序拆分 PR，作者通过反复 rerun 验证通过，但长期看建议补充 direct linker 的定向单测。

风险与影响

• 风险：
  1. 新增核心路径缺少本 PR 配套测试：mooncake_direct_linker.py 是 414 行全新核心代码，但 PR 内没有测试文件变更，验证依赖已有测试和多次重跑，存在覆盖盲区（如线程失败路径、跨 rank 求交）。
  2. 异步错误传播不完整：load_thread_func 中某批次失败只 fail 当前 counter_index 的 futures，但 finally 仍会把 pending 列表放入 completed_loads，调用方可能误以为加载成功；跨批次队列的异常状态没有统一处理。
  3. 启动期硬失败：register_buffers 中任何 buffer 注册出错都会抛 RuntimeError 导致服务启动失败，对 Mooncake 服务可用性有强依赖。
  4. 全局 GC 冻结副作用：freeze_gc_once() 会对整个进程调用 freeze_gc，虽然只执行一次，但会影响进程内其他模块的 GC 行为。
  5. 跨 rank 一致性依赖外部求交：batch_exists_v2 返回 restorable 集合后，最终正确性取决于调用方是否正确实现跨 rank 求交，若某个 rank 的集合与其他 rank 不一致，可能导致加载了不完整的前缀。
  6. 兼容性风险：batch_exists_v2 的返回语义从单一索引改为集合，现有其他调用方（如 3FS/file 后端）如果仍按旧语义消费 PoolTransferResult，可能出现兼容问题；_batch_io_v2 的打包条件放宽可能改变 DEEPSEEK_V4_C4 以外池的打包行为。
     - 影响：用户/系统层面：为统一缓存外部 linker 增加了 Mooncake 后端，使多节点场景下可以跨机器共享 KV 与混合状态（MLA/MHA/SWA/DRAFT/MAMBA），是 multi-node 统一缓存（HiCache）能力的关键拼图。影响范围集中在 mem_cache 模块，但 batch_exists_v2 语义变化会影响所有依赖它的存储后端。团队层面：这是系列拆分 PR 的第 4 个，后续还有更多后端/编排 PR 会基于此落地；代码中大量设计决策（suffix 顺序、restorable 集合、layer-wise 计数）会成为其他 linker 后端的参考模板。整体影响面中等偏大，但只涉及缓存子系统，不直接影响模型前向计算。
     - 风险标记：新增核心路径缺少配套测试, 线程异常传播不完整, 启动期强依赖 Mooncake 服务, 全局 GC 冻结副作用, 跨 rank 一致性依赖外部求交, batch_exists_v2 语义变更兼容风险

关联脉络

• PR #37307 fix(unified-memory): forward the KV-index translator through every wrapper backend: 同一 unified-memory 功能线，修复 wrapper 后端 KV 翻译器转发问题；本 PR 的 MooncakeStore 也依赖 registered_pools 与 KV 组件键的正确命名，两者共同保证混合池在多后端下的索引一致性。
• PR #35177 feat(unified-memory): three sub-pools for mamba + hybrid-SWA models: 为 unified 池引入多子池与浮动池，本 PR 的 MooncakeDirectLinker 正是消费这些混合 pool group 的外部 linker 后端，二者构成同一架构的上下游。
• PR #35158 feat(unified-memory): byte-budget sizing, feasibility floor, and a conservation verifier: 统一内存按字节预算定容并引入守恒校验，是 unified cache 的基础设施；本 PR 的 storage 层查询语义变化需要在后续与这些校验联动。
• PR #37339 [Fix] Use real ReqKvInfo in unit-test req mocks: 修复多个缓存相关单测的 KV mock，使 hicache/unified_radix_tree 等测试在新 linker 后端下更可靠；本 PR 的 CI 复用了这些测试。