昇腾 NPU 独立部署模式 KVCache 自管完整指南：Legacy 连续缓存、手动 Paged 与 MLA 压缩缓存的 Runner 实现

昇腾 NPU 独立部署模式 KVCache 自管完整指南Legacy 连续缓存、手动 Paged 与 MLA 压缩缓存的 Runner 实现【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer本文档面向在 CANN 昇腾 NPU 平台上采用独立部署模式Runner 自管推理循环、不接入executor/core/框架调度的 LLM 推理场景完整讲解 Runner 自管 KVCache 数据结构、block_table/slot_mapping/forward_metadata构造与 modeling 层 FA 融合算子接入的三种骨架连续缓存Legacy、手动 Paged 模式、MLA 压缩缓存。读完本文你将掌握在 runner 脚本中从零分配 KV tensor、按 Prefill / Decode 动态构造写入索引与 FA 入参、并正确驱动kv_len生命周期的完整实现方案。背景独立部署模式与框架部署模式的边界model-infer-kvcache技能将推理部署划分为两种模式见 SKILL.md框架部署模式模型接入executor/core/调度 / 批组装 / KV 管理由框架统一负责KV 路径由模型类通过get_cache_info()cache_entries声明BlockPool动态分配 blockslot_mapping/block_table由executor/core/kv_cache/自动构造独立部署模式Runner 自管 KV tensor 分配、block_table静态预分配、slot_mapping动态计算并自管forward_metadata构造。modeling 内scatter_update_写入 FA 算子调用与框架部署完全一致。判定方式模型注册到executor/core/support_models.py为框架部署模型目录带runner_*.py且不继承 framework 基类、自管推理循环仓内如 models/pangu_7b/runner_openpangu_dense.py、models/glm_5_2/runner_glm.py、models/kimi_k3/runner_kimi_k3.py则为独立部署。本文按 KV 模式分三类骨架连续缓存Legacy、手动 Paged 模式、MLA 压缩缓存可叠加在前两者之上。三者的核心差异在于谁负责构造block_table/slot_mapping/actual_seq_lengths框架部署由框架构造独立部署全部由 Runner 在推理循环内自管。模式一连续缓存Legacy最简KV 以连续 tensor 存储scatter_update_写入FA 直接读取整个缓存。无 paging 概念不需要block_table/slot_mapping。该骨架形态由 model-infer-migrator skill 在独立部署阶段产出是 migrator 阶段最简起点。Runner 内 KV 分配# runner_{model_name}.py 内 _init_kvcache(...) def _init_kvcache(self, num_layers, batch_size, max_seq_len, num_kv_heads_per_rank, head_dim, dtype): self.kv_caches [] # list of (k, v) for _ in range(num_layers): k torch.zeros(batch_size, max_seq_len, num_kv_heads_per_rank * head_dim, dtypedtype, deviceself.device) # BSH layout v torch.zeros_like(k) self.kv_caches.append((k, v))每层分配一组(k, v)shape 为[batch_size, max_seq_len, num_kv_heads_per_rank * head_dim]即 BSH layoutB 为 batch、S 为序列、H 为展平的 head 维度。由于无 paging缓存容量按batch × max_seq_len一次性静态铺满。modeling 写入 FA 调用# attention layer class XxxAttention(nn.Module): def forward(self, hidden_states, kv_len, attention_mask, past_kv, ...): # ... QKV projection、RoPE ... past_key, past_value past_kv # 来自 Runner 注入 # 写入migrator 阶段 BSH legacy 骨架scatter_update_ 按 kv_len 写位置BSH 时 axis1 torch_npu.scatter_update_(past_key, kv_len, key_states, 1) torch_npu.scatter_update_(past_value, kv_len, value_states, 1) # FA 调用FA v1 示例 attn_output, _ torch.ops.npu.npu_fused_infer_attention_score( query_states, past_key, past_value, num_headsnum_heads, num_key_value_headsnum_kv_heads, input_layoutBSH, scale1.0 / math.sqrt(head_dim), actual_seq_lengths_kvactual_seq_lengths_kv, atten_maskattention_mask, sparse_mode0 if not is_prefill else 3, )改造边界只动 attention 计算→ FA和 cache 存储→scatter_update_上游QKV projection / RoPE保持不变layout 不匹配时在接缝处 transpose / reshape 适配。写入算子选择scatter_update_是 BSH Legacy 骨架专用migrator 阶段最简进入 Paged 模式模式二后必须切到npu_scatter_nd_update_接slot_mapping或npu_kv_rmsnorm_rope_cache融合写入。注意Legacy 模式仅支持 offline 推理。若要使用 online / PD 服务或框架数据集评测能力必须升级到 Paged 模式框架部署下实现get_cache_info()接入 KVCacheManager独立部署下采用本文模式二。FA 调用完整代码示例含 GPT-OSS / Qwen3-MoE / DeepSeek-R1 / Kimi-K2 / LongCat-Flash 五种见 references/fa-code-examples.md。模式二手动 Paged 模式 FAKV 按固定大小 Block 存储FA 通过block_table索引分块缓存。Paged 模式必须配合 FA 使用无法走标准 softmax。这是独立部署的性能主路径也是框架部署 Paged Attention 机制的 Runner 侧等价实现。Runner 内 KV block_table 静态预分配def _init_kvcache(self, num_layers, batch_size, max_seq_len, num_kv_heads_per_rank, head_dim, dtype, block_size128): self.block_size block_size num_blocks_per_seq max_seq_len // block_size total_blocks batch_size * num_blocks_per_seq # 物理 cache[total_blocks, block_size, num_kv_heads, head_dim] self.kv_caches [] for _ in range(num_layers): k torch.zeros(total_blocks, block_size, num_kv_heads_per_rank, head_dim, dtypedtype, deviceself.device) v torch.zeros_like(k) self.kv_caches.append((k, v)) # block_table 静态预分配推理全程不变 # block_table[b, i] 第 b 个 batch 的第 i 个逻辑 block 对应的物理 block ID self.block_table torch.arange(0, total_blocks).reshape( batch_size, -1).to(torch.int32).npu() # 预计算 kv_len_offsetslot_mapping 用 self.kv_len_offset torch.arange( 0, batch_size * max_seq_len, max_seq_len, deviceself.device, ).view(-1, 1)物理存储 shape 为[total_blocks, block_size, num_kv_heads_per_rank, head_dim]。分页注意力把逻辑地址(batch_idx, seq_pos)映射到物理 block映射公式为逻辑 block 编号 seq_pos // block_size block 内偏移 seq_pos % block_size 物理 block ID block_table[batch_idx, 逻辑 block 编号] 物理 slot 位置 物理 block ID × block_size block 内偏移block_tableshape[batch_size, num_blocks_per_seq]dtypeint32推理全程不变。独立部署 Runner 用arange(total_blocks).reshape(batch_size, -1)一次性静态预分配并持有self.block_table传给 modeling。Runner 内 slot_mapping 动态计算slot_mapping是缓存写入位置索引给npu_kv_rmsnorm_rope_cache/npu_scatter_nd_update_等写入算子用。独立部署 Runner 通常走 block_table 静态预分配arange模式slot 公式简化为slot(batch_idx, seq_pos) batch_idx × max_seq_len seq_pos 顺序分配模式下等于展平后线性索引框架部署模式下 BlockPool 动态分配 blockslot 计算依赖block_table索引block_id × block_size offset不适用以上静态预分配公式。框架侧实现在 executor/core/kv_cache/cache_utils.py 的prepare_slot_mapping()中按position_ids切分每个 batch用block_indices tmp_position_ids // cur_block_size、position_offsets tmp_position_ids % cur_block_size计算slot_mapping block_ids * cur_block_size position_offsets。两种公式在 block_table 恰为 arange 时等价。Prefill—— 每个 batch 写入多个 tokendef _build_slot_mapping_prefill(self, kv_len, max_seq_len): # kv_len: [batch_size]每条请求的 prompt 长度 all_tensors [] for i, seq_len in enumerate(kv_len): all_tensors.append(torch.arange( max_seq_len * i, seq_len.item() max_seq_len * i, dtypetorch.int32, deviceself.device, )) return torch.cat(all_tensors) # 示例kv_len[512, 256], max_seq_len2048 # batch 0: [0, 1, ..., 511] # batch 1: [2048, 2049, ..., 2303] # 拼接 → shape[768] 一维 tensorDecode—— 每个 batch 写入 1 个 tokendef _build_slot_mapping_decode(self, kv_len): # kv_len: [batch_size]每条请求当前 KV 实际长度 return kv_len.view(-1, 1) self.kv_len_offset # [batch_size, 1] # 示例kv_len[522, 266], kv_len_offset[[0], [2048]] # slot_mapping [[522], [2314]]slot_mapping 与 block_table 的分工slot_mapping供缓存写入算子npu_kv_rmsnorm_rope_cache/npu_scatter_nd_update_使用block_table供 FA注意力读取算子npu_fused_infer_attention_score{,_v2}使用二者寻址逻辑一致、职责不同。actual_seq_lengths 构造FA 算子需要每个 batch 的实际 KV/Q 长度构造方式取决于input_layout| | TND layout多 batch token 拼一维 | BSH layout各 batch 独立 | |--|------|------| |Prefill KV|cumsum(kv_len)→ [512, 768] |kv_len→ [512, 256] | |Prefill Q| 同 KV | 同 KV | |Decode KV|kv_len→ [522, 266] |kv_len→ [522, 266] | |Decode Q|cumsum([1,1])→ [1, 2] |[1, 1]|# TND Prefill 示例 actual_seq_lengths_kv torch.cumsum(kv_len, dim0) actual_seq_lengths_q actual_seq_lengths_kv.clone()在框架部署模式下ForwardMetaData提供两组字段见 executor/utils/forward_metadata.pyactual_seq_lengths_kv/actual_seq_lengths_q为直接形态actual_seq_lengths_cu_kv/actual_seq_lengths_cu_q为 cumulative 形态。FA 调用 TND layout 时actual_seq_qlen取 cu 形态、actual_seq_kvlen取直接形态。modeling 写入 FA 调用class XxxAttention(nn.Module): def forward(self, hidden_states, slot_mapping, block_table, ...): # 写入用 npu_scatter_nd_update_slot_mapping 索引模式cache / states 都 view 为 [total, num_kv_heads, head_dim] torch_npu.npu_scatter_nd_update_( self.k_cache.view(-1, num_kv_heads, head_dim), slot_mapping.view(-1, 1), key_states.view(-1, num_kv_heads, head_dim), ) # 或融合写入MLA 模型推荐一步完成 RMSNorm RoPE Cache 写入 # torch_npu.npu_kv_rmsnorm_rope_cache(..., slot_mapping.view(-1), ...) # FA 调用FA v2 示例传 block_table attn_output, _ torch_npu.npu_fused_infer_attention_score_v2( query_states, self.k_cache, self.v_cache, block_tableblock_table, block_sizeself.block_size, num_query_headsnum_heads, num_key_value_headsnum_kv_heads, actual_seq_kvlenactual_seq_lengths_kv, actual_seq_qlenactual_seq_lengths_q, input_layoutTND_NTD, sparse_mode0 if not is_prefill else 3, softmax_scale1.0 / math.sqrt(head_dim), )仓内 Paged 模式的标准实现可对照两类参考标准 LLM TND FA v1Qwen3-MoEmodels/qwen3_moe/models/modeling_qwen3_moe.py 的exec_qkv函数Prefill Decode 统一sparse_mode3 TND block_table滑窗 sink FA v2GPT-OSSmodels/gpt_oss/models/modeling_gpt_oss.py滑窗层用sparse_mode4pre_tokenssliding_windownext_tokens0Prefill 直读 KV不传block_table、Decode 切到block_table路径。Paged 模式 KV cache 实际 shape 与 input_layout 是两回事input_layout字符串描述 FA 算子对各张量的解释方式paged KV 的实际 tensor shape 与之不同。GQA / MHA 的 cache 逻辑 shape 为[bn, bs, num_kv_heads_per_rank, head_dim]传给 FA 前view(*shape[:2], -1)成 3D[bn, bs, H]MLA 的 nope_cache / rope_cache 逻辑 shape 为[bn, bs, 1, D]写入算子cache_modePA_NZ实际按 NZ 排布FA 调用前 view 为 5D NZNZ_DIM16for bf16int8 量化时 ×2参考 models/deepseek_r1/models/modeling_deepseek.py 中KV_CACHE_NZ_DIM相关实现。FA v1 / v2 关键参数名映射混用不报错会静默落到算子默认值导致精度异常功能FA v1FA v2易错点缩放系数scalesoftmax_scale默认 1.0传错名精度崩溃Q head 数num_headsnum_query_heads传错名走默认值Q 长度actual_seq_lengthsactual_seq_qlenv1 BSH Q_S1 路径下算子内部忽略仍需传KV 长度actual_seq_lengths_kvactual_seq_kvlen名称完全不同KV head 数num_key_value_headsnum_key_value_heads相同sparse_mode 与 atten_mask 组合两者组合错误是精度问题的第一大来源sparse_mode含义atten_mask 要求适用场景0Dense可选通常传NoneMLA absorb Decodeq_len1 单 token由actual_seq_lengths_kv控制有效长度非 MLA 路径不推荐1allMask必传完整矩阵(Q_S, KV_S)特殊场景2leftUpCausal不推荐建议改用 3—3Causal标准因果必传[2048, 2048]bool 下三角标准 LLM TND PAPrefill Decode 统一、MLA absorb Prefill、MTP Decodesq1无滑窗4Band滑动窗口必传[2048, 2048]bool滑窗模型gpt-oss sliding 层 Prefill Decode 统一需配合pre_tokensatten_mask硬约束dtype 只允许torch.bool推荐、torch.int8、torch.uint8浮点类型直接报错shape 在sparse_mode3/4时固定[2048, 2048]与max_position_embeddings无关。本仓库统一用 executor/utils/common_utils.py 的get_init_attn_mask(2048, device)构造~torch.tril(torch.ones((2048, 2048), dtypetorch.bool, devicedevice))返回 bool 上三角为 True 的因果 mask。模式三MLA 压缩缓存叠加在模式一或模式二之上MLAMulti-head Latent Attention将 KV 压缩为低维 latent只缓存压缩后的cache_nope非位置和cache_rope位置两个分量大幅降低显存占用与传输量。Runner 内 KV 分配def _init_kvcache(self, num_layers, batch_size, max_seq_len, kv_lora_rank, qk_rope_head_dim, dtype, block_size128): # MLA 压缩维度kv_lora_rank512远小于完整 KVnum_heads * head_dim16384 num_blocks_per_seq max_seq_len // block_size total_blocks batch_size * num_blocks_per_seq self.cache_nope [] # 非位置部分 self.cache_rope [] # 位置编码部分 for _ in range(num_layers): c_nope torch.zeros(total_blocks, block_size, 1, kv_lora_rank, dtypedtype, deviceself.device) c_rope torch.zeros(total_blocks, block_size, 1, qk_rope_head_dim, dtypedtype, deviceself.device) self.cache_nope.append(c_nope) self.cache_rope.append(c_rope)仓内 MLA 模型的默认维度可从配置确认DeepSeek-R1 的 models/deepseek_r1/models/configuration_deepseek.py 与 Bailing 2.5 的 models/bailing_2_5/models/configuration_bailing_moe_v2_5.py 均设kv_lora_rank512。latent KV 共享不沿attn_tp_size切分每 attn_tp rank 持完整 latent KVQ 按 head 切、attn_dp 按 batch 切。modeling FA absorb 调用# key 和 value 传同一个 cache_nopeabsorb 技术将 V 投影吸收到 O 投影中 attn_output, _ torch_npu.npu_fused_infer_attention_score_v2( q_nope, k_nope_cache, k_nope_cache, # key value cache_nope query_ropeq_pe, key_ropek_rope_cache, # RoPE 单独传入 block_tableblock_table, block_sizeself.block_size, ... )重要约束query_rope和key_rope必须同时传或同时不传rope D 必须为 64MLA query D 仅支持 512 或 128MLA D512 时仅支持sparse_mode为 0、3、4FA 调用必须num_key_value_heads 1latent 单 head与cache_entries.num_head 1严格一致不一致会触发 GQA 比例错误仓内参考实现models/deepseek_r1/models/modeling_deepseek.py 的forward_absorb路径完整展示了 MLA 全链路融合写入npu_kv_rmsnorm_rope_cache_v2(latent_cache, kv_a_layernorm.weight, cos, sin, slot_mapping, rope_cache, nope_cache, epsilon..., cache_modePA_NZ, is_output_kvTrue)一步完成 RMSNorm RoPE Cache 写入输出k_rope/k_nopeabsorb 投影k_nope_out matmul(k_nope.view(1, -1, kv_lora_rank), kv_b_proj_w_k.permute(0, 2, 1))、v_out matmul(k_nope, kv_b_proj_w_v)V 由 latent 经kv_b_proj重投影获得自动消除qk_head_dim ≠ v_head_dim时 V 需 pad 到 qk_head_dim 的问题FA 调用keyvalue 传 latent 派生张量query_rope/key_rope分离传入。GLM-5.2 的 models/glm_5_2/models/modeling_glm.py 同样提供forward_absorb/forward_absorb_cp实现其函数签名直接接收kv_len、actual_seq_lengths_kv、actual_seq_lengths_q、is_prefill、slot_mapping等 Runner 侧注入的参数与本文骨架完全对应。完整 Runner 改造与 forward_metadata 自管独立部署 Runner 在每步推理前需自管构造forward_metadatadef _build_forward_metadata(self, kv_len, is_prefill, max_seq_len): # 直接复用 from executor.utils.forward_metadata import ForwardMetaData # 或参考字段在 runner 内 inline 一个轻量 dataclass if is_prefill: slot_mapping self._build_slot_mapping_prefill(kv_len, max_seq_len) actual_seq_lengths_kv torch.cumsum(kv_len, dim0).to(torch.int32) actual_seq_lengths_q actual_seq_lengths_kv.clone() # Prefill 用 [2048, 2048] bool 因果 mask参考 executor.utils.common_utils.get_init_attn_mask attention_mask self.share_mask_tril else: slot_mapping self._build_slot_mapping_decode(kv_len) actual_seq_lengths_kv kv_len.to(torch.int32) actual_seq_lengths_q torch.arange(1, kv_len.shape[0] 1).to(torch.int32) attention_mask None # full-attention Decode 不需要 mask return ForwardMetaData( is_prefillis_prefill, kv_lenkv_len, slot_mappingslot_mapping, # 单 attn_type 时传 tensor多 attn_type 时按 dict 包装 block_tableself.block_table, # 同上 actual_seq_lengths_kvactual_seq_lengths_kv, actual_seq_lengths_qactual_seq_lengths_q, attention_maskattention_mask, )仓内的 executor/utils/forward_metadata.py 提供了完整的ForwardMetaDatadataclass 定义可直接 import 复用除上述字段外还包括is_warm_up、actual_seq_lengths_cu_kv/actual_seq_lengths_cu_qcumulative 形态、actual_seq_lengths_cu_list_kv/actual_seq_lengths_cu_list_q/actual_seq_lengths_list_kv/actual_seq_lengths_list_qlist 形态、prompt_tokens、cp_metadata上下文并行元数据。模块还提供全局单例get_forward_metadata()/set_forward_metadata(**kwargs)/reset_forward_metadata()供框架部署模式使用独立部署 Runner 可直接构造实例或参考字段 inline 一个轻量 dataclass。kv_len 生命周期kv_len是驱动所有入参计算的核心变量# Prefill从 attention_mask 计算 position_ids attention_mask.long().cumsum(-1) - 1 kv_len torch.max(position_ids, axis1)[0] 1 # 如 [512, 256] # 每次 Decode 后递增在 Runner 层model.forward() 之外 kv_len kv_len 1 # [513, 257], [514, 258], ... # kv_len 驱动三个计算 # 1. slot_mapping → 决定新 token 写到哪 # 2. actual_seq_lengths_kv → 告诉 FA 读多少 KV # 3. position_ids → RoPE 位置编码kv_len是 Runner 层变量各层只读不写。Runner 计算一次kv_len及由它派生的slot_mapping所有 Transformer 层用同一份值写入各自缓存模式一用scatter_update_接kv_len/ 模式二/三用npu_scatter_nd_update_或npu_kv_rmsnorm_rope_cache接slot_mapping。不要在 attention 内部递增kv_len——否则 Prefill 阶段各层写入位置逐层偏移精度损坏。每步推理的数据流可概括为初始化一次性block_table、kv_cache、kv_len_offset 每步 Prefill / Decode kv_len → slot_mapping → 写 cachenpu_scatter_nd_update_ / npu_kv_rmsnorm_rope_cache kv_len → actual_seq_lengths_kv → FA 读 cache via block_table kv_len → position_ids → RoPE Decode 后kv_len 1独立部署模式的验证与常见错误独立部署改造完成后按以下链路验证Prefill 输出与基线对比migrator 阶段已建立的基线Decode 单步耗时合理与基线对比无显著回归多卡场景各 rank 输出形状一致KV 切分无错位长跑无显存泄漏。独立部署模式的高频错误与预防完整错误表见 SKILL.md 的常见错误一节错误模式根因预防Prefill 输出乱码但不报错sparse_mode0maskNone无因果遮蔽改用sparse_mode3[2048, 2048]maskMLA absorb Decode 输出严重偏差误用sparse_mode3改sparse_mode0、maskNone仅 MLA absorb Decode 路径切此配置标准 LLM TND PA PrefillDecode 统一用 sparse_mode3滑窗 Decode 长序列输出偏差误把滑窗层改成sparse_mode0导致看见窗口外 KV保留sparse_mode4设置pre_tokenssliding_window,next_tokens0atten_mask dtype报错mask 用了 float16/bfloat16mask.to(torch.bool)atten_mask shape报错mask 不是[2048, 2048]固定用executor.utils.common_utils.get_init_attn_mask(2048, device)scale默认 1.0 导致精度崩溃FA v1 用scalev2 用softmax_scale传错名静默生效确认参数名与 FA 版本匹配各层共享kv_len但内部递增Prefill 各层写入位置逐层偏移kv_len是 Runner 层变量各层只读不写参考实现索引实现模式参考文件搜索关键词独立部署 Runner 形态models/pangu_7b/runner_openpangu_dense.py、models/glm_5_2/runner_glm.py、models/kimi_k3/runner_kimi_k3.py_init_kvcache、forward_metadataMLA absorb 完整实现models/deepseek_r1/models/modeling_deepseek.py、models/glm_5_2/models/modeling_glm.pyforward_absorb、npu_kv_rmsnorm_rope_cache、PA_NZ标准 LLM TND PagedFA v1models/qwen3_moe/models/modeling_qwen3_moe.pysparse_mode3、block_table滑窗 sinkFA v2models/gpt_oss/models/modeling_gpt_oss.pysliding_window、learnable_sink框架部署 cache 工具对照参考executor/core/kv_cache/cache_utils.pyprepare_block_tables、prepare_slot_mapping、calculate_block_numForwardMetaData 定义executor/utils/forward_metadata.pyForwardMetaData因果 mask 构造executor/utils/common_utils.pyget_init_attn_maskFA 完整调用示例.agents/skills/model-infer-kvcache/references/fa-code-examples.mdnpu_fused_infer_attention_scoreKVCache 技能总纲选型与决策流程.agents/skills/model-infer-kvcache/SKILL.md部署模式判定、第一层快速选型【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考