分层与稀疏KV缓存卸载设计

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

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

该设计基于 RFC #48203。

1. 动机

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

阶段	计算模式	卸载策略	NPU驻留数据
Prefill	每层计算量大	传输完整层并将传输与计算重叠	少量可复用的层缓冲区
Decode	每个token计算量小	将完整KV保留在主机内存中，仅加载选中的条目	索引器缓存和每层热top-k缓冲区

在每一步Decode中加载完整层会增加传输延迟，而Decode计算无法隐藏这些延迟。稀疏Decode卸载通过仅移动稀疏注意力选中的条目来避免这一开销。

以下概念图说明了差异。在Prefill中，少量完整层缓冲区随时间被复用。在Decode中，完整的主KV历史保留在主机内存中，而完整的索引器缓存和稀疏主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内存中的索引器缓存和每层热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拥有的主机内存中，并将索引器数据拉取到rank本地的NPU内存中。可选的LIC8缩放数据跟随索引器的目标位置。

MemFabric 数据路径在启动时通过 kv_connector_extra_config["memfabric_transfer_protocol"] 选择，而不是通过硬件 检测：sdma（默认）和 device_rdma 面向 A3 系列节点，而 device_urma 是 Ascend 950PR&950DT系列产品节点的连接器协议。连接器 协议的可用性并不意味着支持稀疏 Decode 卸载：稀疏 Decode 卸载会拒绝 Ascend 950PR&950DT系列产品。Prefill 和 Decode 必须使用相同的 协议。

传输粒度

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

路径	粒度	用途
Prefill NPU到Decode主机	层/块范围	填充Decode的完整主KV历史
Prefill NPU到Decode NPU	索引器/缩放张量内的块范围	填充rank本地索引器和可选的LIC8缩放数据
Decode当前KV到Decode主机	Token行	追加新生成的KV，无需完整的NPU主缓存
Decode主机到Decode NPU	Top-k未命中行	仅填充每层热缓冲区中缺失的条目

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存储，不包括索引器或其他固定的NPU分配。

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

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

执行流水线

对于每一层：

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

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

复用不变量

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

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

4. 稀疏解码卸载

内存布局

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

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

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

解码步骤

对于每个解码步骤：

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

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

完整历史记录仍保留在主机内存中，但传输量仅与top-k未命中数成正比，而非总序列长度。这使得卸载对低计算量的解码阶段切实可行。

融合 Copy-SFA 算子

Decode 稀疏卸载默认使用融合 LIM 和 Copy-SFA。Python 集成要求 已安装的 _C_ascend 扩展中包含这些算子：

• npu_fused_lightning_indexer_manage
• npu_fused_quant_lightning_indexer_manage (LIM C8)
• npu_fused_scatter_copy_sparse_flash_attention

仓库中包含融合 LIM 和 Copy-SFA 的原生内核、PyTorch 绑定以及 Python 集成。已安装的构建必须包含这些 算子才能运行稀疏 Decode 卸载；仅有 Python 源码 是不够的。

LIM C8 使用 npu_fused_quant_lightning_indexer_manage，该算子已包含在 本分支中。通过 --attention_config.indexer_kv_dtype int8 启用它。 随后 indexer 会对 query 和 key 进行 Hadamard 旋转和 int8 量化， _lim_select 将该 int8 query 和 fp16 反量化缩放因子转发给 C8 LIM。19 参数管理 ABI 与 BF16 LIM 一致，因此 Copy-SFA 和 miss H2D 保持不变。C8 内核仅支持 int8；indexer_kv_dtype fp8 不是 LIM C8 的服务路径。

Copy-SFA 接收 dram_k_rope 和 dram_kv_cache 作为由已注册 MemFabric 内存 支撑的 CPU 张量视图。原生适配器接受这些已注册的 主机视图。其余输入保留在 NPU 上。普通 CPU 分配不能 替代已注册的 DRAM，并且 Python 集成不会 将主机缓存移动到 NPU。

5. 联合预填充与解码设计

连接器组合

预填充使用MultiConnector组合远程D2H完成提供程序和AscendStoreConnector。无论连接器配置中的顺序如何，AscendMultiConnector都会首先调用完成提供程序。对于每个层缓冲区：

• AscendStore将层保存到预填充主机KV池；以及
• 远程D2H发布源元数据并等待解码确认。

缓冲区重用门仅在两条路径都完成后才打开。这防止预填充在解码仍在读取共享物理缓冲区时覆盖它。

请求流程

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

物理缓冲区完成和请求完成是分开的：

• 物理缓冲区完成保护预填充内存重用；以及
• 请求完成控制解码何时可以开始推理。

6. 非均匀张量并行

支持的联合拓扑为：

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

对于预填充秩p_rank：

d_rank = p_rank // ratio
group_member_idx = p_rank % ratio

映射到同一个解码秩的预填充秩构成一个贡献组。

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

7. 验证与故障处理

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

Decode 处理就绪通知后，会回复 READ_DONE 或 READ_FAILED。上报的 MemFabric 读取失败会使 Decode 目标失效，并以错误释放 Prefill 侧的等待者，从而防止 静默损坏。初始 MF_META 握手有 100 秒超时， 如果 Decode 未回复则抛出错误。逐层读取完成 门控是独立的：它检查上报的错误和已停止的发送线程， 但没有整体截止时间。如果某一层的确认丢失而 发送线程仍然存活，Prefill 可能会继续等待。

8. 当前限制

• Ascend 950PR&950DT系列产品不支持稀疏 Decode 卸载，包括 keep-device-cache 和 overlap 路径。配置校验会在卸载分配之前检查已安装 构建的设备元数据；它不要求 NPU 初始化或运行时 SoC 查询。
• 逐层共享缓冲区卸载需要 Memcache 后端和 eager 模式。
• 稀疏 Decode 卸载支持 Model Runner V1 和 V2，配合 SFA/MLA 稀疏注意力模型。主 KV 缓存必须使用 BF16；打包的主缓存 （C8 或 TQ4）会被拒绝。LIC8 量化单独支持 设备常驻的 indexer 缓存。
• 不支持混合 KV 缓存布局。
• 稀疏 Decode 卸载支持 DP 和 TP；不支持 CP 和 PP。
• 联合部署要求 Prefill TP 大于或等于 Decode TP， 且能被 Decode TP 整除。
• MemFabric 是唯一支持的远程 D2H 传输后端。
• 逐层缓冲区复用目前无法与 MooncakeLayerwiseConnector 组合使用，因为它不提供逐缓冲区传输 完成门控。支持计划在后续更新中提供。
• 未实现连接器级的数据读取重试。