external-cannbot-model-model-infer-kvcache

KVCache Optimization Skill

This skill is organized by progressive disclosure: quick selection → implementation patterns → in-depth principles → reference appendix. Choose the reading level according to the depth needed.

---

Layer 1: Quick Selection

According to the model type, choose the KVCache mode and reference implementation:

Model Type	Selected Mode	Reference Implementation	Key Points
Standard LLM (MHA/GQA)	Mode 1: Contiguous cache	gpt-oss, qwen3-moe	Simplest, no block_table needed
High-performance LLM	Mode 2: Paged attention + FA	deepseek-r1, kimi-k2	Need to construct block_table + slot_mapping
MLA architecture (DeepSeek series)	Mode 3: MLA compression	deepseek-r1, deepseek-v3.2	Stacked on top of Mode 1 or Mode 2
Diffusion/video models	FA direct computation	hunyuan-video, wan2.2-i2v	No KVCache paging, BNSD layout

Decision Path：

前置确认：从 progress.md 读取 Phase 0 已验证的架构类型（基于实际 config 值，非代码推断）
         │
需要分页注意力？ ──→ 否 → 模式一（连续缓存，scatter_update_ + FA）
                  └→ 是 → 模式二（分页注意力，block_table + slot_mapping + FA）
                           │
使用 MLA？ ──────→ 是 → 叠加模式三（cache_nope + cache_rope 压缩存储）

---

Layer 2: Implementation Patterns

Mode 1: Contiguous cache (non-paged attention)

KV is stored as a contiguous tensor,scatter_update_ write, FA directly reads the entire cache.

Three-step implementation：

# 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)
)

Modification boundary: Only modify the attention computation (→FA) and cache storage (→scatter_update_). Upstream (QKV projection, RoPE, etc.) remains unchanged; when layouts do not match, adapt with transpose/reshape at the seam.

Reference Files：cann-recipes-infer/models/qwen3_moe/models/modeling_qwen3_moe.py（BSH）、cann-recipes-infer/models/gpt_oss/models/modeling_gpt_oss.py（TND）

---

Mode 2: Paged attention + FA fused operator

KV is stored in fixed-size Blocks, and FA uses block_table to index the block cache.PA must be used together with FA.

Four-step implementation：

# 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),
)

Reference Files：cann-recipes-infer/models/deepseek_r1/models/modeling_deepseek.py、cann-recipes-infer/models/deepseek-v3.2-exp/models/modeling_deepseek.py

---

Mode 3: MLA compressed cache (stacked on top of Mode 1 or Mode 2)

MLA compresses KV into a low-dimensional latent, caching only the compressed cache_nope(non-positional) and cache_rope(positional), with dimensions far smaller than the original KV.

Core Differences：

# 缓存存储：压缩维度（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 单独传入
    ...
)

Important Constraints：

• query_rope and key_rope must be passed together or not passed at all
• rope D must be 64
• MLA query D only supports 512 or 128
• When MLA D=512, only supports sparse_mode being 0, 3, 4

Reference Files：cann-recipes-infer/models/deepseek_r1/models/modeling_deepseek.py（forward_absorb()）

---

Layer 3: In-depth Principles

The following content requires in-depth understanding when implementing paged attention (Mode 2).

3.1 Physical Memory Layout and Logical Mapping

The core of paged attention is mapping the logical address(batch_idx, seq_pos) to the physical 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]

Mapping Formula：

给定: 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 Construction

Static pre-allocation, unchangedthroughout inference：

# 参考: 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 Construction

slot_mapping is the cachewriteposition index, used for npu_kv_rmsnorm_rope_cache and other write operators.

Core Formula：slot(batch_idx, seq_pos) = batch_idx × max_seq_len + seq_pos

In this repository's sequential allocation mode, slot_mapping equals the flattened linear index, with no additional table lookup required.

Prefill — Multiple tokens written per batch:

# 参考: 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 — 1 token written per batch:

# 参考: 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]]

Division of Labor with block_table：

• slot_mapping → Cachewriteoperators (npu_kv_rmsnorm_rope_cache, etc.)
• block_table → FA attention readoperators (npu_fused_infer_attention_score, etc.)
• The two share the same addressing logic but have different responsibilities

3.4 actual_seq_lengths Construction

