Imported from maxwel1o/expert (skills/analyst/model-infer-kvcache/SKILL.md). Install upstream with npx skills add maxwel1o/expert --skill model-infer-kvcache. Copyright stays with the author.

KVCache 优化技能

本技能按渐进式披露组织：快速选型 → 实现模式 → 深入原理 → 参考附录。 根据需要的深度选择阅读层级。

第一层：快速选型

根据模型类型，选择 KVCache 模式和参考实现：

决策路径：

前置确认：从 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, dtype=dtype, device="npu")
past_value = torch.zeros_like(past_key)

# 2. 写入缓存（每步调用）
torch_npu.scatter_update_(past_key, kv_len, key_states, 1)     # axis=1 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_heads=num_heads, num_key_value_heads=num_kv_heads,
    input_layout="BSH",
    scale=1.0 / math.sqrt(head_dim),
    actual_seq_lengths_kv=actual_seq_lengths_kv,    # [batch_size] 每个 batch 的 KV 长度
    atten_mask=attention_mask,
    sparse_mode=0,                                   # 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.py（BSH）、cann-recipes-infer/models/gpt_oss/models/modeling_gpt_oss.py（TND）

模式二：分页注意力 + 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, dtype=dtype, device="npu")
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_table=block_table,                        # 逻辑 block → 物理 block 映射
    block_size=block_size,                          # 块大小
    actual_seq_kvlen=actual_seq_lengths_kv,         # 每个 batch 的实际 KV 长度
    actual_seq_qlen=actual_seq_lengths_q,           # 每个 batch 的实际 Q 长度
    input_layout="TND_NTD",                         # PA 模式常用 layout
    sparse_mode=0,                                  # Decode=0, Prefill=3
    num_query_heads=num_heads, num_key_value_heads=num_kv_heads,
    softmax_scale=1.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_rank=512）远小于完整 KV（num_heads * head_dim=16384）
cache_nope = torch.zeros(..., 1, kv_lora_rank, ...)     # 非位置部分
cache_rope = torch.zeros(..., 1, qk_rope_head_dim, ...)  # 位置编码部分

# FA 调用：key 和 value 传同一个 cache_nope（absorb 技术将 V 投影吸收到 O 投影中）
attn_output, _ = FA_OP(
    q_nope, k_nope_cache, k_nope_cache,   # key = value = cache_nope
    query_rope=q_pe, key_rope=k_rope_cache,  # RoPE 单独传入
    ...
)

重要约束：

• query_rope 和 key_rope 必须同时传或同时不传
• rope D 必须为 64
• MLA query D 仅支持 512 或 128
• MLA D=512 时仅支持 sparse_mode 为 0、3、4

参考文件：cann-recipes-infer/models/deepseek_r1/models/modeling_deepseek.py（forward_absorb()）

第三层：深入原理

以下内容在实现分页注意力（模式二）时需要深入理解。

3.1 物理内存布局与逻辑映射

分页注意力的核心是将 逻辑地址（batch_idx, seq_pos）映射到 物理 block。

KV 缓存物理存储: [total_num_blocks, block_size, num_kv_heads, head_dim]

假设 batch_size=2, max_seq_len=2048, block_size=128:
  每个 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_idx=1, seq_pos=300, block_size=128
  逻辑 block = 300 // 128 = 2, 偏移 = 300 % 128 = 44
  物理 block = block_table[1, 2] = 18
  物理 slot  = 18 × 128 + 44 = 2348

3.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], dtype=int32

3.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_len=2048
#   batch 0: [0, 1, ..., 511]
#   batch 1: [2048, 2049, ..., 2303]
#   拼接为 shape=[768] 的一维 tensor

Decode —— 每个 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：

# 参考: 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, dim=0)    # [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, axis=1)[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-OSS（TND, FA v2, sliding window）、Qwen3-MoE（BSH, FA v1）
• 模式二+三（PA + MLA）：DeepSeek-R1（TND_NTD, FA v2, MLA absorb）、Kimi-K2（FA v1, Prefill/Decode 分离）、LongCat-Flash（BSND_NBSD, KVP）
• 缓存写入融合算子：npu_kv_rmsnorm_rope_cache

附录 A：FA 融合算子版本对照

两个版本均支持：PA（传 block_table）/ 非 PA（不传）、MLA rope 分离（query_rope / key_rope）。

v1/v2 关键参数名映射（混用不会报错，静默使用默认值导致精度异常）：

附录 B：input_layout 选择

附录 C：sparse_mode 与 atten_mask

FA 算子的 sparse_mode 决定注意力遮蔽方式，atten_mask 配合提供掩码。两者组合错误是精度问题的第一大来源。

参数配置规则

atten_mask 硬约束

• dtype：只允许 torch.bool（推荐）、torch.int8、torch.uint8。浮点类型直接报错
• shape：sparse_mode=3/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 不需要 mask：sparse_mode=0 + atten_mask=None

Prefill / 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 分支逻辑

高频错误速查

附录 D：NPU KVCache 算子速查

附录 E：Prefill vs Decode 参数差异

高阶特性

以下特性仅供参考，尚未完整文档化。

• KV 并行（KVP）：KVCache 沿 head 维度切分到多卡。参考 cann-recipes-infer/models/longcat-flash/
• KVCache 量化：INT8/W8A8C8 量化。参考 cann-recipes-infer/models/deepseek-v3.2-exp/models/offload_cache.py
• CPU-GPU Offload：torch_npu.empty_with_swapped_memory + 异步双流。参考 cann-recipes-infer/models/hstu/modules/gpu_kv_cache_manager.py
• 扩散模型缓存（DiT Cache）：TeaCache/FBCache/TaylorSeer。参考 cann-recipes-infer/module/dit_cache/cache_method.py

注意事项

1. PA 必须配合 FA：分页注意力必须通过 FA 融合算子的 block_table 参数索引，无法使用标准 softmax
2. block_table 静态不变：初始化后推理全程不修改，只有 kv_len 在递增
3. slot_mapping 每步重算：基于 kv_len + kv_len_offset，Prefill 和 Decode 公式不同
4. block_table dtype 必须为 int32
5. Prefill/Decode 参数分离：Prefill 和 Decode 需分别配置 FA 参数（sparse_mode、atten_mask、actual_seq_lengths 等）
6. NZ 格式：NPU 上 PA 模式推荐 NZ 格式存储 KVCache（cache_mode="PA_NZ"）
7. inner_precise 行无效：当 mask 存在整行全被遮蔽时（如自定义 Prefill mask），设置 inner_precise=2 开启行无效修正