逐层与稀疏 KV 缓存卸载设计¶

本文档解释了为什么 Prefill 和 Decode 使用不同的 KV 缓存卸载策略，这两种策略如何工作，以及它们一起使用时如何保持正确性。

有关安装和配置，请参阅 逐层与稀疏 KV 缓存卸载指南。

该设计基于 RFC #48203。

1. 动机¶

KV 缓存是长序列推理中主要的 NPU 内存开销。将其移至主机内存可以增加容量，但 Prefill 和 Decode 无法高效地使用相同的传输策略。

在每个 Decode 步骤中加载完整层会增加无法被 Decode 计算隐藏的传输延迟。稀疏 Decode 卸载通过仅移动由稀疏注意力选定的条目来避免该开销。

以下概念图说明了二者的区别。在 Prefill 中，少量完整层缓冲区会随时间复用。在 Decode 中，完整的主 KV 历史保留在主机内存中，而完整的 indexer 缓存和稀疏主 KV 选择保留在 NPU 上。RFC 图中的设备内存对应于本实现中的 NPU HBM。

逐层 Prefill 卸载。来源：RFC #48203.

稀疏 Decode 卸载。来源：RFC #48203.

2. 系统概览¶

生产设计面向分离式 Prefill/Decode 部署。目前不支持混合 KV 缓存布局。P/D 共置仅限于调试路径，不属于受支持的生产工作流。

组合设计包含四个存储区域：

• Prefill NPU 内存中的可复用层缓冲区；
• 由 Memcache 支持的 Prefill 主机 KV 池；
• Decode 主机内存中的完整主 KV 缓存；以及
• Decode NPU 内存中的 indexer 缓存和每层热 top-k 缓冲区。

flowchart LR
    PHost["Prefill host KV pool<br/>Memcache"]
    PNPU["Prefill NPU<br/>reusable layer buffers"]
    DHost["Decode host<br/>full main KV cache"]
    DNPU["Decode NPU<br/>indexer + per-layer hot top-k buffers"]

    PHost <-->|"Layerwise load / save"| PNPU
    PNPU -->|"Remote D2H pull<br/>main KV"| DHost
    PNPU -->|"Remote pull<br/>indexer / LIC8 scale"| DNPU
    DNPU -->|"New-token D2H"| DHost
    DHost -->|"Top-k miss H2D"| DNPU

AscendStoreConnector 通过 Memcache 管理逐层 Prefill 卸载。 SfaRemoteD2HConnector 暴露 Prefill NPU 缓冲区，并允许 Decode 通过 MemFabric 将主 KV 拉取到 Decode 拥有的主机内存中，将 indexer 数据拉取到 rank 本地 NPU 内存中。可选的 LIC8 scale 数据跟随 indexer 的目标位置。

MemFabric 数据路径在启动时通过 kv_connector_extra_config["memfabric_transfer_protocol"] 选择，而不是硬件 检测：sdma（默认）和 device_rdma 面向 A3 系列节点，而 device_urma 面向 950PR&950DT Products 节点。Prefill 和 Decode 必须使用相同 协议。

传输粒度¶

数据路径根据所需地址何时确定而使用不同的传输粒度：

Prefill 传输在每层完成后即可确定，并且可以以更大的范围发出。Decode 未命中地址仅在 top-k 选择和驻留查找之后才确定，因此它们使用稀疏 token 行复制。

3. 逐层 Prefill 卸载¶

缓冲区规划¶

逐层 Prefill 卸载将许多逻辑缓存承载层映射到更少的物理 NPU 缓冲区。某些层可以保留专用缓冲区；其余兼容层共享一个可复用缓冲区池。

对于统一的缓存布局：

N = number of cache-bearing layers
I = number of independent layers
R = N - I
B = configured shared-buffer count

physical buffers = I + min(B, R)

main-KV NPU footprint ratio ~= physical buffers / N

占用率假设统一的缓存布局，且层缓冲区大小相等。它仅描述主 KV 存储，不包括 indexer 或其他固定 NPU 分配。

可复用层按轮询顺序分配到共享缓冲区。主 KV 规格不兼容的层不共享物理缓冲区。MTP 层参与相同的规划。可选的 indexer 缓存跟随其主缓冲区分配，并且仅在需要时分配。

在张量描述符合并之前，规划器会验证缓存规格、张量大小和可选 indexer 布局。不兼容的布局会导致初始化失败，而不是共享大小错误的存储。

执行流水线¶

对于每一层：

1. 在需要时从 Memcache 加载缓存前缀；
2. 等待直到目标物理缓冲区可以安全覆盖；
3. 运行注意力并更新 KV；
4. 将更新后的 KV 保存到主机内存；以及
5. 在计算继续时预取后续层。

多个物理缓冲区使传输和计算能够重叠。Prefill 具有足够的每层计算量来隐藏大部分加载和保存延迟，这就是为什么完整层传输适合该阶段。

复用不变量¶

物理缓冲区不能在其先前内容的每个消费者都完成之前被复用。其 Memcache 保存必须在复用前完成。在联合部署中，Decode 侧远程读取也必须完成。

实现按物理存储槽跟踪完成状态，而不仅仅按逻辑层名称。这很重要，因为多个逻辑层可能引用同一个 NPU 地址。

