Mooncake 分层适配与优化分析¶

关于后续的多组注意力扩展及其验证边界，请参阅 Mooncake 分层混合注意力。

1. 适配基线与结论¶

基线：Eric-dot/vllm-ascend:mooncake，本地提交 0a023b094e9e88ffaca0b1fda02529cef6277f8e。
参考实现：vllm-project/vllm-ascend#12418。
本地工作分支：eric-mooncake-pr12418-port。
当前状态：核心实现、CPU mock 单元测试和静态检查均已通过。在真实昇腾 NPU、Mooncake Master/Metadata 服务以及多 TP 进程上的端到端测试尚未完成。

本次适配并未直接复制早前的 PR。个人分支已包含更新后的 metadata.py、多缓存条目/MTP/SFA 布局、Memcache GVA 分层传输、异步线程异常传播以及层缓冲区复用。因此 Mooncake 路径复用了这些较新的结构，并新增了独立的以 key 为主的 range/session 协议。

2. 为什么远端对象数量从 `block × layer × rank` 降至 `block × rank`¶

早前的逐层 key 方案将每一层视为独立的远端对象：

model@block_hash@layer_0@rank_0
model@block_hash@layer_1@rank_0
...

对于 B 个 block、L 层以及实际存储数据的 R 个 rank，对象数量约为：

B × L × R

Mooncake 分层传输为每个 block/rank 对仅创建一个对象：

model@block_hash@rank_0

所有层按照实际缓存布局连续放置在对象内部：

object(block, rank)
├── layer 0: [cache entry 0][cache entry 1]...
├── layer 1: [cache entry 0][cache entry 1]...
└── ...

每一层计算完成后，仅将该层的字节范围写入同一对象。分层传输改变了传输时机和范围；它并不要求每一层都是独立对象。因此对象数量变为：

B × R

这减少了 Mooncake 元数据对象、key、session 以及 exists 查询项的数量，但并未减少 KV 负载的总大小。每一层的 KV 数据仍必须存储。

3. 当前实现中的数据流与控制流¶

3.1 保存¶

调度器针对每个 block，使用所有存储 rank 的 key 调用 batch_is_exist。
对于缺失的 key，worker 调用 batch_put_session_start(keys, object_sizes, ReplicateConfig)。
LayerBatchBuilder 根据实际缓存条目布局计算以下值：
本地 NPU 缓冲区地址。
该层中每个缓存条目的大小。
该层在远端全层对象中的目标偏移量。
每个注意力层完成后，发送线程调用 batch_put_from_multi_buffer_ranges。
如果某个 key 的范围写入失败，仅撤销该 key；其余 key 继续处理后续层。
最后一层成功后，worker 调用 batch_put_session_end。对于提交失败的 key，它会调用 batch_put_session_revoke。
只有成功提交的 key 才会进入可读集合，供后续 chunked-prefill 步骤使用。

3.2 加载¶

调度器仅将从 block 0 开始、且每个存储 rank 均为 COMPLETE 的连续前缀计为缓存命中。
worker 对去重后的远端 key 调用 batch_get_session_start。
每一层计算前，接收线程调用 batch_get_into_multi_buffer_ranges 将该层的范围加载到本地 block。
某一行失败时会记录对应的本地 block ID 并将其返回给调度器以重新计算，从而防止消费不完整的 KV 数据。
worker 在最后一层之后或请求终止时调用 batch_get_session_end。异常路径根据重试或终止语义释放 owner。

3.3 Chunked Prefill¶

MooncakeSessionTracker 维护三种关系：

尚未提交的 put key 及其 request/block owner。
同一请求的后续 chunk 可以加载的已提交 key。
打开的 get-session key 及其 request owner。

它保证：

对象在提交之前不会被视为可读。
后续 chunk 会续租并逐层恢复先前已提交的前缀，即使它没有新的 load_spec。
重试会释放 get session，同时保留后续重试所需的 key/block 关系。
终止或抢占事件会清除请求状态。
当多个请求共享同一远端 key 时，仅当最后一个 owner 释放后才会执行 get-end。

4. 相对于原始 PR 的适配与修正¶

4.1 使用实际布局偏移量而非固定的 `layer_id × page_size`¶

当前分支支持一个物理层中的多个缓存条目以及 MTP/SFA 布局。远端偏移量根据 group_layer_cache_entry_offsets 以及实际 block_len 值的前缀和计算：

layer_object_offset = sum(block_len before this layer)
entry_offset = layer_object_offset + prefix_sum(entry sizes in this layer)

因此，远端对象大小等于当前 rank 上所有层的所有缓存条目的总字节数。该实现不再假设每一层具有相同大小，也不再错误地将大小再次乘以 num_layers。

4.2 PutStart 继承 Mooncake 放置策略¶

原始 PR 中的 session-start 路径未传递 ReplicateConfig。当前实现与整 key put 路径保持一致，并传递：

preferred_segment。
prefer_alloc_in_same_node。

这防止分层传输绕过现有的本地优先和同节点放置策略。

4.3 针对 Range 调用的传输限流¶

layerwise_max_transfer_blocks 限制单次 range API 调用中的 key/block 行数。
layerwise_max_transfer_bytes 将过大的连续段拆分为更小的 range。

PutStart、GetStart、GetEnd 以及调度器的 exists 查询也会根据 block 限制进行批处理。这可以避免 Python/C++ 参数列表过大，以及长 prompt 场景下元数据的瞬时峰值。

4.4 启动期间快速失败¶

Mooncake layerwise 传输会在启动期间检查所有必需的 session 和 range 方法。如果缺少某个接口，启动会立即失败并提示客户端必须包含 Mooncake PR #2881，而不是在首次请求时于异步线程中失败。

