Mooncake 分层混合注意力

本扩展为 PR #16004 引入的 Mooncake 分层 range/session 实现增加了多 KV 组支持。 该实现基于上游 main 分支的 be427041bf63a620e4a637b60f2656e87dcdf8f6。

它面向 DeepSeek-V4 等注意力布局，其中滑动窗口 KV、 压缩 KV、indexer 缓存和 compressor 状态可以属于不同的 KV 组。它复用了 Memcache 也在使用的共享可达性协调器； 参见 issue #12234 中的设计讨论。 它不会将 Mooncake 操作转换为 Memcache GVA 操作。

适用范围与前提条件

• 使用 AscendStoreConnector，并设置 backend="mooncake" 和 use_layerwise=true。
• 安装带有 range/session 接口的 Mooncake 客户端，接口列表见 单组指南。
• 模型必须提供有效的 KVCacheConfig.kv_cache_groups 和组块表。 运行时组块大小已经描述了原始 token 覆盖范围；不要再 乘以压缩比。
• 本补丁将计算窗口接入昇腾 DSA、FA 和 SFA 注意力 路径。请在 eager 模式下开始验证。图模式执行和其他 注意力后端需要单独的集成验证。
• 拓扑匹配的 PP 支持不均匀分区和阶段本地混合组。 DCP、PCP 以及 prefill/decode TP 不匹配仍不支持。
• 使用 MambaSpec 的循环/线性注意力组需要 mamba_cache_mode="align"。其状态在注意力内核 更新后保存；它们不会在注意力入口提交 PUT。
• 仅发布完整的、协调器对齐的块快照。部分 块卸载和跨层缓存缓冲区复用不属于本 混合实现的一部分。

组感知对象与 range 偏移

单组线格式保持不变。混合布局使用单独的 命名空间：

model@mooncake_hybrid_v1:<layout-digest>@group:<id>@block:<size>@<hash>@<head>

该摘要包含 TP 大小、有序组成员关系以及每个组的页大小 签名。它与组 ID 一起标识缓存组，而无需 在对象键中重复调度器的缓存族分类。它 可防止不兼容的组布局读取相同的对象。它并非 模型权重校验和或租户隔离机制：对于模型名称和缓存规格相同 但权重不同的情况，请使用隔离部署。

每个键为单个组和存储 head/rank 存储一个块。其字节仅包含 该组注册的缓存条目，按物理层和缓存 名称排序。对象大小是它们实际每块字节长度的总和。传输 使用组本地层索引来计算远程偏移，而物理 层索引用于选择计算事件和完成信号。

例如：

组	原始 token/块	物理层	提交边界
Full attention	16	0, 2	Layer 2
Compressed KV	32	1, 2, 3	Layer 3
Window/state	16	0, 3	Layer 3

多个组可以在同一物理层传输。一个组在其自身的 最后一层之后提交，而不一定是模型的最后一层。同一物理层上的多个缓存 条目仍保持为独立的字节范围。

可达性与会话生命周期

调度器针对协调器按组查找掩码选中的键查询 batch_is_readable。 仅当所有必需的阶段和存储 head 键均已提交且可读时，块才可用。 随后协调器确定跨组的公共可达 token 边界。 当对应的窗口或 compressor 状态缺失时，仅命中全注意力 是不够的。

worker 对存储和加载掩码使用相同的协调器。它不会 将稀疏状态缓存解释为连续前缀。掩码生成失败 不会被静默转换为复制每个状态块的请求。

worker 为每个组创建单独的请求视图。会话所有权 以 (request_id, group_id, logical_block_index) 为索引，因此一个组中的块零 无法替换另一个组中的块零。在分块预填充步骤之间，即使没有新的 调度器加载规格，也可以使用当前本地块 ID 恢复已提交的 键。每个组保留独立的活跃传输 状态，并在其最后一个组层重置。

失败的 range put 会撤销受影响的键；其他组中成功的键 仍可提交。准备过程中的异常会撤销该准备过程中 先前打开的对象。失败的 get 会在使用不完整的混合 KV/状态之前停止前向路径。Get 会话在飞行中的读取完成后释放。

传输时序：可配置的异步流水线

重负载操作使用每层注意力计算窗口：

previous collectives / cache updates
              |
       cache-ready NPU event
              |
     +--------+---------------------+
     |                              |
 current attention kernel     put current-layer ranges
                              get future-layer ranges
     |                              |
     +------- policy checkpoint ----+
                         |
         output projection / MoE communication
              || bounded transfers

