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 覆盖范围；不要再次乘以压缩比。
• 此补丁将计算窗口接入 Ascend DSA、FA 和 SFA 注意力路径。请先在 eager 模式下开始验证。图模式执行和额外的注意力后端需要单独的集成验证。
• PP、DCP、PCP 以及 prefill/decode TP 不匹配仍保留上游 main Mooncake 逐层实现的限制。此分支不包含独立的 PP/DCP 适配。
• 明确拒绝循环 Mamba 状态。混合注意力和混合循环/线性注意力状态不能互换。
• 仅发布完整的、与协调器对齐的块快照。部分块卸载和跨层缓存缓冲区复用不属于此混合实现。

组感知对象与范围偏移¶

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

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

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

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

例如：

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

可达性与会话生命周期¶

调度器查询协调器按组查找掩码所选键的 batch_is_exist。只有当所有必需的存储头键都存在时，块才可用。然后协调器确定跨组的公共可达 Token 边界。如果相应的窗口或压缩器状态缺失，仅命中全注意力是不够的。

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

Worker 为每个组创建独立的请求视图。会话所有权以 (request_id, group_id, logical_block_index) 为索引，因此一个组中的块零不能替换另一个组中的块零。在分块 prefill 步骤之间，即使没有新的调度器加载规范，也可以使用当前本地块 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 等待该事件。Worker 记录保存事件，并在同一边界提交当前层的 put。然后主机启动注意力内核，并在返回到后续输出投影或 MoE 通信之前应用配置的队列策略。

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

在 DSA 中，该边界位于压缩器/索引器/缓存更新之后、稀疏注意力算子之前。在 SFA 中，它围绕稀疏注意力算子；在 FA 中，它围绕 paged/FIA 注意力路径。它不围绕整个 Transformer 层。

影响：

• 在默认策略下，传输尾部可能与后续的 HCCL 或专家并行通信竞争。应使用设备跟踪来验证该固定策略在目标部署上的表现。
• 初始的、未预取的按需加载没有前置注意力窗口。它会等待更早的计算流工作，然后在注意力开始前完成。
• 会话分配、存在性查询和其他元数据 RPC 不是 payload DMA，并且可以发生在窗口之外。
• 这是一个本地队列策略，而不是集群范围的带宽调度器。无关的 worker 和独立的通信流不会被全局串行化。使用设备跟踪验证多流和多 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 单元测试覆盖组感知字节往返、不同块大小、不等的缓存条目数、稀疏掩码、独立提交边界、组会话所有权、使用重映射本地块的延续、负传输结果、分配回滚、设备事件门控以及致命线程传播。它们不会执行注意力 kernel 或真实的 Mooncake 服务。

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

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

1. 真实混合模型上的冷/热 token 一致性、正向远程命中以及多个 chunk 边界。
2. 每个组（包括稀疏状态组和索引器组）都具有有效的加载范围。将 Mooncake range-debug 日志与设备跟踪一起检查。
3. Put/get payload 范围在缓存就绪事件之后开始，并在被测设备上的后续 HCCL 之前完成；检查所有相关流和 rank。
4. 在相同模型、TP、Prompt 和池条件下，相对于非逐层基线的 TTFT、吞吐量、暴露的传输尾部和 HCCL 持续时间。
5. 额外的并发请求、抢占、模型特定状态和服务故障测试。仅冒烟脚本不足以证明生产就绪。

真实 NPU/Mooncake 验证仍在等待中；实现环境仅支持 CPU/mock 检查。