The FA operator requires the actual KV/Q length of each batch; the construction method depends on input_layout：

	TND layout (multiple batch tokens concatenated into one dimension)	BSH layout (each batch independent)
Prefill KV	cumsum(kv_len) → [512, 768]	kv_len → [512, 256]
Prefill Q	Same as KV	Same as 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, dim=0)    # [512, 768]
actual_seq_lengths_q = actual_seq_lengths_kv.clone()

3.5 kv_len Lifecycle

kv_len is the core variable driving the computation of all input parameters:

# 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 is a Runner-layer variable; each layer only reads it and does not write to it. kv_len is passed as a parameter to all Transformer layers,scatter_update_ uses the same kv_len to write to each layer's respective cache. kv_len should not be incremented inside attention/cache—otherwise, during the Prefill phase, the write position of each layer shifts layer by layer, corrupting precision. Reference cann-recipes-infer/models/qwen3_moe/models/modeling_qwen3_moe.py。

3.6 Data Flow Overview

初始化（一次性）
  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 之后递增

---

Complete Code Examples

The above modes one/two/three provide streamlined implementation skeletons. For complete FA call examples for each model in the repository, see references/fa-code-examples.md, covering:

• Mode 1 (contiguous cache): GPT-OSS (TND, FA v2, sliding window), Qwen3-MoE (BSH, FA v1)
• Mode 2 + 3 (PA + MLA): DeepSeek-R1 (TND_NTD, FA v2, MLA absorb), Kimi-K2 (FA v1, Prefill/Decode separation), LongCat-Flash (BSND_NBSD, KVP)
• Cache write fusion operators:npu_kv_rmsnorm_rope_cache

---

Appendix A: FA Fusion Operator Version Comparison

Feature	FA v1 (npu_fused_infer_attention_score)	FA v2 (npu_fused_infer_attention_score_v2)
Call Method	torch.ops.npu.npu_fused_infer_attention_score(...)	self.fa_ops.npu_fused_infer_attention_score_v2(...)
Quantized KV	antiquant_mode / antiquant_scale	dequant_scale_key/value + query_quant_mode
Sink token	Not supported	learnable_sink parameters
Typical Usage	Qwen3-MoE, LongCat-Flash, Kimi-K2	DeepSeek-R1, GPT-OSS

