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 的自动融合策略，以及一些高级的融合技巧。

← 上一章 📖 返回目录 下一章 →