CANN/cannbot-skills KVCache优化技能

【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skillsname: model-infer-kvcache description: 基于 PyTorch 框架的昇腾 NPU 模型推理 KVCache 优化技能。分析和优化 LLM 文本推理模型中的 KVCache 实现包括连续缓存、分页注意力Paged Attention配合 FA 融合算子、MLA 压缩缓存。触发场景包括KVCache 管理实现、paged attention、KV 压缩、FA 融合算子、OOM/性能问题、block_table/slot_mapping 构造。基于本仓库已有模型的 KVCache 实现经验按模型类型和场景推荐最佳方案。KVCache 优化技能本技能按渐进式披露组织快速选型 → 实现模式 → 深入原理 → 参考附录。 根据需要的深度选择阅读层级。第一层快速选型根据模型类型选择 KVCache 模式和参考实现模型类型选择模式参考实现要点标准 LLMMHA/GQA模式一连续缓存gpt-oss, qwen3-moe最简单无需 block_table高性能 LLM模式二分页注意力 FAdeepseek-r1, kimi-k2需构造 block_table slot_mappingMLA 架构DeepSeek 系列模式三MLA 压缩deepseek-r1, deepseek-v3.2叠加在模式一或模式二之上扩散/视频模型FA 直接计算hunyuan-video, wan2.2-i2v无 KVCache 分页BNSD layout决策路径前置确认从 progress.md 读取 Phase 0 已验证的架构类型基于实际 config 值非代码推断 │ 需要分页注意力 ──→ 否 → 模式一连续缓存scatter_update_ FA └→ 是 → 模式二分页注意力block_table slot_mapping FA │ 使用 MLA ──────→ 是 → 叠加模式三cache_nope cache_rope 压缩存储第二层实现模式模式一连续缓存非分页注意力KV 以连续 tensor 存储scatter_update_写入FA 直接读取整个缓存。三步实现# 1. 初始化缓存BSH layout 示例 past_key torch.zeros(batch_size, max_seq_len, num_kv_heads * head_dim, dtypedtype, devicenpu) past_value torch.zeros_like(past_key) # 2. 写入缓存每步调用 torch_npu.scatter_update_(past_key, kv_len, key_states, 1) # axis1 for BSH torch_npu.scatter_update_(past_value, kv_len, value_states, 1) # 3. FA 注意力计算FA v1 示例不传 block_table attn_output, _ torch.ops.npu.npu_fused_infer_attention_score( query_states, past_key_states, past_value_states, num_headsnum_heads, num_key_value_headsnum_kv_heads, input_layoutBSH, scale1.0 / math.sqrt(head_dim), actual_seq_lengths_kvactual_seq_lengths_kv, # [batch_size] 每个 batch 的 KV 长度 atten_maskattention_mask, sparse_mode0, # Decode; Prefill 用 3(causal) )改造边界只改 attention 计算→FA和 cache 存储→scatter_update_。上游QKV projection、RoPE 等保持不变layout 不匹配时在接缝处 transpose/reshape 适配。参考文件cann-recipes-infer/models/qwen3_moe/models/modeling_qwen3_moe.pyBSH、cann-recipes-infer/models/gpt_oss/models/modeling_gpt_oss.pyTND模式二分页注意力 FA 融合算子KV 按固定大小 Block 存储FA 通过block_table索引分块缓存。PA 必须配合 FA 使用。四步实现# 1. 初始化缓存和 block_table一次性推理全程不变 block_size 128 num_blocks_per_seq max_seq_len // block_size total_blocks batch_size * num_blocks_per_seq kv_cache torch.zeros(total_blocks, block_size, num_kv_heads, head_dim, dtypedtype, devicenpu) block_table torch.arange(0, total_blocks).reshape(batch_size, -1).to(torch.int32).npu() # block_table[b, i] 第 b 个 batch 的第 i 个逻辑 block 对应的物理 block ID # 2. 计算 slot_mapping每步重算公式见「深入原理」 slot_mapping kv_len kv_len_offset # Decode: [batch_size, 1] # 3. 写入缓存 torch_npu.scatter_update_(kv_cache, kv_len, key_states, -2) # 或融合写入: torch_npu.npu_kv_rmsnorm_rope_cache(..., slot_mapping, ...) # 4. FA 注意力计算FA v2 示例传 block_table attn_output, _ FA_OP( query, key_cache, value_cache, block_tableblock_table, # 逻辑 block → 物理 block 映射 block_sizeblock_size, # 块大小 actual_seq_kvlenactual_seq_lengths_kv, # 每个 batch 的实际 KV 长度 actual_seq_qlenactual_seq_lengths_q, # 每个 batch 的实际 Q 长度 input_layoutTND_NTD, # PA 模式常用 layout sparse_mode0, # Decode0, Prefill3 num_query_headsnum_heads, num_key_value_headsnum_kv_heads, softmax_scale1.0 / math.sqrt(head_dim), )参考文件cann-recipes-infer/models/deepseek_r1/models/modeling_deepseek.py、cann-recipes-infer/models/deepseek-v3.2-exp/models/modeling_deepseek.py模式三MLA 压缩缓存叠加在模式一或模式二之上MLA 将 KV 压缩为低维 latent只缓存压缩后的cache_nope非位置和cache_rope位置维度远小于原始 KV。核心差异# 缓存存储压缩维度kv_lora_rank512远小于完整 KVnum_heads * head_dim16384 cache_nope torch.zeros(..., 1, kv_lora_rank, ...) # 非位置部分 cache_rope torch.zeros(..., 1, qk_rope_head_dim, ...) # 位置编码部分 # FA 调用key 和 value 传同一个 cache_nopeabsorb 技术将 V 投影吸收到 O 投影中 attn_output, _ FA_OP( q_nope, k_nope_cache, k_nope_cache, # key value cache_nope query_ropeq_pe, key_ropek_rope_cache, # RoPE 单独传入 ... )重要约束query_rope和key_rope必须同时传或同时不传rope D 必须为 64MLA query D 仅支持 512 或 128MLA D512 时仅支持sparse_mode为 0、3、4参考文件cann-recipes-infer/models/deepseek_r1/models/modeling_deepseek.pyforward_absorb()第三层深入原理以下内容在实现分页注意力模式二时需要深入理解。3.1 物理内存布局与逻辑映射分页注意力的核心是将逻辑地址batch_idx, seq_pos映射到物理 block。KV 缓存物理存储: [total_num_blocks, block_size, num_kv_heads, head_dim] 假设 batch_size2, max_seq_len2048, block_size128: 每个 batch 需要: 2048 / 128 16 个 block 总共需要: 2 × 16 32 个 block Block 0~15: batch_0 的 token [0, 2047] Block 16~31: batch_1 的 token [0, 2047]映射公式给定: batch_idx, seq_pos, block_size, block_table 逻辑 block 编号 seq_pos // block_size block 内偏移 seq_pos % block_size 物理 block ID block_table[batch_idx, 逻辑 block 编号] 物理 slot 位置 物理 block ID × block_size block 内偏移 示例: batch_idx1, seq_pos300, block_size128 逻辑 block 300 // 128 2, 偏移 300 % 128 44 物理 block block_table[1, 2] 18 物理 slot 18 × 128 44 23483.2 block_table 构造静态预分配推理全程不变# 参考: cann-recipes-infer/models/deepseek-v3.2-exp/models/modeling_deepseek.py (line 809-820) num_blocks_per_seq max_seq_len // block_size total_blocks batch_size * num_blocks_per_seq block_table torch.arange(0, total_blocks).reshape(batch_size, -1).to(torch.int32).npu() # 结果: [[0,1,...,15], [16,17,...,31]] shape[batch_size, num_blocks_per_seq], dtypeint323.3 slot_mapping 构造slot_mapping是缓存写入位置索引用于npu_kv_rmsnorm_rope_cache等写入算子。核心公式slot(batch_idx, seq_pos) batch_idx × max_seq_len seq_pos在本仓库的顺序分配模式下slot_mapping 等于展平后的线性索引无需额外查表。Prefill—— 每个 batch 写入多个 token# 参考: cann-recipes-infer/models/deepseek-v3.2-exp/models/modeling_deepseek.py (line 2275-2282) 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, ...)) slot_mapping 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 个 token# 参考: cann-recipes-infer/models/deepseek-v3.2-exp/models/modeling_deepseek.py (line 2283-2284) # 预计算偏移一次性 kv_len_offset torch.arange(0, batch_size * max_seq_len, max_seq_len, ...).view(-1, 1) # 每步计算 slot_mapping kv_len.view(batch_size, -1) kv_len_offset # 示例: kv_len[522, 266], kv_len_offset[[0], [2048]] # slot_mapping [[522], [2314]]与 block_table 的分工slot_mapping→ 缓存写入算子npu_kv_rmsnorm_rope_cache 等block_table→ FA注意力读取算子npu_fused_infer_attention_score 等二者寻址逻辑一致职责不同3.4 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]|# 参考: cann-recipes-infer/models/deepseek-v3.2-exp/models/modeling_deepseek.py (line 2286-2301) # TND Prefill 示例: actual_seq_lengths_kv torch.cumsum(kv_len, dim0) # [512, 768] actual_seq_lengths_q actual_seq_lengths_kv.clone()3.5 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]具体是否 1 视框架约定 # 每次 Decode 后递增在 Runner 层model.forward() 之外每个 forward pass 只做一次 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 层变量各层只读不写。kv_len作为参数传入所有 Transformer 层scatter_update_使用同一个kv_len写入每层各自的缓存。不应在 attention/cache 内部递增 kv_len——否则 Prefill 阶段各层写入位置逐层偏移精度损坏。参考cann-recipes-infer/models/qwen3_moe/models/modeling_qwen3_moe.py。3.6 数据流总览初始化一次性 block_table arange().reshape(batch, -1) ← 推理全程不变 kv_cache zeros(total_blocks, block_size, heads, dim) kv_len_offset arange(0, batch*max_seq_len, max_seq_len) │ ▼ Prefill kv_len 从 attention_mask 计算 ← 如 [512, 256] slot_mapping [0..511, 2048..2303] ← 多 token 写入 actual_seq_lengths cumsum(kv_len) 或 kv_len → 每层: 写入 kv_cache FA 读取 kv_cache │ ▼ Decode循环kv_len 递增时机视具体 Runner 实现 slot_mapping kv_len kv_len_offset ← 写入位置 actual_seq_lengths_kv kv_len 1 ← FA 需读到刚写入的 token → 每层: 写入 kv_cache FA 读取 kv_cache kv_len 1 ← forward 之后递增完整代码示例以上模式一/二/三给出了精简的实现骨架。仓库中各模型的完整 FA 调用示例见references/fa-code-examples.md覆盖模式一连续缓存GPT-OSSTND, FA v2, sliding window、Qwen3-MoEBSH, FA v1模式二三PA MLADeepSeek-R1TND_NTD, FA v2, MLA absorb、Kimi-K2FA v1, Prefill/Decode 分离、LongCat-FlashBSND_NBSD, KVP缓存写入融合算子npu_kv_rmsnorm_rope_cache附录 AFA 融合算子版本对照特性FA v1 (npu_fused_infer_attention_score)FA v2 (npu_fused_infer_attention_score_v2)调用方式torch.ops.npu.npu_fused_infer_attention_score(...)self.fa_ops.npu_fused_infer_attention_score_v2(...)量化 KVantiquant_mode/antiquant_scaledequant_scale_key/valuequery_quant_modeSink token不支持learnable_sink参数典型使用Qwen3-MoE, LongCat-Flash, Kimi-K2DeepSeek-R1, GPT-OSS两个版本均支持PA传 block_table/ 非 PA不传、MLA rope 分离query_rope/key_rope。v1/v2 关键参数名映射混用不会报错静默使用默认值导致精度异常功能FA v1FA v2易错点缩放系数scalesoftmax_scale默认 1.0传错名精度崩溃Q head 数num_headsnum_query_heads传错名默认 1Q 长度actual_seq_lengthsactual_seq_qlenv1 Decode 时忽略但仍需传KV 长度actual_seq_lengths_kvactual_seq_kvlen名称完全不同KV head 数num_key_value_headsnum_key_value_heads相同附录 Binput_layout 选择layoutQ 格式KV 格式适用场景TND[T, N, D][T, N, D]非 PAflattened batchBSH[B, S, N*D][B, S, N*D]非 PA标准 batchBNSD[B, N, S, D][B, N, S, D]非 PA扩散模型TND_NTD[T, N, D][N, T, D]PA 模式NZ 格式缓存BSND_NBSD[B, S, N, D][N, B, S, D]PA 模式KVP 场景附录 Csparse_mode 与 atten_maskFA 算子的sparse_mode决定注意力遮蔽方式atten_mask配合提供掩码。两者组合错误是精度问题的第一大来源。参数配置规则sparse_mode含义atten_mask 要求适用场景0Dense可选通常传NoneDecode标准由actual_seq_lengths_kv控制有效长度1allMask必传完整矩阵(Q_S, KV_S)特殊场景2leftUpCausal不推荐建议改用 3—3Causal标准因果必传[2048, 2048]bool 下三角PrefillDecoder-only LLM 标准选择MTP Decodesq14Band滑动窗口必传[2048, 2048]boolPrefill / Decode滑窗模型如 gpt-oss需配合pre_tokensatten_mask 硬约束dtype只允许torch.bool推荐、torch.int8、torch.uint8。浮点类型直接报错shapesparse_mode3/4时固定[2048, 2048]与max_position_embeddings无关。本仓库统一用get_init_attn_mask(2048, device)构造见cann-recipes-infer/executor/utils/common_utils.py的get_init_attn_mask函数Decode 不需要 masksparse_mode0atten_maskNonePrefill / Decode 标准配置if q_len 1: # Prefill或 MTP Decode sparse_mode 3 atten_mask share_mask_tril # [2048, 2048] bool else: # Decode sparse_mode 0 atten_mask None参考cann-recipes-infer/models/deepseek_r1/models/modeling_deepseek.py的 Prefill/Decode 分支逻辑高频错误速查症状根因修复Prefill 输出乱码但不报错sparse_mode0maskNone无因果遮蔽改用sparse_mode3[2048, 2048]maskDecode 输出严重偏差Decode 误用sparse_mode3Decode 改sparse_mode0,maskNoneatten_mask dtype报错mask 用了 float16/bfloat16mask.to(torch.bool)atten_mask shape报错mask 不是[2048, 2048]固定用get_init_attn_mask(2048, device)scale默认 1.0 导致精度崩溃FA v1 用scalev2 用softmax_scale传错名静默生效确认参数名与 FA 版本匹配见附录 A 映射表附录 DNPU KVCache 算子速查算子功能接收入参torch_npu.scatter_update_KV 按位置写入缓存kv_lentorch_npu.npu_kv_rmsnorm_rope_cache融合 RMSNormRoPECache 写入slot_mappingtorch_npu.npu_fused_infer_attention_scoreFA v1 融合注意力block_tabletorch_npu.npu_fused_infer_attention_score_v2FA v2 融合注意力block_table附录 EPrefill vs Decode 参数差异参数PrefillDecodesparse_mode3因果0denseMTP 时 sq1 需用 3因果actual_seq_qlen输入序列长度1单 tokenMTP 时为 next_n1actual_seq_kvlen actual_seq_qlen累计的 KV 长度slot_mapping多 token 拼接单 token 偏移atten_mask因果 maskdtype 必须为 bool/int8/uint8或 None通常 NoneFA 参数配置Prefill 配置sparse_mode3 maskDecode 配置sparse_mode0 maskNone高阶特性以下特性仅供参考尚未完整文档化。KV 并行KVPKVCache 沿 head 维度切分到多卡。参考cann-recipes-infer/models/longcat-flash/KVCache 量化INT8/W8A8C8 量化。参考cann-recipes-infer/models/deepseek-v3.2-exp/models/offload_cache.pyCPU-GPU Offloadtorch_npu.empty_with_swapped_memory 异步双流。参考cann-recipes-infer/models/hstu/modules/gpu_kv_cache_manager.py扩散模型缓存DiT CacheTeaCache/FBCache/TaylorSeer。参考cann-recipes-infer/module/dit_cache/cache_method.py注意事项PA 必须配合 FA分页注意力必须通过 FA 融合算子的block_table参数索引无法使用标准 softmaxblock_table 静态不变初始化后推理全程不修改只有 kv_len 在递增slot_mapping 每步重算基于 kv_len kv_len_offsetPrefill 和 Decode 公式不同block_table dtype 必须为 int32Prefill/Decode 参数分离Prefill 和 Decode 需分别配置 FA 参数sparse_mode、atten_mask、actual_seq_lengths 等NZ 格式NPU 上 PA 模式推荐 NZ 格式存储 KVCachecache_modePA_NZinner_precise 行无效当 mask 存在整行全被遮蔽时如自定义 Prefill mask设置inner_precise2开启行无效修正【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考