Both versions support: PA (pass block_table) / non-PA (do not pass), MLA rope separation (query_rope / key_rope）。

v1/v2 Key Parameter Name Mapping (mixing them will not raise an error, silently using default values causes precision anomalies)：

Function	FA v1	FA v2	Common Pitfalls
Scaling Factor	scale	softmax_scale	Defaults to 1.0; passing the wrong name causes precision collapse
Number of Q heads	num_heads	num_query_heads	Passing the wrong name defaults to 1
Q length	actual_seq_lengths	actual_seq_qlen	Ignored in v1 Decode but still must be passed
KV length	actual_seq_lengths_kv	actual_seq_kvlen	Completely different names
Number of KV heads	num_key_value_heads	num_key_value_heads	Same

Appendix B: input_layout Selection

layout	Q Format	KV Format	Applicable Scenarios
"TND"	[T, N, D]	[T, N, D]	Non-PA, flattened batch
"BSH"	[B, S, N*D]	[B, S, N*D]	Non-PA, standard batch
"BNSD"	[B, N, S, D]	[B, N, S, D]	Non-PA, diffusion models
"TND_NTD"	[T, N, D]	[N, T, D]	PA mode, NZ format cache
"BSND_NBSD"	[B, S, N, D]	[N, B, S, D]	PA mode, KVP scenario

Appendix C: sparse_mode and atten_mask

The FA operator's sparse_mode determines the attention masking method,atten_mask Provided together with the mask. The combination error of the two is the number one source of precision issues.

Parameter configuration rules

sparse_mode	Meaning	atten_mask requirement	Applicable scenario
0	Dense	Optional, usually pass None	Decode(standard), with actual_seq_lengths_kv controlling the effective length
1	allMask	RequiredFull matrix (Q_S, KV_S)	Special scenarios
2	leftUpCausal	Not recommended, recommend switching to 3	—
3	Causal (standard causal)	Required [2048, 2048] bool lower triangular	Prefill(standard choice for Decoder-only LLM); MTP Decode (sq>1)
4	Band (sliding window)	Required [2048, 2048] bool	Prefill / Decode (sliding window models, such as gpt-oss, need to be combined with pre_tokens）

atten_mask hard constraints

• dtype: only torch.bool(recommended),torch.int8、torch.uint8are allowed. Floating-point types directly error out
• shape：sparse_mode=3/4 fixed at [2048, 2048]when max_position_embeddings , unrelated to get_init_attn_mask(2048, device) . This repository uniformly uses cann-recipes-infer/executor/utils/common_utils.py construction (see the get_init_attn_mask function of
• Decode does not need a mask：sparse_mode=0 + atten_mask=None

Prefill / Decode standard configuration

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

Reference:cann-recipes-infer/models/deepseek_r1/models/modeling_deepseek.py Prefill/Decode branch logic

High-frequency error quick reference

Symptom	Root cause	Fix
Prefill output garbled but no error	sparse_mode=0 + mask=None, no causal masking	Switch to sparse_mode=3 + [2048, 2048] mask
Decode output severely deviates	Decode mistakenly uses sparse_mode=3	Decode change sparse_mode=0, mask=None
atten_mask dtype errors out	mask uses float16/bfloat16	mask.to(torch.bool)
atten_mask shape errors out	mask is not [2048, 2048]	fixed use get_init_attn_mask(2048, device)
scale defaults to 1.0 causing precision collapse	FA v1 uses scale, v2 uses softmax_scale, wrong name silently takes effect	Confirm the parameter name matches the FA version (see Appendix A mapping table)

Appendix D: NPU KVCache operator quick reference

Operator	Function	Accepted input parameters
torch_npu.scatter_update_	KV is written to the cache by position	kv_len
torch_npu.npu_kv_rmsnorm_rope_cache	Fused RMSNorm+RoPE+Cache write	slot_mapping
torch_npu.npu_fused_infer_attention_score	FA v1 fused attention	block_table
torch_npu.npu_fused_infer_attention_score_v2	FA v2 fused attention	block_table

Appendix E: Prefill vs Decode parameter differences

Parameter	Prefill	Decode
sparse_mode	3 (causal)	0 (dense); for MTP when sq>1, 3 (causal) is required
actual_seq_qlen	Input sequence length	1 (single token); next_n+1 for MTP
actual_seq_kvlen	= actual_seq_qlen	Accumulated KV length
slot_mapping	Multi-token concatenation	Single token offset
atten_mask	Causal mask (dtype must be bool/int8/uint8) or None	Usually None
FA parameter configuration	Prefill configuration (sparse_mode=3 + mask)	Decode configuration (sparse_mode=0 + mask=None)

---

Advanced features

The following features are for reference only and are not yet fully documented.

• KV Parallelism (KVP): KVCache is partitioned along the head dimension across multiple cards. See cann-recipes-infer/models/longcat-flash/
• KVCache Quantization: INT8/W8A8C8 quantization. See cann-recipes-infer/models/deepseek-v3.2-exp/models/offload_cache.py
• CPU-GPU Offload：torch_npu.empty_with_swapped_memory + asynchronous dual-stream. See cann-recipes-infer/models/hstu/modules/gpu_kv_cache_manager.py
• Diffusion Model Cache (DiT Cache): TeaCache/FBCache/TaylorSeer. See cann-recipes-infer/module/dit_cache/cache_method.py

---

Notes

1. PA must be used with FA: Paged attention must be indexed through the FA fused operator's block_table parameter, and standard softmax cannot be used
2. block_table is statically immutable: After initialization, it is not modified throughout inference; only kv_len increases
3. slot_mapping is recomputed each step: Based on kv_len + kv_len_offset; the formulas for Prefill and Decode are different
4. block_table dtype must be int32
5. Prefill/Decode parameters are separated: Prefill and Decode require separately configuring FA parameters (sparse_mode, atten_mask, actual_seq_lengths, etc.)
6. NZ format: On NPU, PA mode recommends storing KVCache in NZ format (cache_mode="PA_NZ"）
7. inner_precise row invalid: When the mask has an entire row fully masked (such as a custom Prefill mask), set inner_precise=2 to enable row invalid correction