面向 GLM-5.2 的共享 Indexer KV Cache OffloadCANN 平台上从 KV 容量释放到系统吞吐提升的完整方案【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer说明本文基于 docs/models/glm_5_2/glm_5_2_offload_guide.md 编写并以仓库内 GLM-5.2 样例源码模型实现、缓存管理、算子加载、配置与评测为佐证进行深度扩充。本文围绕 CANN / cann-recipes-infer 仓库中 GLM-5.2 样例的共享 Indexer KV Cache Offload 方案展开先梳理其从面向 DeepSeek V3.2 的基础按需加载 Offload 出发的设计演进再深入讲解更大的 Device 侧缓存池、跨层共享规划、组相联 FIFO 替换、SFA 优先回填后置四个关键技术决策最后给出三个定制 AscendC 算子的职责划分、仓库内完整的启用与算子编译流程以及双机 Atlas A3 上的端到端评测数据。读完本文你可以掌握长上下文场景下 KV Offload 从容量收益走向吞吐收益的完整设计思路与可复现的仓库实操路径。一、问题背景长上下文让 KV Cache 先于算力触及 HBM 上限大语言模型的上下文窗口正在持续增长。长文档问答、代码库分析、长时间多轮对话等场景正将推理服务推向数十 K 乃至上百 K token 的上下文规模。这一趋势给推理系统带来一个结构性的挑战KV Cache 的显存占用。在自回归 decode 过程中每个生成的 token 都会在 HBM 中产生 Key-Value 状态序列越长、并发请求越多累积的 KV Cache 就越庞大且增长是线性的。HBM 容量是固定的。当 KV Cache 占据显存的大半即便计算单元负载尚低系统也无法接纳更多并发请求——内存先于计算触及上限。KV Cache Offload 的思路很直接将完整的 KV Cache 放到容量更大、成本更低的 Host 内存中仅在注意力计算真正需要时将相关数据搬运回 HBM。这打破了 HBM 的容量限制理论上可以将 batch size 翻倍甚至更高。但在实际部署中一个问题随之而来Host 到 Device 的数据搬运H2D直接叠加在 decode 延迟上。如果新增的 batch 容量被传输开销抵消Offload 便只解决了内存容量却未能转化为吞吐提升。本文介绍的方案沿一条递进路径展开其出发点是此前面向 DeepSeek V3.2 实现的基础 KV Offload 方案相关实现已开源见 deepseek_v3.2_exp_inference_guide.md最终落地为面向 GLM-5.2 的共享 Indexer KV Cache Offload实现代码位于 models/glm_5_2。二、此前方案面向 DeepSeek V3.2 的按需加载式 KV Offload2.1 稀疏注意力的数据加载模型此前面向 DeepSeek V3.2 实现的 KV Offload 方案是本工作的出发点。其 decode 路径有两个紧密配合的组件Lightning Indexer从全部历史 token 中筛选出与当前 query 最相关的 2048 个位置即 Top-K 稀疏索引Sparse Flash AttentionSFA仅在这 2048 个位置上执行注意力计算。Indexer 决定哪些 token 参与注意力SFA 执行实际计算。两者之间有一个数据搬运环节SFA 所需的 2048 条 KV 数据散布在完整序列的 KV Cache 中需要被收集到一起。当前 query | v Lightning Indexer ---- Top-K token 位置 | v Host 侧完整 KV Cache ------- 选中的 KV ------- SFA该基础方案将完整 KV Cache 放置在 Host 内存中。Prefill 阶段模型在 Device 侧逐层生成 KV 后将其传输至 Hostdecode 阶段GatherSelectionKvCache算子根据 Indexer 的 Top-K 结果从 Host 侧完整 KV Cache 中加载对应的 KV 数据并交由 SFA 使用。这一流程释放了完整序列 KV Cache 此前占据的大部分 HBM腾出的空间可支撑更大的 batch 或更长的上下文。在仓库中GatherSelectionKvCache算子位于 ops/ascendc/src/gather_selection_kv_cache包含op_host算子原型、tiling与op_kernel内核实现含split_bs_reuse复用变体两部分并有独立的 torch 绑定 与 文档说明。在 GLM-5.2 样例中custom_op_loader.py 的load_offload_gather_op()会在 import 时把该算子挂载为torch_npu.npu_gather_selection_kv_cache。2.2 相邻 step 复用Gather 算子还利用了一个重要的局部性相邻 decode step 的 Top-K 往往高度重叠。上一个 step 已搬到 Device 侧的 KV 可直接复用无需再次从 Host 读取。全模型 Top-K 轨迹分析显示相邻 step 的平均复用率约为 60%。这意味着每次 decode约六成 KV 可从 Device 侧直接获得Host 只需提供剩余约四成。当前 Top-K 上一步选集中命中的部分 从 Host 侧完整 KV Cache 加载的缺失部分这一复用将 H2D 数据传输量降至全部从 Host 加载的约 40%但复用范围仅限于上一步的 2048 个位置。仍有近五分之二的选中 token 需要进行离散的 H2D 搬运——每次未命中对应一个独立位置的小段 KV 数据。2.3 容量收益与吞吐缺口该基础方案带来了明确的容量收益相同上下文长度下可承载的 batch size 约可翻倍。但 batch 扩大后每步需要同时处理的请求和数据规模随之增加叠加剩余的 Gather 搬运开销decode 延迟显著上升。在端到端评估中扩展的 batch 容量尚未充分转化为吞吐提升。这明确了下一步的优化目标提升缓存命中率将 H2D 数据搬运移出关键路径。三、更大的 KV 缓存池瓶颈从数据搬运转向缓存管理基础方案的选择缓冲区只有 2048 个条目因为它的容量等于 SFA 的 Top-K 输入量。如果进一步扩大复用范围让 Device 侧保留一个更大的 KV 缓存池——例如 8192 或 16384 个条目下一轮的 Top-K 就能以更高概率命中 Device 侧已有的数据。我们构建了一个包含 8192 个条目、采用 LRU 策略管理的 Device 侧 KV 缓存池。在相同 Top-K 轨迹下平均命中率从 58.85% 提升至 90.50%。平均而言仅有 9.50% 的选中 token 需要从 Host 加载。在这一命中率水平上H2D 数据搬运已不再是主要开销。但大容量缓存池改变了开销的构成。每层仍需独立处理自己的 Top-K id逐条判定数据位于缓存池还是 Host、对重复未命中去重、选定替换槽位、维护缓存元数据。LRU 的细粒度管理还需要维护访问时序和淘汰决策。这些操作多属于 Vector Core 不擅长的标量和控制密集型计算。KV 数据路径虽已通畅逐层重复的缓存管理操作却形成了新的 decode 瓶颈。Device 侧缓存设计平均命中率主导压力仅保留上一步的 2K Top-K 选集58.85%离散 H2D KV 加载8K KV 缓存池 LRU90.50%命中分类与缓存管理这一瓶颈迁移指明了后续工作的方向增大 Device 侧缓存池缓解了数据搬运问题但将缓存池管理做成高效算子还需要在控制面上寻找进一步的优化机会。GLM-5.2 的架构恰好提供了这样的切入点。四、GLM-5.2 共享 Indexer从逐层规划到逐组复用4.1 IndexShare四层一组的 Top-K 复用结构GLM-5.2 在 attention 设计中引入了一项关键特性共享 IndexerIndexShare。默认配置下一层独立执行 Indexer 生成新的 Top-K紧接着的三层直接复用这组结果构成四层一组的 Indexer 共享组。在仓库源码中这一结构由 configuration_glm.py 显式建模每层通过indexer_types标记为full本层运行 Lightning Indexer、持有 indexer 权重或shared复用上一个 full 层的 top-k、不持有 indexer 权重。若未显式指定则按如下规则推导# models/glm_5_2/models/configuration_glm.py # full iff (max(i - index_skip_topk_offset 1, 0) % index_topk_freq) 0 self.indexer_types [ full if (max(i - offset 1, 0) % freq) 0 else shared for i in range(num_hidden_layers) ]即默认表现为每 N 层一组、组首层为 full、其余层为 shared的周期结构。模型侧实现modeling_glm.py中shared层通过skip_topk indexer_types[layer_idx] shared跳过本层 Indexer 计算直接复用上一层输出的prev_topk_indices。四层共享同一组 Top-K id因此也共享了基于这些 id 的缓存决策。某个 token id 在组内任一层命中 Device 侧缓存池在其余三层也会映射到相同的缓存槽位某个 token id 未命中整组也需要执行相同的回填规划。命中分类和未命中规划是标量最密集、Vector Core 最不擅长的操作。共享 Indexer 结构将它们从逐层执行压缩为逐组执行——规划做一次四层复用。标量开销被分摊到四层而各层独立的 KV 数据搬运不受影响。4.2 共享规划规划器DsaPlan对共享 Top-K 执行一次性处理逐条记录数据来自缓存池还是 Host同时输出一份紧凑的未命中条目列表用于后续回填。组内后续层直接沿用同一份规划。各层拥有各自的 KV 和 RoPE 数据搬运过程独立进行但命中检测、来源分类、未命中去重和替换位置选择均不再重复计算。在仓库实现中这一一次规划、整组复用的语义由 offload_cache.py 的共享组元数据管理支撑build_dsa_shared_group_specs()按indexer_types把各层归入共享组并标记owner_layer组首 full 层规划结果通过set_dsa_shared_group_plan()/get_dsa_shared_group_plan()暂存组内共享层直接取用待组末层处理完毕后再clear_dsa_shared_group_plan()清除。4.3 组相联 FIFO大容量缓存池需要配套的替换策略。LRU 能够较好地匹配访问模式但在 NPU 上高效实现 LRU需要持续维护条目的最近访问状态和淘汰顺序执行开销较高。我们采用组相联 FIFO组织缓存池。所谓组相联set-associative是指先将缓存池划分为若干独立的组。每个 token id 通过哈希映射到唯一一组但可以存放在该组内任意一个槽位。相比每个 token 只能对应一个固定槽位的直接映射方式组相联可以降低哈希冲突相比允许条目存放在整个缓存池任意位置的全相联方式它又将查找和替换范围限制在较小的组内。每组包含固定数量的槽位替换仅发生在组内。每组维护一个 FIFO 指针仅在装入新条目时推进一次。当组内需要装入新条目时指针指向的槽位被替换随后指针移动到下一个槽位缓存命中不会改变该指针。因此实现上采用轮转替换语义上等价于按照装入顺序执行组内 FIFO。该方案在 NPU 上有三个适配性优势查找和替换状态被限制在较小的组内元数据开销可控不同组之间相互独立天然支持多核划分与并行处理FIFO 仅在新条目装入时更新无需在每次命中时维护访问时序。仓库中的常量定义印证了这一设计offload_cache.py默认池大小_DSA_DEFAULT_POOL_SIZE 8192enlarge_pool_size: True时切换到_DSA_LARGE_POOL_SIZE 16384组相联度_DSA_SET_ASSOCIATIVITY 16lru_counter的形状为(batch_size_per_rank, dsa_pool_size // 16)即每个 set 维护一个替换计数FIFO 指针同时以id_to_slot哈希表dsa_id_range不小于 131072完成 token id 到槽位的映射。缓存数据侧每个共享组在 Device 上维护独立的常驻池张量dsa_pool_key_valuesnope 与 rope 两部分shape 为(batch, pool_size, dim)见init_cache()。组相联在降低哈希冲突的同时将查找和替换范围限制在组内FIFO 则将替换机制做得足够轻量适配并行执行的要求。4.4 SFA 优先回填后置当前层注意力计算直接依赖选中的 KV——数据未就绪SFA 就无法启动。因此执行调度优先准备选中的 KV数据就绪后立即推进 SFA。缓存池回填的时限要宽松得多新装入的条目是供后续 decode step 使用的。回填因此在辅流上发起与 SFA 和 MoE 并行推进。共享 Indexer Top-K | v 共享规划 | ---- 逐层 KV 准备 ---- SFA ---- MoE | ---- 辅流缓存池回填这一调度将 Offload 的关键阻塞段压缩到准备选中 KV这一步缓存维护开销则尽可能与计算重叠不单独占用 decode 延迟。五、实现三个算子三条路径上述设计在代码层面映射为三个定制 AscendC 算子源码位于 ops/ascendc/src/dsa_plan、ops/ascendc/src/dsa_serve、ops/ascendc/src/dsa_install算子执行频率职责DsaPlan每个 Indexer 共享组一次分类 Top-K、判定缓存池命中或未命中、生成紧凑的未命中记录DsaServe每层一次从缓存池和 Host 侧完整 KV Cache 中准备 SFA 所需的 KV 和 RoPEDsaInstall每层一次回填未命中条目、执行 FIFO 淘汰、更新缓存池状态DsaPlan承载所有共享的标量工作将 Top-K id 和缓存池元数据转化为紧凑的执行规划。DsaServe依照规划执行数据搬运优先交付 SFA 所需的数据。DsaInstall依照未命中记录在辅流上执行延迟回填。三个算子的划分形成三条职责分明的执行路径规划在组级别完成不随层数放大数据服务每层独立进行但不再包含重复的决策逻辑回填与计算并行不占用关键延迟。5.1 模型侧编排以源码为证模型侧的编排逻辑集中在 shared_indexer_offload.py 的run_shared_indexer_offload()中与上述设计一一对应逐组规划offload_cache.is_dsa_shared_group_owner(layer_idx)判定当前层是否为组首 owner 层owner 层调用torch.ops.custom.dsa_plan生成plan与install_records并存入共享组非 owner 层直接复用get_dsa_shared_group_plan()取回的同一份规划逐层数据服务每层独立调用torch.ops.custom.dsa_serve将plan作用到本层的full_kv_cache/full_k_rope与缓存池pool_kv_view/pool_rope_view产出 SFA 所需的selection_kv_cache/selection_k_ropecompact_layout1表示池按[batch, pool_size, dim]紧凑布局存储辅流回填非 owner 层回填本层、组末层额外回填 owner 层的数据metadata_update仅在组末层触发用于推进 FIFO 指针并更新pool_ids/id_to_slot/lru_counterenable_install_stream开启时回填经npu_stream_switch切到独立辅流执行并通过record_event/wait_event与主计算流同步从而与 SFA、MoE 计算重叠。在 eager / npugraph_ex / ge_graph 三种执行模式下算子加载路径由 custom_op_loader.py 统一处理load_offload_gather_op()挂载npu_gather_selection_kv_cache基础 Offload 算子register_dsa_shared_ge_converters()ge_graph 模式注册 DsaPlan/DsaServe/DsaInstall 的 torchair GE converter含dsa_functionalization辅助模块register_dsa_shared_npugraph_reinplace()npugraph_ex 模式注册dsa_serve_functional/dsa_serve、dsa_install_functional/dsa_install的 inplace 配对使图捕获阶段能够正确执行 host-device 数据搬运。Runner 侧runner_glm.py在shared_indexer_offload开启时按exe_mode选择上述加载路径。此外offload_cache.py 还实现了 MTP 场景的 gather 复用enable_mtp_gather_reuseMTP draft 步冻结首个 decode step 的 top-k 后其选中的历史位置 KV 不可变首个 step 收集到的 selection buffer 在后续 draft 步中保持有效可跳过冗余的重复 gather。六、在仓库中启用共享 Indexer Offload6.1 依赖的定制算子编译安装KV Offload 依赖仓内自定义 AscendC 算子非 torch_npu 内置须先编译安装详见 models/glm_5_2/README.md# 1) 编译算子内核A3 默认 ascend910_93bisheng 随 CANN 提供 cd ops/ascendc bash build.sh -n gather_selection_kv_cache;dsa_plan;dsa_serve;dsa_install # 安装包位于 output/CANN-custom_ops-*-linux.arch.run # 2) 安装内核到 CANN opp cd output ./CANN-custom_ops-*-linux.*.run --quiet --install-path${ASCEND_HOME_PATH}/opp # 3) 编译并安装 torch 绑定缺 ninja 时先 pip install ninja cd ../torch_ops_extension bash build_and_install.sh # 生成并 pip 安装 custom_ops wheel # 4) 运行前 source 自定义算子环境设置 ASCEND_CUSTOM_OPP_PATH / LD_LIBRARY_PATH source ${ASCEND_HOME_PATH}/opp/vendors/customize/bin/set_env.bashmodeling_glm.py在 import 时自动把基础 offload 算子挂到torch_npu.npu_gather_selection_kv_cache开启shared_indexer_offload时Runner 初始化会按需加载共享 IndexShare offload 相关自定义算子。因此模型侧无需改动只要上面的算子已安装、且运行前 source 了第 4 步的环境即可。若遇到_OpNamespace custom object has no attribute类报错可参考 ops/ascendc/README.md 重新编译对应算子。6.2 配置项offload 三开关与执行模式仓库提供了两份 offload 示例配置models/glm_5_2/configglm_5_2_rank_32_32ep_w8a8_offload.yaml基础 KV Offload 路径glm_5_2_rank_32_32ep_w8a8_offload_mtp.yaml开启共享 Indexer Offload 与 MTP3 的完整路径。在model_config中设置 offload 选项详细参数说明见 config/README.mdmodel_config: enable_offload: True # [False, True] 开启 KV Offload长序列、大 batch 场景 shared_indexer_offload: True # [False, True] 使用 GLM-5.2 IndexShare 共享规划路径 enlarge_pool_size: False # [False, True] 设备侧常驻 token 池从 8K 扩大到 16K next_n: 3 # MTP 步数支持 [0, 1, 2, 3] pa_block_size: 128 # PagedAttention 块大小支持 [128, 256] with_ckpt: True # 是否加载权重 enable_multi_streams: True # 多流并行 enable_online_split_weight: True # 在线切分权重 exe_mode: npugraph_ex # [ge_graph, eager, npugraph_ex]decode 执行模式三个开关的语义与作用enable_offload: True把全量 MLA KV 卸载到 Host swapped memorydecode 时按 DSA top-k 把命中的 block gather 回 Device面向全量 KV 放不下 HBM 的长上下文场景shared_indexer_offload: True利用 GLM-5.2 的 IndexShare 特性复用同一组 top-k 规划减少共享层重复的命中判断和缓存管理开销并将常驻池回填与 SFA、MoE 计算并行enlarge_pool_size: True把设备侧常驻 token 池从 8K 扩大到 16K进一步提高命中率并减少 host 到 device 的 KV 搬运。该开关经 offload_cache.py 的resolve_dsa_pool_size()生效8192 → 16384。配套的并行配置参考A3 单卡双 die、world_size 2 * chip_numembed_tp_size/lmhead_tp_size切到 32W8A8 权重下 embedding 与 lm_head 走 TPattn_tp_size/dense_tp_size/moe_tp_size为 1MoE 经 32 路 EP 展开。数据侧示例input_max_len: 4096、batch_size: 32可按评测需求调整。6.3 运行流程各节点同步执行 infer.sh 即可拉起多卡推理# infer.sh 中指定配置文件 export YAML_FILE_NAMEglm_5_2_rank_32_32ep_w8a8_offload_mtp.yaml bash infer.sh运行前需在 executor/scripts/set_env.sh 中填写各节点 IP第 1 个为 master与 CANN 包路径并在 yaml 中把model_path指向转换好的 W8A8 权重转换命令见 models/glm_5_2/README.md 的weight_convert.sh部分。七、双机 Atlas A3 评测7.1 评测环境与口径评测环境为生产级配置双机 Atlas A3共 16 卡、GLM-5.2-W8A8 模型、64K 输入序列、MTP3。针对 64K 输入序列在相同设备资源约束下不使用 Offload 时KV Cache 容量仅能支持全局 batch size 上限为 64启用 Offload 后该上限提升至 128。实验分别测试 8K 和 16K 两种 Device 侧 KV 缓存池容量并覆盖 GE Graph 和 NPUGraph Ex 两种执行模式。吞吐按两种口径报告统一假设口径为在相同 MTP 接收长度假设下比较不同配置将 MTP 平均接收长度统一设为 2.7。该数值为经验假设仅用于归一化估算不代表各配置的实测接收长度实测口径使用推理过程中实际观测到的 MTP 平均接收长度。计算方式吞吐 全局 batch size × MTP 平均接收长度 /主模型 MTP3 单步延迟7.2 完整评测数据执行模式配置全局 batch缓存池容量主MTP3 延迟吞吐2.7 假设吞吐实测GE Graph无 Offload64-71.21 ms2,427 token/s2,504 token/sGE Graph共享 Indexer Offload1288K100.55 ms3,437 token/s41.6%3,642 token/s45.4%GE Graph共享 Indexer Offload12816K92.03 ms3,755 token/s54.8%3,982 token/s59.0%NPUGraph Ex无 Offload64-71.59 ms2,414 token/s2,507 token/sNPUGraph Ex共享 Indexer Offload1288K105.04 ms3,290 token/s36.3%3,452 token/s37.7%NPUGraph Ex共享 Indexer Offload12816K95.67 ms3,612 token/s49.6%3,791 token/s51.2%7.3 从数据中归纳的三点结论可承载的 batch 上限翻倍延迟增长可控。针对 64K 输入序列Offload 将可承载的全局 batch size 上限从 64 提升至 128。16K 缓存池下主模型 MTP3 单步延迟相较无 Offload 配置增长 29.2%GE Graph和 33.6%NPUGraph Ex。延迟增长幅度远低于 batch 上限的增幅表明 Offload 的附加开销处于有效控制之下。吞吐净增约 50%。16K 缓存池下按统一的 2.7 经验假设估算GE Graph 和 NPUGraph Ex 的吞吐分别提升 54.8% 和 49.6%按各配置的实测接收长度计算吞吐分别提升 59.0% 和 51.2%。batch 容量的扩大切实转化为了系统吞吐收益。16K 缓存池优于 8K 缓存池。在最终的组相联 FIFO 方案下缓存池容量从 8K 扩大至 16K 后平均命中率由约 85% 提升至 90% 以上需要从 Host 加载并回填的未命中数据随之减少。在相同全局 batch size 128 下16K 缓存池相较 8K 缓存池吞吐进一步提升约 9.3%GE Graph和 9.8%NPUGraph Ex按 2.7 MTP 接受长度经验假设估算。这表明扩大缓存池容量所带来的命中率提升可以进一步降低 Offload 的端到端开销。这些对比印证了 Offload 设计的核心判断Host 内存用于承载完整 KV Cache从而提高可支持的全局 batch 上限Device 侧缓存池、跨层复用和调度优化则用于压低附加开销。两者协同使 batch 上限翻倍最终转化为约 50% 的吞吐净增。八、扩展思考KV Cache 管理的系统化方向GLM-5.2 的共享 Indexer 为本次 Offload 方案提供了关键的结构性支撑。四层一组的 Top-K 复用使标量规划开销得以分摊这一点离开了模型的特定架构是无法成立的。换言之KV Offload 的上限很大程度上由模型结构决定。不同模型的注意力设计、层间关联、计算与带宽的配比各不相同缓存大小、替换策略和执行调度的最优选择也会随之变化。从更长远的视角看KV Cache 的管理最终需要走向一个系统化的方向。Offload 解决了容量问题但单靠 Offload 还不足以应对上下文持续增长带来的全方位压力。一个更完整的思路是将 KV Cache 视为可分级、可调度的资源在 Host 内存与 Device HBM 之间建立缓存层次依据访问热度动态调配数据位置结合淘汰策略自动适应工作负载变化。与此同时稀疏注意力压缩每步计算量投机解码通过并行生成均摊延迟——三者各自有效但要系统性地应对长上下文下的吞吐和延迟挑战需要将它们协同部署每项技术的短板恰可由其他技术补足。参考本文涉及的仓库关键路径技术文档本体docs/models/glm_5_2/glm_5_2_offload_guide.md相关图示位于 docs/models/glm_5_2/figuresGLM-5.2 样例说明与算子编译步骤models/glm_5_2/README.md配置示例glm_5_2_rank_32_32ep_w8a8_offload.yaml、glm_5_2_rank_32_32ep_w8a8_offload_mtp.yaml、config/README.md模型侧编排shared_indexer_offload.py、offload_cache.py、custom_op_loader.py、configuration_glm.py、runner_glm.py定制算子ops/ascendc/src/gather_selection_kv_cache、ops/ascendc/src/dsa_plan、ops/ascendc/src/dsa_serve、ops/ascendc/src/dsa_install启动脚本models/glm_5_2/infer.sh、executor/scripts/set_env.sh【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 面向 GLM-5.2 的共享 Indexer KV Cache Offload：CANN 平台上从 KV 容量释放到系统吞吐提升的完整方案
面向 GLM-5.2 的共享 Indexer KV Cache OffloadCANN 平台上从 KV 容量释放到系统吞吐提升的完整方案【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer说明本文基于 docs/models/glm_5_2/glm_5_2_offload_guide.md 编写并以仓库内 GLM-5.2 样例源码模型实现、缓存管理、算子加载、配置与评测为佐证进行深度扩充。本文围绕 CANN / cann-recipes-infer 仓库中 GLM-5.2 样例的共享 Indexer KV Cache Offload 方案展开先梳理其从面向 DeepSeek V3.2 的基础按需加载 Offload 出发的设计演进再深入讲解更大的 Device 侧缓存池、跨层共享规划、组相联 FIFO 替换、SFA 优先回填后置四个关键技术决策最后给出三个定制 AscendC 算子的职责划分、仓库内完整的启用与算子编译流程以及双机 Atlas A3 上的端到端评测数据。读完本文你可以掌握长上下文场景下 KV Offload 从容量收益走向吞吐收益的完整设计思路与可复现的仓库实操路径。一、问题背景长上下文让 KV Cache 先于算力触及 HBM 上限大语言模型的上下文窗口正在持续增长。长文档问答、代码库分析、长时间多轮对话等场景正将推理服务推向数十 K 乃至上百 K token 的上下文规模。这一趋势给推理系统带来一个结构性的挑战KV Cache 的显存占用。在自回归 decode 过程中每个生成的 token 都会在 HBM 中产生 Key-Value 状态序列越长、并发请求越多累积的 KV Cache 就越庞大且增长是线性的。HBM 容量是固定的。当 KV Cache 占据显存的大半即便计算单元负载尚低系统也无法接纳更多并发请求——内存先于计算触及上限。KV Cache Offload 的思路很直接将完整的 KV Cache 放到容量更大、成本更低的 Host 内存中仅在注意力计算真正需要时将相关数据搬运回 HBM。这打破了 HBM 的容量限制理论上可以将 batch size 翻倍甚至更高。但在实际部署中一个问题随之而来Host 到 Device 的数据搬运H2D直接叠加在 decode 延迟上。如果新增的 batch 容量被传输开销抵消Offload 便只解决了内存容量却未能转化为吞吐提升。本文介绍的方案沿一条递进路径展开其出发点是此前面向 DeepSeek V3.2 实现的基础 KV Offload 方案相关实现已开源见 deepseek_v3.2_exp_inference_guide.md最终落地为面向 GLM-5.2 的共享 Indexer KV Cache Offload实现代码位于 models/glm_5_2。二、此前方案面向 DeepSeek V3.2 的按需加载式 KV Offload2.1 稀疏注意力的数据加载模型此前面向 DeepSeek V3.2 实现的 KV Offload 方案是本工作的出发点。其 decode 路径有两个紧密配合的组件Lightning Indexer从全部历史 token 中筛选出与当前 query 最相关的 2048 个位置即 Top-K 稀疏索引Sparse Flash AttentionSFA仅在这 2048 个位置上执行注意力计算。Indexer 决定哪些 token 参与注意力SFA 执行实际计算。两者之间有一个数据搬运环节SFA 所需的 2048 条 KV 数据散布在完整序列的 KV Cache 中需要被收集到一起。当前 query | v Lightning Indexer ---- Top-K token 位置 | v Host 侧完整 KV Cache ------- 选中的 KV ------- SFA该基础方案将完整 KV Cache 放置在 Host 内存中。Prefill 阶段模型在 Device 侧逐层生成 KV 后将其传输至 Hostdecode 阶段GatherSelectionKvCache算子根据 Indexer 的 Top-K 结果从 Host 侧完整 KV Cache 中加载对应的 KV 数据并交由 SFA 使用。这一流程释放了完整序列 KV Cache 此前占据的大部分 HBM腾出的空间可支撑更大的 batch 或更长的上下文。在仓库中GatherSelectionKvCache算子位于 ops/ascendc/src/gather_selection_kv_cache包含op_host算子原型、tiling与op_kernel内核实现含split_bs_reuse复用变体两部分并有独立的 torch 绑定 与 文档说明。在 GLM-5.2 样例中custom_op_loader.py 的load_offload_gather_op()会在 import 时把该算子挂载为torch_npu.npu_gather_selection_kv_cache。2.2 相邻 step 复用Gather 算子还利用了一个重要的局部性相邻 decode step 的 Top-K 往往高度重叠。上一个 step 已搬到 Device 侧的 KV 可直接复用无需再次从 Host 读取。全模型 Top-K 轨迹分析显示相邻 step 的平均复用率约为 60%。这意味着每次 decode约六成 KV 可从 Device 侧直接获得Host 只需提供剩余约四成。当前 Top-K 上一步选集中命中的部分 从 Host 侧完整 KV Cache 加载的缺失部分这一复用将 H2D 数据传输量降至全部从 Host 加载的约 40%但复用范围仅限于上一步的 2048 个位置。仍有近五分之二的选中 token 需要进行离散的 H2D 搬运——每次未命中对应一个独立位置的小段 KV 数据。2.3 容量收益与吞吐缺口该基础方案带来了明确的容量收益相同上下文长度下可承载的 batch size 约可翻倍。但 batch 扩大后每步需要同时处理的请求和数据规模随之增加叠加剩余的 Gather 搬运开销decode 延迟显著上升。在端到端评估中扩展的 batch 容量尚未充分转化为吞吐提升。这明确了下一步的优化目标提升缓存命中率将 H2D 数据搬运移出关键路径。三、更大的 KV 缓存池瓶颈从数据搬运转向缓存管理基础方案的选择缓冲区只有 2048 个条目因为它的容量等于 SFA 的 Top-K 输入量。如果进一步扩大复用范围让 Device 侧保留一个更大的 KV 缓存池——例如 8192 或 16384 个条目下一轮的 Top-K 就能以更高概率命中 Device 侧已有的数据。我们构建了一个包含 8192 个条目、采用 LRU 策略管理的 Device 侧 KV 缓存池。在相同 Top-K 轨迹下平均命中率从 58.85% 提升至 90.50%。平均而言仅有 9.50% 的选中 token 需要从 Host 加载。在这一命中率水平上H2D 数据搬运已不再是主要开销。但大容量缓存池改变了开销的构成。每层仍需独立处理自己的 Top-K id逐条判定数据位于缓存池还是 Host、对重复未命中去重、选定替换槽位、维护缓存元数据。LRU 的细粒度管理还需要维护访问时序和淘汰决策。这些操作多属于 Vector Core 不擅长的标量和控制密集型计算。KV 数据路径虽已通畅逐层重复的缓存管理操作却形成了新的 decode 瓶颈。Device 侧缓存设计平均命中率主导压力仅保留上一步的 2K Top-K 选集58.85%离散 H2D KV 加载8K KV 缓存池 LRU90.50%命中分类与缓存管理这一瓶颈迁移指明了后续工作的方向增大 Device 侧缓存池缓解了数据搬运问题但将缓存池管理做成高效算子还需要在控制面上寻找进一步的优化机会。GLM-5.2 的架构恰好提供了这样的切入点。四、GLM-5.2 共享 Indexer从逐层规划到逐组复用4.1 IndexShare四层一组的 Top-K 复用结构GLM-5.2 在 attention 设计中引入了一项关键特性共享 IndexerIndexShare。默认配置下一层独立执行 Indexer 生成新的 Top-K紧接着的三层直接复用这组结果构成四层一组的 Indexer 共享组。在仓库源码中这一结构由 configuration_glm.py 显式建模每层通过indexer_types标记为full本层运行 Lightning Indexer、持有 indexer 权重或shared复用上一个 full 层的 top-k、不持有 indexer 权重。若未显式指定则按如下规则推导# models/glm_5_2/models/configuration_glm.py # full iff (max(i - index_skip_topk_offset 1, 0) % index_topk_freq) 0 self.indexer_types [ full if (max(i - offset 1, 0) % freq) 0 else shared for i in range(num_hidden_layers) ]即默认表现为每 N 层一组、组首层为 full、其余层为 shared的周期结构。模型侧实现modeling_glm.py中shared层通过skip_topk indexer_types[layer_idx] shared跳过本层 Indexer 计算直接复用上一层输出的prev_topk_indices。四层共享同一组 Top-K id因此也共享了基于这些 id 的缓存决策。某个 token id 在组内任一层命中 Device 侧缓存池在其余三层也会映射到相同的缓存槽位某个 token id 未命中整组也需要执行相同的回填规划。命中分类和未命中规划是标量最密集、Vector Core 最不擅长的操作。共享 Indexer 结构将它们从逐层执行压缩为逐组执行——规划做一次四层复用。标量开销被分摊到四层而各层独立的 KV 数据搬运不受影响。4.2 共享规划规划器DsaPlan对共享 Top-K 执行一次性处理逐条记录数据来自缓存池还是 Host同时输出一份紧凑的未命中条目列表用于后续回填。组内后续层直接沿用同一份规划。各层拥有各自的 KV 和 RoPE 数据搬运过程独立进行但命中检测、来源分类、未命中去重和替换位置选择均不再重复计算。在仓库实现中这一一次规划、整组复用的语义由 offload_cache.py 的共享组元数据管理支撑build_dsa_shared_group_specs()按indexer_types把各层归入共享组并标记owner_layer组首 full 层规划结果通过set_dsa_shared_group_plan()/get_dsa_shared_group_plan()暂存组内共享层直接取用待组末层处理完毕后再clear_dsa_shared_group_plan()清除。4.3 组相联 FIFO大容量缓存池需要配套的替换策略。LRU 能够较好地匹配访问模式但在 NPU 上高效实现 LRU需要持续维护条目的最近访问状态和淘汰顺序执行开销较高。我们采用组相联 FIFO组织缓存池。所谓组相联set-associative是指先将缓存池划分为若干独立的组。每个 token id 通过哈希映射到唯一一组但可以存放在该组内任意一个槽位。相比每个 token 只能对应一个固定槽位的直接映射方式组相联可以降低哈希冲突相比允许条目存放在整个缓存池任意位置的全相联方式它又将查找和替换范围限制在较小的组内。每组包含固定数量的槽位替换仅发生在组内。每组维护一个 FIFO 指针仅在装入新条目时推进一次。当组内需要装入新条目时指针指向的槽位被替换随后指针移动到下一个槽位缓存命中不会改变该指针。因此实现上采用轮转替换语义上等价于按照装入顺序执行组内 FIFO。该方案在 NPU 上有三个适配性优势查找和替换状态被限制在较小的组内元数据开销可控不同组之间相互独立天然支持多核划分与并行处理FIFO 仅在新条目装入时更新无需在每次命中时维护访问时序。仓库中的常量定义印证了这一设计offload_cache.py默认池大小_DSA_DEFAULT_POOL_SIZE 8192enlarge_pool_size: True时切换到_DSA_LARGE_POOL_SIZE 16384组相联度_DSA_SET_ASSOCIATIVITY 16lru_counter的形状为(batch_size_per_rank, dsa_pool_size // 16)即每个 set 维护一个替换计数FIFO 指针同时以id_to_slot哈希表dsa_id_range不小于 131072完成 token id 到槽位的映射。缓存数据侧每个共享组在 Device 上维护独立的常驻池张量dsa_pool_key_valuesnope 与 rope 两部分shape 为(batch, pool_size, dim)见init_cache()。组相联在降低哈希冲突的同时将查找和替换范围限制在组内FIFO 则将替换机制做得足够轻量适配并行执行的要求。4.4 SFA 优先回填后置当前层注意力计算直接依赖选中的 KV——数据未就绪SFA 就无法启动。因此执行调度优先准备选中的 KV数据就绪后立即推进 SFA。缓存池回填的时限要宽松得多新装入的条目是供后续 decode step 使用的。回填因此在辅流上发起与 SFA 和 MoE 并行推进。共享 Indexer Top-K | v 共享规划 | ---- 逐层 KV 准备 ---- SFA ---- MoE | ---- 辅流缓存池回填这一调度将 Offload 的关键阻塞段压缩到准备选中 KV这一步缓存维护开销则尽可能与计算重叠不单独占用 decode 延迟。五、实现三个算子三条路径上述设计在代码层面映射为三个定制 AscendC 算子源码位于 ops/ascendc/src/dsa_plan、ops/ascendc/src/dsa_serve、ops/ascendc/src/dsa_install算子执行频率职责DsaPlan每个 Indexer 共享组一次分类 Top-K、判定缓存池命中或未命中、生成紧凑的未命中记录DsaServe每层一次从缓存池和 Host 侧完整 KV Cache 中准备 SFA 所需的 KV 和 RoPEDsaInstall每层一次回填未命中条目、执行 FIFO 淘汰、更新缓存池状态DsaPlan承载所有共享的标量工作将 Top-K id 和缓存池元数据转化为紧凑的执行规划。DsaServe依照规划执行数据搬运优先交付 SFA 所需的数据。DsaInstall依照未命中记录在辅流上执行延迟回填。三个算子的划分形成三条职责分明的执行路径规划在组级别完成不随层数放大数据服务每层独立进行但不再包含重复的决策逻辑回填与计算并行不占用关键延迟。5.1 模型侧编排以源码为证模型侧的编排逻辑集中在 shared_indexer_offload.py 的run_shared_indexer_offload()中与上述设计一一对应逐组规划offload_cache.is_dsa_shared_group_owner(layer_idx)判定当前层是否为组首 owner 层owner 层调用torch.ops.custom.dsa_plan生成plan与install_records并存入共享组非 owner 层直接复用get_dsa_shared_group_plan()取回的同一份规划逐层数据服务每层独立调用torch.ops.custom.dsa_serve将plan作用到本层的full_kv_cache/full_k_rope与缓存池pool_kv_view/pool_rope_view产出 SFA 所需的selection_kv_cache/selection_k_ropecompact_layout1表示池按[batch, pool_size, dim]紧凑布局存储辅流回填非 owner 层回填本层、组末层额外回填 owner 层的数据metadata_update仅在组末层触发用于推进 FIFO 指针并更新pool_ids/id_to_slot/lru_counterenable_install_stream开启时回填经npu_stream_switch切到独立辅流执行并通过record_event/wait_event与主计算流同步从而与 SFA、MoE 计算重叠。在 eager / npugraph_ex / ge_graph 三种执行模式下算子加载路径由 custom_op_loader.py 统一处理load_offload_gather_op()挂载npu_gather_selection_kv_cache基础 Offload 算子register_dsa_shared_ge_converters()ge_graph 模式注册 DsaPlan/DsaServe/DsaInstall 的 torchair GE converter含dsa_functionalization辅助模块register_dsa_shared_npugraph_reinplace()npugraph_ex 模式注册dsa_serve_functional/dsa_serve、dsa_install_functional/dsa_install的 inplace 配对使图捕获阶段能够正确执行 host-device 数据搬运。Runner 侧runner_glm.py在shared_indexer_offload开启时按exe_mode选择上述加载路径。此外offload_cache.py 还实现了 MTP 场景的 gather 复用enable_mtp_gather_reuseMTP draft 步冻结首个 decode step 的 top-k 后其选中的历史位置 KV 不可变首个 step 收集到的 selection buffer 在后续 draft 步中保持有效可跳过冗余的重复 gather。六、在仓库中启用共享 Indexer Offload6.1 依赖的定制算子编译安装KV Offload 依赖仓内自定义 AscendC 算子非 torch_npu 内置须先编译安装详见 models/glm_5_2/README.md# 1) 编译算子内核A3 默认 ascend910_93bisheng 随 CANN 提供 cd ops/ascendc bash build.sh -n gather_selection_kv_cache;dsa_plan;dsa_serve;dsa_install # 安装包位于 output/CANN-custom_ops-*-linux.arch.run # 2) 安装内核到 CANN opp cd output ./CANN-custom_ops-*-linux.*.run --quiet --install-path${ASCEND_HOME_PATH}/opp # 3) 编译并安装 torch 绑定缺 ninja 时先 pip install ninja cd ../torch_ops_extension bash build_and_install.sh # 生成并 pip 安装 custom_ops wheel # 4) 运行前 source 自定义算子环境设置 ASCEND_CUSTOM_OPP_PATH / LD_LIBRARY_PATH source ${ASCEND_HOME_PATH}/opp/vendors/customize/bin/set_env.bashmodeling_glm.py在 import 时自动把基础 offload 算子挂到torch_npu.npu_gather_selection_kv_cache开启shared_indexer_offload时Runner 初始化会按需加载共享 IndexShare offload 相关自定义算子。因此模型侧无需改动只要上面的算子已安装、且运行前 source 了第 4 步的环境即可。若遇到_OpNamespace custom object has no attribute类报错可参考 ops/ascendc/README.md 重新编译对应算子。6.2 配置项offload 三开关与执行模式仓库提供了两份 offload 示例配置models/glm_5_2/configglm_5_2_rank_32_32ep_w8a8_offload.yaml基础 KV Offload 路径glm_5_2_rank_32_32ep_w8a8_offload_mtp.yaml开启共享 Indexer Offload 与 MTP3 的完整路径。在model_config中设置 offload 选项详细参数说明见 config/README.mdmodel_config: enable_offload: True # [False, True] 开启 KV Offload长序列、大 batch 场景 shared_indexer_offload: True # [False, True] 使用 GLM-5.2 IndexShare 共享规划路径 enlarge_pool_size: False # [False, True] 设备侧常驻 token 池从 8K 扩大到 16K next_n: 3 # MTP 步数支持 [0, 1, 2, 3] pa_block_size: 128 # PagedAttention 块大小支持 [128, 256] with_ckpt: True # 是否加载权重 enable_multi_streams: True # 多流并行 enable_online_split_weight: True # 在线切分权重 exe_mode: npugraph_ex # [ge_graph, eager, npugraph_ex]decode 执行模式三个开关的语义与作用enable_offload: True把全量 MLA KV 卸载到 Host swapped memorydecode 时按 DSA top-k 把命中的 block gather 回 Device面向全量 KV 放不下 HBM 的长上下文场景shared_indexer_offload: True利用 GLM-5.2 的 IndexShare 特性复用同一组 top-k 规划减少共享层重复的命中判断和缓存管理开销并将常驻池回填与 SFA、MoE 计算并行enlarge_pool_size: True把设备侧常驻 token 池从 8K 扩大到 16K进一步提高命中率并减少 host 到 device 的 KV 搬运。该开关经 offload_cache.py 的resolve_dsa_pool_size()生效8192 → 16384。配套的并行配置参考A3 单卡双 die、world_size 2 * chip_numembed_tp_size/lmhead_tp_size切到 32W8A8 权重下 embedding 与 lm_head 走 TPattn_tp_size/dense_tp_size/moe_tp_size为 1MoE 经 32 路 EP 展开。数据侧示例input_max_len: 4096、batch_size: 32可按评测需求调整。6.3 运行流程各节点同步执行 infer.sh 即可拉起多卡推理# infer.sh 中指定配置文件 export YAML_FILE_NAMEglm_5_2_rank_32_32ep_w8a8_offload_mtp.yaml bash infer.sh运行前需在 executor/scripts/set_env.sh 中填写各节点 IP第 1 个为 master与 CANN 包路径并在 yaml 中把model_path指向转换好的 W8A8 权重转换命令见 models/glm_5_2/README.md 的weight_convert.sh部分。七、双机 Atlas A3 评测7.1 评测环境与口径评测环境为生产级配置双机 Atlas A3共 16 卡、GLM-5.2-W8A8 模型、64K 输入序列、MTP3。针对 64K 输入序列在相同设备资源约束下不使用 Offload 时KV Cache 容量仅能支持全局 batch size 上限为 64启用 Offload 后该上限提升至 128。实验分别测试 8K 和 16K 两种 Device 侧 KV 缓存池容量并覆盖 GE Graph 和 NPUGraph Ex 两种执行模式。吞吐按两种口径报告统一假设口径为在相同 MTP 接收长度假设下比较不同配置将 MTP 平均接收长度统一设为 2.7。该数值为经验假设仅用于归一化估算不代表各配置的实测接收长度实测口径使用推理过程中实际观测到的 MTP 平均接收长度。计算方式吞吐 全局 batch size × MTP 平均接收长度 /主模型 MTP3 单步延迟7.2 完整评测数据执行模式配置全局 batch缓存池容量主MTP3 延迟吞吐2.7 假设吞吐实测GE Graph无 Offload64-71.21 ms2,427 token/s2,504 token/sGE Graph共享 Indexer Offload1288K100.55 ms3,437 token/s41.6%3,642 token/s45.4%GE Graph共享 Indexer Offload12816K92.03 ms3,755 token/s54.8%3,982 token/s59.0%NPUGraph Ex无 Offload64-71.59 ms2,414 token/s2,507 token/sNPUGraph Ex共享 Indexer Offload1288K105.04 ms3,290 token/s36.3%3,452 token/s37.7%NPUGraph Ex共享 Indexer Offload12816K95.67 ms3,612 token/s49.6%3,791 token/s51.2%7.3 从数据中归纳的三点结论可承载的 batch 上限翻倍延迟增长可控。针对 64K 输入序列Offload 将可承载的全局 batch size 上限从 64 提升至 128。16K 缓存池下主模型 MTP3 单步延迟相较无 Offload 配置增长 29.2%GE Graph和 33.6%NPUGraph Ex。延迟增长幅度远低于 batch 上限的增幅表明 Offload 的附加开销处于有效控制之下。吞吐净增约 50%。16K 缓存池下按统一的 2.7 经验假设估算GE Graph 和 NPUGraph Ex 的吞吐分别提升 54.8% 和 49.6%按各配置的实测接收长度计算吞吐分别提升 59.0% 和 51.2%。batch 容量的扩大切实转化为了系统吞吐收益。16K 缓存池优于 8K 缓存池。在最终的组相联 FIFO 方案下缓存池容量从 8K 扩大至 16K 后平均命中率由约 85% 提升至 90% 以上需要从 Host 加载并回填的未命中数据随之减少。在相同全局 batch size 128 下16K 缓存池相较 8K 缓存池吞吐进一步提升约 9.3%GE Graph和 9.8%NPUGraph Ex按 2.7 MTP 接受长度经验假设估算。这表明扩大缓存池容量所带来的命中率提升可以进一步降低 Offload 的端到端开销。这些对比印证了 Offload 设计的核心判断Host 内存用于承载完整 KV Cache从而提高可支持的全局 batch 上限Device 侧缓存池、跨层复用和调度优化则用于压低附加开销。两者协同使 batch 上限翻倍最终转化为约 50% 的吞吐净增。八、扩展思考KV Cache 管理的系统化方向GLM-5.2 的共享 Indexer 为本次 Offload 方案提供了关键的结构性支撑。四层一组的 Top-K 复用使标量规划开销得以分摊这一点离开了模型的特定架构是无法成立的。换言之KV Offload 的上限很大程度上由模型结构决定。不同模型的注意力设计、层间关联、计算与带宽的配比各不相同缓存大小、替换策略和执行调度的最优选择也会随之变化。从更长远的视角看KV Cache 的管理最终需要走向一个系统化的方向。Offload 解决了容量问题但单靠 Offload 还不足以应对上下文持续增长带来的全方位压力。一个更完整的思路是将 KV Cache 视为可分级、可调度的资源在 Host 内存与 Device HBM 之间建立缓存层次依据访问热度动态调配数据位置结合淘汰策略自动适应工作负载变化。与此同时稀疏注意力压缩每步计算量投机解码通过并行生成均摊延迟——三者各自有效但要系统性地应对长上下文下的吞吐和延迟挑战需要将它们协同部署每项技术的短板恰可由其他技术补足。参考本文涉及的仓库关键路径技术文档本体docs/models/glm_5_2/glm_5_2_offload_guide.md相关图示位于 docs/models/glm_5_2/figuresGLM-5.2 样例说明与算子编译步骤models/glm_5_2/README.md配置示例glm_5_2_rank_32_32ep_w8a8_offload.yaml、glm_5_2_rank_32_32ep_w8a8_offload_mtp.yaml、config/README.md模型侧编排shared_indexer_offload.py、offload_cache.py、custom_op_loader.py、configuration_glm.py、runner_glm.py定制算子ops/ascendc/src/gather_selection_kv_cache、ops/ascendc/src/dsa_plan、ops/ascendc/src/dsa_serve、ops/ascendc/src/dsa_install启动脚本models/glm_5_2/infer.sh、executor/scripts/set_env.sh【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考