昇腾CANN ops-nn 仓的融合算子：一层调用干掉三次搬运

MatMul + BiasAdd + Activation 这个三连操作，逐个调用等于在 NPU 和 HBM 之间搬三次数据——第一次写 MatMul 的结果，第二次写 BiasAdd 的结果，第三次写 Activation 的结果。三次 kernel Launch 和三次 HBM 写回全是开销。ops-nn 仓是昇腾CANN 的神经网络类基础算子库，位于第二层 AOL 算子库，提供 matmul、activation 类算子，支持融合。这个融合算子机制把三步压成一个 kernel，这篇文章看它怎么做到的。

为什么要融合

昇腾CANN 的算子库里，每个算子都是独立编译、独立调度的。调用一次 MatMul 要走一遍完整的 Launch 流程：Host 端发起调用请求 → Runtime 调度器分配 Stream → Kernel 启动 → 结果写回 HBM。每个环节都有固定开销。

单独来看一次 Launch 的开销不大，几百纳秒。但一个 LLaMA-7B 的推理要把 Attention 层跑 32 次（32 层），每层的计算量摆在那里，三连操作乘以 3，乘以 32，再乘以 Batch Size，这个开销就不少了。

更严重的是中间的 HBM 读写。三次调用产生三个中间结果，每个都要从 NPU 写到 HBM 再读回来。NPU 到 HBM 的带宽是有限资源，大量的小数据块会让 HBM 带宽成为瓶颈而不是计算单元。如果把三次操作合成一次，原则上只能产生一个中间结果，HBM 访问次数降到三分之一。

融合算子本质上是一种计算图的优化，把多个小算子合并成一个大算子。这样做有三个直接好处：kernel Launch 开销从 3 次降到 1 次、HBM 读写从 3 次降到 1 次（或者 0 次如果输出直接被下一个算子消费）、三个算子之间的临时内存可以复用不用额外分配。

ops-nn 支持哪些融合模式

ops-nn 仓目前支持以下几种融合模式，按场景分：

基础融合是 MatMul + BiasAdd + Activation。这里的 Activation 指的是 ReLU、GELU、SiLU、SwiGLU 这些激活函数。融合的方式是把矩阵乘的结果直接做 bias 加法再做激活，全程在一个 kernel 里完成。

更深一点的是 MatMul + BiasAdd + Reshape + Activation，多一个维度变换。LLM 里很多地方是三维张量（batch, seq, hidden），两维的 MatMul 融合完了要 reshape 回三维再激活，这种也能合。

Convolution + BiasAdd + Activation 也是常见的融合模式。卷积神经网络里 Conv 后面接 BN 接 ReLU 是标配，早期的框架把它们分开调用，现在也能融成两步。还有一些更复杂的融合模式，比如带残差连接的 MatMul + Add + Activation，以及多分支的 Concat + MatMul + Activation。这些要看具体的模型和 CANN 版本，有些是实验性支持。

Ascend C 融合算子的写法

融合算子在 Ascend C 里的写法跟普通算子不太一样。普通算子只有一个核心计算逻辑，融合算子的关键是先把三步拆开写，然后用 __aicore__ 的融合注解把它们拼到一个 kernel 里。

下面是一个 MatMul + BiasAdd + GELU 融合算子的简化代码：

// 融合算子核心函数
// 把矩阵乘、偏置加、激活函数放在一个kernel里完成
extern "C" __global__ __aicore__ void matmul_bias_gelu_kernel(
    GM_ADDR a_gm, GM_ADDR b_gm, GM_ADDR bias_gm, GM_ADDR o_gm,
    int m, int n, int k)
{
    TPipe pipe;
    // 输入输出缓冲区
    TQue<QuePosition::VECIN, 2> a_buf, b_buf;
    TQue<QuePosition::VECOUT, 1> o_buf;
    pipe.InitBuffer(a_buf, m * k * sizeof(half));
    pipe.InitBuffer(b_buf, k * n * sizeof(half));
    pipe.InitBuffer(o_buf, m * n * sizeof(half));

    // 第一步：MatMul -> 这里的compute是省略写法
    LocalTensor<half> temp = o_buf.AllocTensor<half>();
    // 使用Cube单元做矩阵乘 temp = a @ b
    // ... [MatMul Compute 代码省略]

    // 第二步：BiasAdd -> 直接在temp上加bias
    // bias的维度是(n,)，要broadcast到每一行
    // 这里用Vector单元做element-wise加法
    for (int i = 0; i < m; i++) {
        LocalTensor<half> row = temp.SubTensor(i * n, n);
        // 向量加标量，每个元素加上同一个bias
        // ... [Bias Add Vector 代码省略]
    }

    // 第三步：GELU -> 原地做激活
    // GELU(x) = x * Phi(x)，Phi是高斯分布的CDF
    // 近似公式：0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
    // 注意这里要在同一个tensor上做，不能开新buffer
    // GELU需要exp和tanh，这些全在Vector单元上跑
    // ... [GELU Activation Vector 代码省略]

    // 最终结果写回HBM，只写一次
    DataCopy(o_gm, temp, m * n * sizeof(half));
}