4. 稀疏 Decode 卸载¶

内存布局¶

稀疏 Decode 卸载将完整的主 KV 缓存保留在固定的 Decode 主机池中。Decode NPU 内存包含：

• 用于选择重要 token 的秩本地索引器缓存；以及
• 包含活动 Decode 行最近选中的主 KV 条目的每层 K/V 热缓冲区。

每个 Decode DP 秩拥有独立的主主机池。在一个 DP 秩内，其 TP 秩共享该池，而每个 TP 秩拥有其本地注意力计算所需的索引器数据。控制面端口按 kv_port + dp_rank * d_tp_size + tp_rank 分配，使 DP 和 TP 端点保持分离。

Decode 步骤¶

对于每个 Decode 步骤：

1. Decode 生成新 token 的 K/V，而不将其写入完整的 NPU 分页主缓存；
2. 新的 K/V 行被复制到共享主机池中的相应逻辑槽位；
3. 索引器选择逻辑 top-k token 位置；
4. 每层 LRU 驻留表识别命中，为未命中分配物理热缓冲区槽位，并在必要时选择驱逐槽位；
5. 仅未命中行从主机内存复制到 NPU；
6. 逻辑 top-k 位置被重映射到物理热缓冲区槽位；以及
7. 稀疏注意力消费驻留 K/V，同时保留驻留元数据供下一步使用。

在张量并行下，Decode 产生的 K/V 会被复制。TP 秩 0 分配共享主机池并写入新 token 行。所有 TP 秩通过 MemFabric 卸载路径提供的广播全局虚拟地址访问该分配。

完整历史仍保留在主机内存中，但传输量与 top-k 未命中数成正比，而不是与总序列长度成正比。这使得卸载在低计算量的 Decode 阶段变得可行。

5. 联合 Prefill 与 Decode 设计¶

连接器组成¶

Prefill 使用 MultiConnector 组合 Remote D2H 完成提供者和 AscendStoreConnector。AscendMultiConnector 会优先调用完成提供者，无论它们在连接器配置中的顺序如何。对于每个层缓冲区：

• AscendStore 将该层保存到 Prefill 主机 KV 池；以及
• Remote D2H 发布源元数据并等待 Decode 确认。

缓冲区复用门仅在两条路径都完成后打开。这可防止 Prefill 在 Decode 仍在读取时覆盖共享物理缓冲区。

请求流程¶

1. Decode 分配其主主机目的地址和秩本地索引器目的地址。
2. Decode 通过代理元服务器通告其端点和并行拓扑。
3. Prefill 逐层计算请求。
4. 每层 KV 写入完成后，Prefill 启动 Memcache 保存并发布 Remote D2H 就绪状态。
5. Decode 将该层拉取到其目的地址并确认完成。
6. 在所有必需完成信号到达后，Prefill 释放物理缓冲区以供复用。
7. 只有当完整请求在所有必需秩上达到终态传输状态后，Decode 才会调度该请求。

物理缓冲区完成与请求完成是相互独立的：

• 物理缓冲区完成保护 Prefill 内存复用；以及
• 请求完成控制 Decode 何时可以开始推理。

6. 不等张量并行¶

支持的联合拓扑为：

p_tp_size >= d_tp_size
p_tp_size % d_tp_size == 0
ratio = p_tp_size // d_tp_size

对于 Prefill 秩 p_rank：

d_rank = p_rank // ratio
group_member_idx = p_rank % ratio

映射到一个 Decode 秩的 Prefill 秩构成一个贡献者组。

• 主 KV 在 Prefill 上被复制，因此只有贡献者成员 0 传输该 Decode 秩拥有的主份额。
• 贡献者传输互不重叠的索引器范围，其并集构成完整的秩本地 Decode 索引器。
• 没有块的贡献者仍会确认该层，以避免完成状态死锁。

7. 验证与故障处理¶

在 Remote D2H 读取之前，Decode 会验证主张量数量和大小、索引器存在性和大小、可选 LIC8 缩放布局以及目的块范围。不匹配会使该层失败，而不是留下部分有效的目的数据。

Decode 处理就绪通知后，会回复 READ_DONE 或 READ_FAILED。报告的 MemFabric 读取失败会使 Decode 目的地址失效，并以错误释放 Prefill 侧等待者，防止静默损坏。丢失的确认或停止的 Decode 进程不被视为终态，因此 Prefill 可以继续等待。

8. 当前边界¶

• 逐层共享缓冲区卸载需要 Memcache 后端和 eager 模式。
• 稀疏 Decode 卸载需要 Model Runner V1 和 SFA/MLA 稀疏注意力模型。主 KV 缓存必须使用 BF16；LIC8 量化仅支持设备驻留索引器缓存。
• 不支持混合 KV 缓存布局。
• 稀疏 Decode 卸载支持 DP 和 TP；不支持 CP 和 PP。
• 联合部署要求 Prefill TP 大于或等于 Decode TP，且可被 Decode TP 整除。
• MemFabric 是唯一支持的 Remote D2H 传输后端。
• 逐层缓冲区复用目前不能与 MooncakeLayerwiseConnector 组合使用，因为它不提供按缓冲区的传输完成门。计划在后续更新中支持。
• 未实现连接器级别的数据读取重试。