一个 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 个 kernelQ/K/V 三个 Linear 共享同一个输入 xMergedMatMul 把它们合成一次 Batch GEMM。RotaryEmbedding 直接塞在 QKV 的输出上在片上缓存完成旋转。FlashAttention 把 MatMulSoftmaxMatMul 融成一个 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(GateUp) 1 个 kernel SiLU Elementwise Down Linear graph-autofusion 自动融合1-2 个 kernelGate 和 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.8121.6ops-transformer 融合1.238.4减少了 68% 的显存读写。在 HBM 带宽固定的前提下减少读写 减少延迟。推理服务的完整链路从请求到响应的全链路1. 请求进来tokenizationCPU 2. Token Embedding 查表1 次 HBM 读取 3. 32 层 Transformer每层 2-3 个融合 kernel 4. LayerNorm Linear 输出 logits2 个 kernel 5. SamplingCPU 1 次 HBM 读取 6. 响应返回步骤 3 占了 85% 以上的时间。在这 85% 里ops-transformer 的融合算子处理了其中 60% 的计算Attention FFN 的融合部分剩下 40% 是 RMSNorm、残差连接这些 ops-nn 的基础算子。KV Cache 对算子的影响推理的 prefill 阶段首次生成和 decode 阶段逐 token 生成走不同的代码路径Prefill完整序列进 FlashAttentionQ/K/V 全量计算DecodeQ 只有 1 个 tokenK/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_decodeTrue# Decode 阶段outtorch_npu.npu.flash_attention(q_decode,k_cache,v_cache,is_decodeTrue)搞清楚单层 Transformer 调了哪些算子调优才有抓手。NPU 利用率低先看 HBM 读写次数——如果融合没生效读写量会多两三倍。仓库在这里https://atomgit.com/cann/ops-transformer