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

本文档解释了为什么Prefill和Decode使用不同的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。

Layerwise Prefill KV cache offload concept

分层Prefill卸载。来源：RFC #48203。

Sparse Decode KV cache offload concept

稀疏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系列产品节点。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层参与相同的规划。可选的索引器缓存跟随其主缓冲区分配，仅在需要时分配。

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

执行流水线¶

对于每一层：

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

多个物理缓冲区使传输和计算可以重叠。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端点互不重叠。

解码步骤¶

对于每个解码步骤：

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

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

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

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

连接器组合¶

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

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

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

请求流程¶

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

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

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

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缩放布局以及目标块范围。任何不匹配都会使该层失败，而不会留下部分有效的目标数据。

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

8. 当前限制¶

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