CANN-ops-transformer遇上graph-autofusion-昇腾NPU算子自动融合的秘密

ops-transformer 的融合算子把 Attention 内部的操作合成了一个 kernel。但 Attention 前后还有 RMSNorm、残差连接这些操作——它们各自是独立的 kernel，每次执行都有 HBM 读写开销。graph-autofusion 在 GE（图引擎）编译阶段自动把这些相邻算子也融合了，用户不需要手动指定。

两种融合的区别

算子内融合（ops-transformer 做的）： 算子开发者用 Ascend C 手写的融合 kernel。FlashAttention 把 MatMul+Softmax+MatMul 写在一个 .cpp 文件里，编译时就是一个不可分割的 kernel。这种融合是确定性的，不依赖运行时上下文。

算子间融合（graph-autofusion 做的）： GE 编译器在构建计算图时，发现两个相邻算子的数据流满足融合条件，自动生成一个融合 kernel。比如 FlashAttention 的输出直接进 RMSNorm——GE 把它们合成一个 kernel，中间结果不写回 HBM。

算子内融合：[MatMul → Softmax → MatMul]  ← 1个kernel，手动实现
算子间融合：[FlashAttention → RMSNorm → 残差] ← 1个kernel，自动生成

graph-autofusion 的融合规则

不是所有相邻算子都能融合。GE 的融合引擎有一套规则：

能融合的条件：

• 数据流是线性的（A 的输出是 B 的唯一输入）
• 两个算子在同一个 stream 上
• 融合后的 register pressure 不超过 AI Core 的上限

不能融合的情况：

• 中间结果被多个下游算子消费（fork 点）
• 两个算子的数据格式不兼容（比如一个 NHWC 一个 NCHW）
• 融合后 register 超限

实际上，Transformer 模型里有大量可以自动融合的机会。残差连接（elementwise add）跟它前面的 RMSNorm 或 Linear，几乎总是可以融合的。ops-transformer 的算子输出格式是固定的，GE 知道它的输出可以直接喂给哪些算子。

自动融合 vs 手动融合

维度	手动融合（ops-transformer）	自动融合（graph-autofusion）
谁决定融合	算子开发者	GE 编译器
融合粒度	算子内部	算子之间
灵活性	低，写死在代码里	高，编译时动态决定
性能上限	高，针对性优化	中，通用融合策略
开发成本	高，需要写 Ascend C	零，自动完成

最好的效果是两者配合：ops-transformer 做算子内融合（大头），graph-autofusion 做算子间融合（锦上添花）。

实际收益

Llama2-7B 推理，单层 Transformer 的 kernel 数量：

配置	kernel 数/层	HBM 读写/层 (GB)
无融合	11	3.8
仅 ops-transformer 融合	5	1.2
ops-transformer + graph-autofusion	3	0.8

自动融合额外减少了 2 个 kernel 和 0.4 GB 的 HBM 读写。收益没有算子内融合大，但零成本——你不用改任何代码。

怎么确认自动融合生效

GE 编译时会输出融合日志：

# 设置 GE 日志级别
export GE_LOG_LEVEL=0  # DEBUG 模式

# 运行推理，检查日志
grep "AutoFusion" ge_log.txt
# 如果看到类似：
# AutoFusion: FlashAttention + Add + RMSNorm -> FusedKernel
# 说明自动融合生效了

如果日志里没有 AutoFusion 记录，常见原因：

1. 模型用的是 PyTorch eager mode，GE 没介入——需要走 torch.compile 或 ATB
2. 中间结果被 tensor detach 了，GE 认为数据流断了
3. CANN 版本太旧，graph-autofusion 在 CANN 8.0 才正式默认启用

什么时候自动融合不靠谱

graph-autofusion 的融合策略是通用的，不针对特定模型优化。在以下场景它可能做错选择：

• MoE 模型的 expert 分支。 每个 expert 的 FFN 是一个独立子图，GE 可能不会跨 expert 做融合。这种情况需要用 ops-transformer 的 MergedMatMul 手动合并。
• 动态 shape。 如果序列长度在运行时才确定，GE 可能放弃融合（需要重新编译）。CANN 8.5 对动态 shape 的支持有改善，但还不完美。
• 自定义算子。 如果你在模型里插了一个自己写的 Ascend C 算子，GE 不一定能判断它的融合安全性，保守策略是不融合。

---

大多数场景下，graph-autofusion 是白送的性能提升。你只要确保推理走的是 GE 编译路径（ATB 或 torch.compile），自动融合就会生效。如果追求极致性能，在自动融合的基础上再用 ops-transformer 做算子内融合。两个仓库：

https://atomgit.com/cann/ops-transformer