在注意力入口，一个 NPU 事件保护缓存写入以及计算流上 之前的工作。预取的 get 会等待该事件。对于 full/SWA 注意力， worker 记录保存事件并在同一边界提交当前层的 put。随后 host 启动 注意力内核，并在返回 到后续输出投影或 MoE 通信之前应用配置的队列策略。

包含 MambaSpec 的物理层（包括 GDN 和 Kimi KDA）在注意力内部 更新 conv/循环状态。它们改为从现有的计算后保存钩子 记录保存事件并提交 PUT。该钩子还会对没有计算窗口上下文管理器的 注意力路径施加发送背压。 步骤完成会等待所有排队的 PUT，包括最后一层 没有需要保存的范围时。

默认情况下，未来层的 get 保持排队，最多八个发送任务可能保持 未完成。这从每层的关键路径中移除了整队列排空， 但这也意味着传输可以与后续的输出投影或 MoE 通信重叠。这是内部调度策略，而非公开的 运行时设置。错误和拆除路径仍会完全排空两个队列。

在 DSA 中，该边界位于 compressor/indexer/cache 更新之后、稀疏注意力算子之前。在 SFA 中，它包围稀疏注意力算子；在 FA 中，它包围分页/FIA 注意力路径。它并不包围整个 transformer 层。

影响：

• 在默认策略下，传输尾部可能与后续的 HCCL 或专家并行通信产生竞争。应使用设备 trace 在目标部署上验证这一固定策略。
• 初始的、未预取的按需加载没有前置的注意力窗口。它会等待更早的计算流工作，然后在注意力开始前完成。
• 会话分配、存在性查询以及其他元数据 RPC 不属于 payload DMA，可以发生在窗口之外。
• 这是一种本地队列策略，而非集群范围的带宽调度器。不相关的工作进程和独立的通信流不会被全局串行化。应使用设备 trace 验证多流和多 rank 行为。
• 在没有硬件测量的情况下，不声称任何延迟或吞吐量改进。较短的注意力 kernel 可能暴露出显著的传输尾部。

配置

export MOONCAKE_CONFIG_PATH=/path/to/mooncake.json
vllm serve /path/to/hybrid-model \
  --tensor-parallel-size 8 \
  --enforce-eager \
  --enable-prefix-caching \
  --enable-chunked-prefill \
  --max-num-batched-tokens 4096 \
  --kv-transfer-config '{
    "kv_connector": "AscendStoreConnector",
    "kv_role": "kv_both",
    "kv_connector_extra_config": {
      "backend": "mooncake",
      "use_layerwise": true,
      "layerwise_prefetch_layers": 2,
      "layerwise_max_transfer_blocks": 64,
      "layerwise_max_transfer_bytes": 16777216
    }
  }'

根据实际硬件/模型选择 TP 大小、chunk 大小、量化方式和模型参数。从两个未来层预取窗口开始。本示例并非经过验证的 DeepSeek-V4 部署方案，也不构成内存容量保证。

验证

CPU/mock 单元测试覆盖了组感知的字节往返、不同的 block 大小、不相等的 cache 条目数量、稀疏掩码、独立的提交边界、组会话归属、使用重映射本地 block 的续传、传输失败结果、分配回滚、设备事件门控以及致命线程传播。它们不会执行注意力 kernel 或真实的 Mooncake 服务。

仍然需要真实的 NPU 和分布式 Mooncake 验证，本次变更并未将其自动化。请使用项目批准的部署和基准测试流程，针对目标模型和硬件进行验证。在其余设置完全相同且测试池隔离的情况下，对比分层和非分层 Mooncake，并在不启用 profiling 或 range-debug 日志的情况下运行性能测量。

在声称支持硬件或性能改进之前，请验证：

1. 在真实混合模型上验证冷/热 token 一致性、正向远程命中以及多个 chunk 边界。
2. 每个组（包括稀疏状态组和 indexer 组）都具有有效的加载范围。结合设备 trace 检查 Mooncake range-debug 日志。
3. 在受测设备上，put/get payload 范围在 cache-ready 事件之后开始，并在后续 HCCL 之前完成；检查所有相关的流和 rank。
4. 在相同模型、TP、prompt 和池条件下，对比非分层基线的 TTFT、吞吐量、暴露的传输尾部以及 HCCL 持续时间。
5. 额外的并发请求、抢占、模型特定状态以及服务故障测试。仅凭冒烟脚本无法确立生产就绪性。

真实的 NPU/Mooncake 验证尚待进行；当前实现环境仅支持 CPU/mock 检查。