1. 算子融合基础:什么是算子融合、为什么需要融合、融合的数学原理与计算图优化

大家好,我是你们的老朋友。今天咱们聊聊算子融合——这个在推理加速里绕不开的话题。

我记得刚入行那会儿,调模型推理速度,恨不得把每一行代码都抠一遍。后来发现,真正拖慢速度的,往往不是单个算子有多慢,而是算子之间那些「看不见的搬运工」——内存读写。说白了,你算得快没用,数据搬来搬去才是真瓶颈。

算子融合,就是把这些搬运工给优化掉。

1.1 什么是算子融合

先看个最简单的例子。你写了个 Conv + ReLU 的组合:

# 伪代码示意
def forward(x):
    y = conv2d(x, weight, bias)   # 卷积
    z = relu(y)                   # 激活
    return z

在传统框架里,这两步是分开执行的。Conv 算完,把结果写回显存;ReLU 再读出来,算完再写回去。一次读写就是一次显存带宽的消耗。

算子融合,就是把 Conv 和 ReLU 合并成一个算子。ReLU 直接在 Conv 的输出寄存器上做,不用来回倒腾。嗯,就这么简单粗暴。

核心定义:算子融合是将多个连续的计算操作合并为一个单一操作,减少中间结果的显存读写,从而降低推理延迟。

我在项目中遇到过最典型的场景:一个 MobileNet 的 depthwise conv + BN + ReLU6,三个算子分开跑,延迟 3.2ms;融合成一个算子后,直接降到 1.8ms。你想想看,这还没改任何数学逻辑,只是把「搬运」省了。

1.2 为什么需要融合

原因其实就三个字:带宽瓶颈。

现在的 GPU 算力越来越猛,但显存带宽的增长速度远远跟不上。我打个比方:算力是工厂里的机器,带宽是门口的传送带。机器再快,传送带送不过来料,也是白搭。

具体来说,融合带来的好处有:

  • 减少显存访问:中间结果不用写回显存,直接在片上缓存里完成
  • 降低 kernel launch 开销:每个算子启动都要 CPU 发指令,融合后启动次数变少
  • 提升计算密度:融合后的算子计算量更大,能更好地利用 GPU 的并行能力

我的经验:在 NVIDIA T4 上做过测试,一个 224x224 的 ResNet-50,不做融合时推理延迟约 12ms;做了 Conv+BN+ReLU 融合后,降到 8ms 左右。省下来的 4ms,够你再加一个轻量级后处理了。

1.3 融合的数学原理

融合不是乱合的,得有数学依据。说白了,就是看两个操作能不能「合并计算」。

最常见的融合模式是 线性变换 + 非线性激活。比如:

Conv:    y = W * x + b
BN:      z = γ * (y - μ) / σ + β
ReLU:    out = max(0, z)

这三步可以合并成一步:

out = max(0, (γ * W / σ) * x + (γ * (b - μ) / σ + β))

你看,新的权重和偏置可以提前算好,推理时直接一个算子搞定。

注意:不是所有算子都能融合。比如 Conv + Softmax 就不适合,因为 Softmax 需要全局的统计信息,没法在局部完成。我曾经踩过这个坑,硬合了半天,结果精度掉了 0.5%,得不偿失。

常见的可融合组合:

前序算子 后续算子 是否可融合 说明
Conv BN + ReLU ✅ 是 线性变换可合并
Conv Add + ReLU ✅ 是 残差结构常用
MatMul Add + ReLU ✅ 是 全连接层常见
Conv Softmax ❌ 否 需要全局归约
Conv Pooling ⚠️ 部分 需看具体实现

1.4 计算图优化

讲完数学,咱们看看工程上怎么落地。计算图优化,就是框架层面自动帮你做融合。

以 TensorRT 为例,它的图优化流程大致是:

  1. 解析模型:读入 ONNX 或 TensorFlow 的图结构
  2. 图分析:遍历所有节点,找出可融合的模式
  3. 模式匹配:比如匹配到 Conv + BN + ReLU 的三连击
  4. 图重写:把三个节点替换成一个融合节点
  5. 常量折叠:把能提前算的权重、偏置都算好

下面我用 SVG 画个图,帮你直观理解这个过程:

算子融合前后的计算图对比 融合前 Conv BN ReLU 融合 融合后 ConvBNReLU 融合算子 中间结果不再写回显存,减少 2 次显存读写 kernel launch 从 3 次减少到 1 次

你看,融合前三个算子各算各的,中间结果来回倒腾。融合后一个算子搞定,干净利落。

避坑指南:我曾经在 ONNX 转 TensorRT 时,发现某些算子没被融合。查了半天,原来是 BN 层的 epsilon 精度不一致导致的。框架做模式匹配时,对数值精度很敏感。建议统一用 float32 导出,等融合完再转 float16。

1.5 常见的融合模式

除了 Conv+BN+ReLU,还有几种常见的融合模式,我列出来供你参考:

  • Conv + Add + ReLU:残差网络里的标配,比如 ResNet 的 shortcut 连接
  • MatMul + Add + ReLU:全连接层的标准融合,Transformer 里也常用
  • Conv + BN + SiLU:EfficientNet 等新网络里常见,SiLU 也是 element-wise 的
  • LayerNorm + Add:Transformer 里的残差连接,可以合并计算

嗯,这里要注意:融合不是越多越好。我见过有人把整个 block 都合到一个算子里的,结果因为寄存器不够用,反而变慢了。融合的粒度要适中,一般 2-4 个算子一组比较合适。

好了,第一章就聊到这儿。算子融合的核心思想其实很简单:少搬数据,多算数。后面的章节我们会深入具体的融合实现,包括 CUDA 层面的手动融合、TensorRT 的自动融合策略,以及一些高级的融合技巧。

想要更多资料,联系可可