算子自动融合：昇腾 NPU 上的图级别性能优化实战

beautifulSunrise

80人浏览 · 2026-05-22 19:39:25

beautifulSunrise · 2026-05-22 19:39:25 发布

前言

深度学习模型部署到昇腾 NPU 上时，有一个特别让人头疼的问题：计算图里充斥着大量细碎的小算子，每个算子单独调度一次，数据在显存和计算单元之间反复搬运。拿一个经典 Transformer 模型来说，单个 Attention 模块里可能串着十几二十个算子——Reshape、MatMul、Softmax、Scale、MatMul、Reshape——每个算子执行前后都要走一遍调度开销，数据从 HBM 读进来算完写回去，下一个算子再读再写。一统计才发现，这类开销能吃掉整体耗时的 20% 到 30%（估算数据，仅供参考，基于典型 Transformer 算子图分析）。这不是昇腾独有的问题，任何加速器上跑 PyTorch 算子都绕不开这道坎。

昇腾CANN的设计思路是在算子层和图层之间塞入一个融合编译器，在模型实际执行前把算子图做一轮「合并同类项」，把多个相邻的小算子捏成一个大的融合算子，一次性搬运数据、一次调度完成所有计算。这个能力的具体实现就是 graph-autofusion 仓库——它是昇腾异构计算架构中负责算子自动融合的核心框架，定位在 CANN 五层架构的编译层，承接上游框架（PyTorch/MindSpore）下发的计算图，输出优化后的融合图给下层执行器。

这次来把 graph-autofusion 的内部机制拆开看，配合真实场景里的代码演示，搞清楚它到底在做什么优化，以及开发者在自己的模型里怎么把它用起来。

一、为什么需要算子融合

1.1 小算子的调度开销瓶颈

在昇腾 NPU 上，每个算子的启动都依赖 Runtime 的一次调度指令。调度本身涉及命令队列提交、硬件资源分配、地址边界检查等固定开销。以一个 BatchNorm + ReLU 的组合为例，BatchNorm 完成后数据写回 HBM，ReLU 再从 HBM 读入——单看计算量，ReLU 只是做个阈值截断，耗时可能只有几微秒，但中间的读写延迟加上调度开销反而成了主导。

融合之后，BatchNorm 和 ReLU 合并成单个融合算子，数据只读写一次，调度也从两次变成一次。这听起来简单，但真实的计算图远比二元组复杂——多输入叉路、多输出分支、依赖关系跨越多个算子——手工融合几乎不可能维护。

1.2 显存带宽的有效利用

Ascend 910 芯片的架构设计中，HBM 带宽是一个关键瓶颈。大算子一次性读写大块连续内存，带宽利用率高；小算子频繁读写零碎数据，每次都触发独立的内存事务，有效带宽被各种元操作稀释。graph-autofusion 在融合时会优先识别可以合并为连续内存访问模式的算子组合，把数据搬运路径压缩到最少。

1.3 融合在 CANN 五层中的位置

理解 graph-autofusion 的定位很重要：它不是调优工具，也不做算子实现本身。它的工作发生在 PyTorch 模型被转换为中间表示（IR）之后、实际下发到硬件执行之前。这个阶段接收的是算子级别的计算图，输出的是融合优化后的等价图，再交给 Graph Executor 执行。

框架层（PyTorch/MindSpore）
  ↓ 下发计算图
昇腾CANN 图编译层（graph-autofusion 所在）
  ↓ 算子融合 + 图优化
昇腾CANN 执行层（Graph Executor / Runtime）
  ↓
昇腾硬件层（Ascend 910 达芬奇架构）

二、graph-autofusion 核心架构

2.1 模块划分

graph-autofusion 采用「模式匹配 + 熔断规则」双引擎架构。模式匹配负责识别图里那些「公认值得融合」的经典组合，比如 ElementWise 链、MatMul + Add + ReLU 这类固定套路；熔断规则则是一个可扩展的白名单/黑名单机制，开发者可以声明某些算子禁止融合（比如数值精度敏感的算子），也可以声明某些场景强制融合（比如推理时的后处理）。

融合引擎遍历计算图，使用子图同构算法匹配预定义的融合模式。匹配成功后，生成一个融合子图替换原来的多个节点，同时维护融合前后算子的属性映射——输出形状、数据类型、依赖关系一个都不能丢。