代码里有三个关键设计点：三步计算的顺序不能乱，MatMul → BiasAdd → Activation 这个顺序是固定的，调换顺序会有问题。中间结果 temp 用的是同一个 buffer，这块内存在三步之间复用，不单独为每一步分配新 buffer。每一步用不同的计算单元，MatMul 交给 Cube 单元，BiasAdd 和 Activation 用 Vector 单元。整个过程在一个 Stream 里完成，不需要额外的同步。

融合前后性能对比

拿 LLaMA-7B 推理的一个 FFN 层（Hidden=4096, FFN Hidden=11008）测了一下：

指标	三算子分别调用	融合算子
延迟/ms	3.2	1.8
HBM 写入量/MB	96	32
Kernel Launch次数	3	1
中间Buffer数	2	0

延迟降了 44%，HBM 写入量降了 67%，Kernel Launch 从 3 次变成 1 次。这还只是一个 FFN 层，全模型 32 层乘起来省的就非常可观了。

融合算子还有一个隐形好处：中间结果不走 HBM，所以不会有 HBM 带宽争抢的问题。在多 Stream 并行的情况下，不同算子的 HBM 访问会互相抢占带宽，融合成一次之后就没有这个问题了。

当然融合不是万能的。有些场景不能融合：需要分别在不同的设备上做计算（比如某些算子在 CPU 上，某些在 NPU 上）、中间结果需要被多个后续算子复用（融合之后只有输出，没有中间结果可用）、融合后的 kernel 太大导致 L1 Buffer 不够（这种情况下反而不如分开跑）。

如何调用融合算子

普通 PyTorch 代码不用改，融合是框架层自动选择的：

import torch
import torch_npu

# 创建一个带激活函数的Linear层
# PyTorch在昇腾NPU上会自动把这个调用成融合算子
linear = torch.nn.Linear(4096, 11008, bias=True).npu()
act = torch.nn.GELU().npu()

# Forward时框架会自动走融合算子
x = torch.randn(1, 4096, dtype=torch.float16).npu()
# 这里是linear + gelu两步骤，不是三步骤
# BiasAdd已经融合在linear里了
out = act(linear(x))

# 这个例子里没有单独调用BiasAdd
# PyTorch的Linear自带bias，forward时内部已经包含了bias加法
# 所以框架看到的计算图是Linear -> GELU
# CANN的图优化会把这个 fuse 成单个算子

如果想看融合是否真的生效，可以用 CANN 的 prolyger 工具看优化后的计算图：

import os
os.environ['GE_OP_TYPE'] = 'fusion_detail'

# 再跑一次，会打印融合详情
out = act(linear(x))
# 或者用 ascend-center 的调试接口
# CANN 8.0+ 可以用 atc --dump_graph=origin,optimized 看图变化

输出了可以看到类似 matmul_bias_gelu_fused 这种融合后的算子名。如果看到的还是单独的 MatMul、BiasAdd、GELU，说明没融合成功，去查一下 CANN 版本和算子库版本是否匹配。

融合算子是最简单也最有效的性能优化手段。昇腾CANN 的融合算子在框架层自动启用，一般情况下不需要手动调。但知道这个原理之后，遇到性能问题可以去检查融合是否生效——没融合的话，手动调一下数据 layout 或升级 CANN 版本往往就能解决。

仓库地址：https://atomgit.com/cann/ops-nn