Mooncake 逐层适配与优化分析¶

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

1. 适配基线与结论¶

基线：Eric-dot/vllm-ascend:mooncake，本地 commit 为 0a023b094e9e88ffaca0b1fda02529cef6277f8e。
参考实现：vllm-project/vllm-ascend#12418。
本地工作分支：eric-mooncake-pr12418-port。
当前状态：核心实现、CPU 模拟单元测试和静态检查已通过。在真实 Ascend 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 个块、L 层和 R 个 rank，近似对象数量为：

B × L × R

Mooncake 逐层传输仅为每个块/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 保存¶

调度器使用每个块的所有存储 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 会进入后续分块 Prefill 步骤的可读集合。

3.2 加载¶

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

3.3 分块 Prefill¶

MooncakeSessionTracker 维护三类关系：

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

它保证：

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

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 范围调用的传输限流¶

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

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

4.4 启动时快速失败¶

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

当前 key schema 仅编码模型、块哈希和 TP/head rank，因此实现会显式拒绝：

流水线并行大小大于 1。
Prefill 或 decode 的上下文并行大小大于 1。
混合或多组 KV 缓存布局。
TP 与逐层传输不匹配。

4.5 异常与回退¶

批处理结果必须与 key 一一对应。缺失项、布尔值和非整数结果会被拒绝。
range 异常会在唤醒计算线程之前发布 invalid-block 和 abort 状态。
put 异常会撤销 PROCESSING 对象。
get 异常只有在确认 range 调用已退出后，才会结束 session。
仅消费者 worker 会在 load hook 中推进层游标，而不是依赖一个永远不会被调用的 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 大小、块数量和预取深度。

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

6. 后续优化建议¶

P0：生产环境前必需¶

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

P1：高影响性能改进¶

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

P2：进一步演进¶

通过在 key 和对象头中记录组布局，并为每个组维护独立的完整性位图，来支持混合和多组布局。
通过在 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、多节点网络以及性能基准测试。

本页原文 Markdown：在 AtomGit 查看·内容源自开源项目 vl/vllm-ascend