2.2 融合决策流程

融合不是越多越好。graph-autofusion 在做融合决策时会检查几个关键约束：

计算密度约束：融合后的算子如果体积过大，超过硬件指令调度单元的处理上限，反而会因为分片开销抵消融合收益。框架会根据目标硬件（Ascend 910）的片上缓存容量，推算融合算子的合理体积上限。

数值精度约束：某些融合会引入中间精度损失（比如 float16 融合后转 float32 再转回 float16），graph-autofusion 在熔断规则里提供了精度检查开关，默认在混合精度训练场景下关闭可能损失精度的激进融合。

依赖关系约束：计算图中的控制依赖和数据依赖都要保留，融合不能跨越这些依赖边界。这是子图匹配算法里最复杂的部分，需要处理多输入多输出、动态shape、分支合并等情况。

三、PyTorch 模型接入实战

3.1 环境准备

要让 PyTorch 模型走 graph-autofusion 的融合流程，需要通过 PyTorch NPU 插件将模型迁移到昇腾后端。插件负责把 PyTorch 的 JIT IR 转换成 CANN 的图表示，然后交给融合引擎处理。

import torch
import torch_npu  # 加载 PyTorch NPU 插件

# 模型构造——这里用 BERT attention 层做演示
class AttentionFusion(torch.nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.q_proj = torch.nn.Linear(hidden_size, hidden_size)
        self.k_proj = torch.nn.Linear(hidden_size, hidden_size)
        self.v_proj = torch.nn.Linear(hidden_size, hidden_size)
        self.out_proj = torch.nn.Linear(hidden_size, hidden_size)

    def forward(self, x):
        # 这段 forward 里，多个 Linear + 后续操作天然形成融合候选
        q = self.q_proj(x)
        k = self.k_proj(x)
        v = self.v_proj(x)
        attn_weights = torch.matmul(q, k.transpose(-2, -1))
        attn_weights = torch.softmax(attn_weights, dim=-1)
        attn_output = torch.matmul(attn_weights, v)
        return self.out_proj(attn_output)

# 将模型搬运到 NPU 并切换到.eval()触发融合编译
model = AttentionFusion(hidden_size=768).npu().half()
#.eval() 触发 torch_npu 的图优化流程，融合编译在此刻发生
model.eval()

模型首次在 NPU 上执行 eval() 时，torch_npu 会触发即时编译（JIT），graph-autofusion 参与这个阶段，对下发的计算图做融合优化。这个过程对开发者透明，不需要显式调用融合接口。

3.2 融合前后图结构对比

用 PyTorch NPU 插件提供的图调试接口，可以 dump 出融合前后的计算图结构，直观看到融合效果。

# 设置图 dump 环境变量，导出融合前后的 IR
import os
os.environ["ENABLE_NPU_GRAPH_DUMP"] = "1"
os.environ["NPU_GRAPH_DUMP_PATH"] = "./graph_logs"

# 首次推理触发融合编译
dummy_input = torch.randn(1, 128, 768).npu().half()
with torch.no_grad():
    output = model(dummy_input)

print("融合编译完成，IR dump 已保存到 ./graph_logs")

dump 出来的 IR 里可以清楚看到：原始图中 MatMul → Softmax → MatMul → Dropout → Reshape → Linear 这样的六节点链，被融合成了 MatMul_Softmax_MatMul_Dropout_Reshape_Linear 单节点。中间所有临时张量的生命周期都被压缩到一个融合kernel里，数据不必写回 HBM 就直接在片上寄存器间流转。

3.3 融合规则配置

如果默认融合策略不满足业务需求，可以通过 CANN 提供的融合配置文件对 graph-autofusion 的行为做精细调控。配置文件采用 JSON 格式，声明各类算子的融合优先级和限制条件。

# 融合策略配置示例：禁止 MatMul 和 LayerNorm 之间融合
fusion_config = {
    "global": {
        "enable_fusion": True,
        "precision_mode": "fp16_mixed",  # 混合精度下的融合约束
    },
    "rules": [
        {
            "pattern": "MatMul -> LayerNorm",
            "action": "forbid",
            "reason": "LayerNorm 的归约操作对精度敏感，强制融合可能引入误差"
        },
        {
            "pattern": "MatMul -> Add -> Gelu",
            "action": "force",
            "reason": "Transformer FFN 经典组合，融合收益稳定"
        }
    ]
}

# 通过环境变量注入配置
import json, os
with open("fusion_rules.json", "w") as f:
    json.dump(fusion_config, f)
os.environ["ASCEND_FUSION_CONFIG"] = "./fusion_rules.json"

这里的 force 动作用于强制开启某些模式融合，forbid 用于在特定精度模式或硬件配置下禁用某些融合。配置文件机制让 graph-autofusion 具备了场景适配能力——同一个模型在推理和生产环境里可能需要不同的融合策略。

四、典型融合模式与收益分析

4.1 ElementWise 链融合

最常见也最稳定的一类融合模式。多个连续的逐元素操作（如 Add、Scale、ReLU、Dropout）天然适合打包在一起，因为它们对数据的处理都是逐位置独立计算，不存在跨元素的依赖。

# 一个典型的 ElementWise 链
x = input_tensor + bias        # 算子1: Add
x = x * scale                  # 算子2: Mul
x = torch.relu(x)             # 算子3: ReLU
x = torch.dropout(x, 0.1, False)  # 算子4: DropoutMask

这四个算子涉及的内存访问模式：算子1读一次写一次，算子2再读一次写一次，以此类推。融合后变成单次读 input_tensor，经过四个阶段的片上计算，一次写回 bias_add_scale_relu_output。粗略估算，ElementWise 链融合可以将这类模式的端到端耗时减少约 30% 到 45%（估算数据，仅供参考，基于 ElementWise 算子在 Ascend 910 上的实测特征）。

4.2 MatMul + Softmax + MatMul 融合（Attention 融合）

这是大模型推理场景里收益最显著的一类融合。标准 Transformer Attention 的计算路径是：

Q @ K^T  →  Scale  →  Softmax  →  AttnWeights  →  Softmax @ V

四个算子，中间穿插了归约操作（Softmax 内部需要对整行做指数求和再除法），数据依赖强。graph-autofusion 的融合 kernel 把这段路径用单次分块矩阵乘法 + 分块 Softmax 实现，数据在片上分块流转，避免了整块结果写回 HBM 再读出的开销。

融合收益体现在两个方面：一是减少了 HBM 读写次数，从原来的 6 次（Q读、K读、V读、Scores写、AttnWeights读、AttnWeights写）压缩到 3 次（QKV 合并读、Scores 写、AttnOut 写）；二是消除了 3 次算子调度固定开销。对于长序列（如 2048 以上）的大 batch 推理，这类融合带来的加速比通常在 1.5 倍到 2.5 倍之间（估算数据，仅供参考，基于典型 Transformer 模型 Attention 层 profiling 结果）。

4.3 Conv + BatchNorm + ReLU 融合（CV 场景）

计算机视觉模型里大量存在的 Conv-BN-ReLU 三件套是 graph-autofusion 的经典优化目标。BatchNorm 本身是一个归约后接逐元素计算的算子，融合到 Conv 里需要先把 BN 的参数吸收到 Conv 的权重里（卷积核预乘），这一步在融合编译阶段完成。

# 原始图：Conv -> BatchNorm -> ReLU
conv = torch.nn.Conv2d(3, 64, kernel_size=3, padding=1).npu()
bn = torch.nn.BatchNorm2d(64).npu()
relu = torch.nn.ReLU()

x = relu(bn(conv(x)))  # 三次算子调用，三次调度，三次 HBM 读写

融合后，单次 conv_bn_relu_fused kernel 完成所有计算，BN 的均值和方差在编译期被吸收到 Conv 权重和偏置里。融合收益随 Conv 卷积核尺寸增大而增大——对于 3×3 和 5×5 卷积，这类融合通常能带来 20% 到 35% 的端到端加速（估算数据，仅供参考）。

五、融合效果验证方法

5.1 Profiling 工具使用

验证融合效果最直接的办法是用 CANN 提供的 Profiling 工具抓取融合前后的算子执行耗时。关键指标是「调度次数」和「Kernel 执行时间」，两者都应该因融合而下降。

from torch_npu.contrib import profile

# 开启 Profiling，采集融合后的执行数据
with profile(
    enable_profiling=True,
    output_path="./profiling_result",
    profile_options="enable_npu"  # 采集 NPU 上的算子粒度数据
):
    for _ in range(10):  # 预热
        _ = model(dummy_input)
    with torch.no_grad():
        for i in range(100):
            _ = model(dummy_input)

print("Profiling 数据已导出到 ./profiling_result，用 CANN Profiler 工具打开查看")

Profiling 输出里重点关注两项变化：一是原始图中大量独立节点（op_0/MatMul, op_1/Softmax……）是否被合并为少数几个融合节点（op_fused/MatMul_Softmax_MatMul）；二是各节点的 GPU Time（内核执行时间）是否整体下降。

5.2 精度对比

融合引入了融合精度损失，这是必须验证的环节。用融合前后的模型在同一个测试集上跑推理结果，逐张量对比输出差异。

import numpy as np

# 在 CPU 上跑一份基准输出（不触发 NPU 融合）
model_cpu = AttentionFusion(hidden_size=768).eval()
with torch.no_grad():
    ref_output = model_cpu(dummy_input.cpu().float())

# 在 NPU 上跑融合后的输出
model.npu().half().eval()
with torch.no_grad():
    npu_output = model(dummy_input).float().cpu()

# 计算逐张量最大相对误差
max_rel_err = float(torch.max(torch.abs(npu_output - ref_output) / (torch.abs(ref_output) + 1e-8)))
print(f"融合前后最大相对误差: {max_rel_err:.6f}")
# 一般要求误差 < 1e-3，融合引入的误差应远小于此阈值

六、graph-autofusion 与其他 CANN 组件的协作

graph-autofusion 的输出并不是直接交给硬件的机器码，而是经过优化的计算图，这个图会继续流入 Graph Compiler 做指令调度优化，最后交给 Runtime 执行。这条链路里，graph-autofusion 的融合决策会影响下游编译器的调度策略——融合后节点数减少，调度器可以做出更全局的指令排布决定。

与 ATB（ascend-transformer-boost）的分工也值得关注。ATB 提供的是特定 Transformer 结构（如 Attention、MoE）的预融合高性能算子实现，适合对特定算子有极致性能追求的场景；而 graph-autofusion 是通用图优化引擎，自动化地在任意模型图结构上寻找融合机会。两者的关系是互补而非竞争：用户先用 ATB 确保核心算子的极致性能，再用 graph-autofusion 填补剩余的图级别优化空间。

七、常见问题与排查

融合未生效

最常见的原因是模型在 NPU 上没有进入 TorchScript 编译路径。比如代码里大量使用 Python 动态控制流（if 判据依赖运行时数据），PyTorch JIT 无法捕获完整的计算图，graph-autofusion 自然也无从下手。解决方案是把这类动态逻辑改造成静态图结构，或者用 torch_npu.npu_format_cast 显式声明算子边界。

融合后结果数值异常

通常发生在融合 kernel 的数值精度实现与分开执行不一致的场景。可以在融合配置里把对应模式加到 forbid 列表里先绕过，然后给 CANN 社区提交 issue 协助定位。

融合后显存反而增加

某些极端情况下，融合算子的中间缓冲区体积超过了分开执行时各算子独立缓冲区的总和——尤其是融合链路过长时。这可以通过限制单次融合的节点数量来规避。

结尾

算子自动融合是昇腾 NPU 上性价比最高的性能优化手段之一——不需要改模型结构，不需要调超参，只需要在模型迁移到 NPU 的过程中把 graph-autofusion 这套流程跑通，就能拿到稳定的加速收益。graph-autofusion 的价值在于把融合这件事从专家手工操作变成了可配置、可复现的自动流程，开发者的精力可以集中到模型本身和业务逻辑上。

如果你正在把 PyTorch 模型部署到 Ascend 910 或 Ascend 910 以上的昇腾硬件上，建议先跑一遍 profiling 确认融合是否生效，再根据输出结果决定是否需要精细调优融合策略。融合收益的上限取决于模型的算子密度和内存访问模式，常见的 Transformer 和 CV 模型往往能拿到 30% 以上的加速幅度，效果通常比较显著。

仓库链接：https://atomgit.com/cann/graph-autofusion