CANN-ops-transformer-从输入到输出昇腾NPU跑了多少个融合算子
一个 Llama2-7B 的推理请求进来，token 经过 32 层 Transformer 最终输出。每一层里，昇腾NPU到底调了 ops-transformer 仓库的哪些算子？搞清楚这个问题，推理性能调优才有方向。
单层 Transformer 的算子调用链
Llama2-7B 的每一层结构：
输入 x → RMSNorm → Q/K/V Linear → RotaryEmbedding → Self-Attention (GQA) → Attention Output Linear → 残差连接 → RMSNorm → Gate/Up Linear → SiLU 激活 → Down Linear → 残差连接 → 输出在 ops-transformer 融合算子介入之前，这一层有 11 个独立 kernel。融合后呢？
融合后的算子调用
Attention 部分：3 个 kernel 合成 1 个
原来： Q Linear → K Linear → V Linear （3 个 MatMul kernel） RotaryEmbedding （1 个逐元素 kernel） Q·K^T → Softmax → ·V （3 个 kernel） Attention Output Linear （1 个 MatMul kernel） 融合后： QKV + RotaryEmbedding 融合算子 （1 个 kernel） FlashAttention 融合算子 （1 个 kernel） Output Linear （1 个 kernel）Q/K/V 三个 Linear 共享同一个输入 x，MergedMatMul 把它们合成一次 Batch GEMM。RotaryEmbedding 直接塞在 QKV 的输出上，在片上缓存完成旋转。FlashAttention 把 MatMul+Softmax+MatMul 融成一个 kernel。
FFN 部分：3 个 kernel 合成 1-2 个
原来： Gate Linear → Up Linear （2 个 MatMul kernel） SiLU （1 个逐元素 kernel） Gate × Up （1 个逐元素 kernel） Down Linear （1 个 MatMul kernel） 融合后： MergedMatMul(Gate+Up) （1 个 kernel） SiLU + Elementwise + Down Linear （graph-autofusion 自动融合，1-2 个 kernel）Gate 和 Up Linear 共享输入，MergedMatMul 合并。SiLU 和 elementwise multiply 如果被 graph-autofusion 捕获，可以跟 Down Linear 融成 1 个 kernel；如果没有被捕获，SiLU 和 multiply 单独跑。
数据在 NPU 内的流转
HBM（显存） ↓ DMA 搬入 AIC Cube 单元 ← MatMul 在这里算 ↓ 片上缓存 AIV Vector 单元 ← Softmax/SiLU/RoPE 在这里算 ↓ 片上缓存 Cube 单元 ← 下一个 MatMul ↓ DMA 搬出 HBM融合算子的核心价值：让数据在 Cube→Vector→Cube 之间流转时不出 HBM。标准实现每一步都要写回 HBM 再读出来，融合后只在片上缓存里传递。
一次 forward pass 的 HBM 读写量对比：
| 配置 | 每层 HBM 读写 (GB) | 32 层总计 (GB) |
|---|---|---|
| 无融合 | 3.8 | 121.6 |
| ops-transformer 融合 | 1.2 | 38.4 |
减少了 68% 的显存读写。在 HBM 带宽固定的前提下，减少读写 = 减少延迟。
推理服务的完整链路
从请求到响应的全链路：
1. 请求进来，tokenization（CPU） 2. Token Embedding 查表（1 次 HBM 读取） 3. 32 层 Transformer（每层 2-3 个融合 kernel） 4. LayerNorm + Linear 输出 logits（2 个 kernel） 5. Sampling（CPU + 1 次 HBM 读取） 6. 响应返回步骤 3 占了 85% 以上的时间。在这 85% 里，ops-transformer 的融合算子处理了其中 60% 的计算（Attention + FFN 的融合部分），剩下 40% 是 RMSNorm、残差连接这些 ops-nn 的基础算子。
KV Cache 对算子的影响
推理的 prefill 阶段（首次生成）和 decode 阶段（逐 token 生成）走不同的代码路径：
- Prefill：完整序列进 FlashAttention，Q/K/V 全量计算
- Decode：Q 只有 1 个 token，K/V 从 KV Cache 读取
Decode 阶段的 FlashAttention 实现不同。Q 是[1, heads, 1, dim]，K/V 是[1, kv_heads, seq, dim]。这种形状下 FlashAttention 的分块策略退化为单行计算——分块没有意义，因为 Q 只有一行。
ops-transformer 对 decode 阶段有专门的 kernel 变体：flash_attention_decode。它省掉了分块逻辑，直接做一次 MatMul + Softmax + MatMul。性能比通用 FlashAttention 在 decode 形状下快 15%。
ATB 会自动判断 prefill/decode 并切换 kernel。手动调 API 时需要传is_decode=True：
# Decode 阶段out=torch_npu.npu.flash_attention(q_decode,k_cache,v_cache,is_decode=True)搞清楚单层 Transformer 调了哪些算子，调优才有抓手。NPU 利用率低？先看 HBM 读写次数——如果融合没生效，读写量会多两三倍。仓库在这里：
https://atomgit.com/cann/ops-transformer