最初的单组实现仅使用 model、block hash 和 TP/head rank。混合注意力与拓扑匹配的 PP 请参见混合指南。该实现仍然明确拒绝：

Prefill 或 decode 的 context parallel size 大于 1。
混合 layerwise 模式下的 Recurrent Mamba 状态，以及投影缓存组为空的 PP stage。
TP 与 layerwise 传输不匹配。

4.5 异常与回退¶

批处理结果必须与 key 一一对应。缺失项、布尔值以及非整数结果都会被拒绝。
range 异常会在唤醒计算线程之前发布 invalid-block 和 abort 状态。
put 异常会撤销 PROCESSING 对象。
get 异常只有在确认 range 调用已退出后才会结束 session。
仅作为 consumer 的 worker 会在 load hook 中推进 layer 游标，而不是依赖一个永远不会被调用的 save hook。

5. 运行该特性¶

5.1 Mooncake 版本¶

安装提供以下方法的 Mooncake Python 客户端：

batch_put_session_start
batch_put_from_multi_buffer_ranges
batch_put_session_end
batch_put_session_revoke
batch_get_session_start
batch_get_into_multi_buffer_ranges
batch_get_session_end

如果当前发布的 wheel 未提供这些接口，请在 PR #2881 合并后从源码构建 Mooncake。这些接口会在启动期间自动检查。

5.2 配置示例¶

export MOONCAKE_CONFIG_PATH=/path/to/mooncake.json

python -m vllm.entrypoints.openai.api_server \
  --model /path/to/model \
  --tensor-parallel-size 2 \
  --enforce-eager \
  --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=1、layerwise_prefetch_layers=1 并禁用 hybrid/CP/PP 开始。然后逐步增大 TP size、block 数量和预取深度。

在集成测试期间，可临时设置 VLLM_ASCEND_KVPOOL_RANGE_DEBUG=1，以输出 whole-key 操作、逐层 range 和 commit 的 JSON 审计日志。正常运行时请保持默认值 0，以避免逐层日志开销。

6. 后续优化建议¶

P0：生产环境前必须完成¶

固定最低 Mooncake commit 或版本
检查方法可用性可以防止使用过时的客户端，但无法检测不兼容的 ABI 或返回码语义。
在安装文档和 CI 镜像中固定包含 PR #2881 的确切 commit，或包含该 PR 的首个官方发布版本。
在 key 命名空间中添加 schema/layout 指纹
为兼容原始 PR，当前实现仍使用模型 basename。
具有相同 basename 的不同 revision、数据类型、block size 或 KV layout 可能发生冲突。
将 tenant、模型 revision、数据类型、block size、TP layout 和 schema 版本的稳定摘要添加到 key 中。
添加真实 NPU 端到端测试与故障注入
至少覆盖 TP=½、kv_both、P/D、chunked prefill、请求抢占、单 key range 失败、 commit 失败以及 Master 重启。
验证加载 KV 数据产生的 logits 和 tokens 与本地计算完全一致。

P1：高影响力性能改进¶

自适应预取深度
固定的 layerwise_prefetch_layers 值无法适应不同的逐层计算时间或网络抖动。
使用近期的 transfer_time / compute_time 比值、队列深度和可用 buffer 动态控制窗口。
滑动 session 窗口
当前实现会一次性为 batch 中的每个 block 打开 session。长 prompt 仍可能创建大量同时活跃的 lease 和 session。
仅为未来的 layer 或 block 窗口打开 session，然后在完成后推进窗口，以减少 Master 状态和超时压力。
减少逐层 Python 列表构建
当前实现仍会为每一层创建 all_buffers、all_sizes 和 all_offsets。
为每一层预计算 range 模板，仅应用向量化的 block-ID 地址偏移。这些描述符最终可以缓存在 C++/pybind 层中。
增量式调度器命中查询
当前实现会查询每一个候选 block，然后才定位到第一个未命中。
按窗口查询，并在第一个不完整的 block 处停止，从而大幅减少低命中率长 prompt 场景下的 Master RPC 和 key 数量。
基于总字节数和后端反馈的动态批处理
当前的 block 和 segment-byte 限制是静态的。
联合限制每次调用的总 range 数和字节数，并根据队列延迟、返回码和带宽自动调整。
共享 key 的本地扇出
当多个请求命中同一个远程 block 时，一个 key 目前可能产生多个远程读取行。
先读取到共享的 staging buffer，再在本地复制到多个目标 block。这是否有帮助取决于远程带宽相对于本地 H2D 带宽的情况。

P2：进一步演进¶

通过在 key 和对象头中记录 group 布局，并为每个 group 维护独立的完整性位图，支持混合布局和多 group 布局。
通过在 key 中编码并行坐标，并由调度器校验参与的 rank 集合，支持 PP/PCP/DCP。
支持逐层就绪位图，使消费方可以在整个对象变为 COMPLETE 之前读取已完成的早期层。这需要 Mooncake 提供可见性和一致性支持，并会引入相当大的复杂度。
在对象头中记录 schema、版本和校验和，并在加载前执行低开销的兼容性检查，以防止出现静默的布局不匹配。

7. 验证记录¶

py_compile：核心变更文件均通过。
Ruff lint 和格式检查：通过。
相关 CPU mock pytest 测试套件：285 passed, 106 subtests passed。
VLLM_ASCEND_KVPOOL_RANGE_DEBUG 的严格 0/1 校验：通过。
尚未完成：使用真实 NPU、真实 Mooncake 客户端和 Master、多节点网络的测试，以及性能基准测试。