Mooncake 分层混合注意力

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

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

适用范围与前提条件

• Use AscendStoreConnector with backend="mooncake" and use_layerwise=true.
• Install a Mooncake client with the range/session interfaces listed in the single-group guide.
• The model must provide valid KVCacheConfig.kv_cache_groups and group block tables. Runtime group block sizes already describe raw-token coverage; do not multiply them by the compression ratio again.
• This patch wires compute windows into the Ascend DSA, FA, and SFA attention paths. Start validation in eager mode. Graph-mode execution and additional attention backends require separate integration validation.
• Topology-matched PP supports uneven partitions and stage-local hybrid groups. DCP, PCP and prefill/decode TP mismatch remain unsupported.
• Recurrent/linear-attention groups using MambaSpec require mamba_cache_mode="align". Their state is saved after the attention kernels update it; they do not submit PUTs at attention entry.
• Only complete, coordinator-aligned block snapshots are published. Partial block offloading and cross-layer cache-buffer reuse are not part of this hybrid implementation.

组感知对象与 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

At attention entry, an NPU event protects cache writes and preceding work on the compute stream. Prefetched gets wait for that event. For full/SWA attention, the worker records the save event and submits current-layer puts at the same boundary. The host then launches the attention kernel and applies the configured queue policy before returning to the following output projection or MoE communication.

Physical layers containing MambaSpec (including GDN and Kimi KDA) update conv/recurrent state inside attention. They record the save event and submit PUTs from the existing post-compute save hook instead. That hook also applies send backpressure to attention paths without a compute-window context manager. Step completion waits for all queued PUTs, including when the final layer has no ranges to save.

默认情况下，未来层的 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 检查。