常见融合模式：Conv+Bn+ReLU、Conv+Add+ReLU、MatMul+Softmax、LayerNorm
各位好，我是老张。今天咱们聊聊算子融合里最经典的几种模式。说实话，这几种融合模式我几乎在每个部署项目里都会遇到。你想想看，一个典型的CNN网络里，Conv+Bn+ReLU这种组合能占到总计算量的60%以上。不把它们融合掉，推理速度根本提不起来。
我个人习惯把融合分为两类：一类是“计算密集型”融合，比如Conv和它的后续操作；另一类是“访存密集型”融合，比如LayerNorm这种。前者省计算，后者省带宽。咱们一个一个来看。
Conv+Bn+ReLU 融合
这个组合太常见了。几乎每个CNN模型里都有。为什么能融合？说白了，BatchNorm在推理时就是一个线性变换：y = gamma * (x - mean) / sqrt(var + eps) + beta。而ReLU就是个简单的max(0, x)。
我在项目中遇到过一个问题：有些框架在训练时会把BN的四个参数（gamma、beta、mean、var）分开存储，但推理时我们只需要一个等效的线性系数。嗯，这里要注意，融合的关键就是把BN的线性变换“吸收”进Conv的权重里。
融合公式（简洁版）：
// 原始：Conv(x) -> BN -> ReLU
// 融合后：Conv_fused(x) -> ReLU
// 新权重：W_fused = W * gamma / sqrt(var + eps)
// 新偏置：b_fused = (b - mean) * gamma / sqrt(var + eps) + beta我的经验：融合后精度几乎不变，但速度能提升20%-30%。不过要注意，如果BN层后面跟着的是LeakyReLU或者PReLU，融合公式要稍微调整一下——把斜率也吸收进去。
Conv+Add+ReLU 融合
这个模式在ResNet、DenseNet这些带残差连接的网络里特别多。你想想看，一个残差块里，主路径上的Conv输出要和旁路的输入相加，然后过ReLU。
为什么能融合？因为Add操作本身没有非线性，就是个逐元素加法。我们可以把Add的输入“提前”融合到Conv的计算里。具体怎么做？
- 方法一：如果Add的输入是前一层Conv的输出，且两个Conv的输入尺寸相同，可以把两个Conv合并成一个更大的Conv（这个比较少见，因为会改变感受野）。
- 方法二（更常用）：把Add操作和后面的ReLU合并成一个“带偏置的ReLU”。说白了，就是先算Conv，然后加上残差输入，最后过ReLU。这个在硬件上可以做成一个算子，减少一次kernel launch。
我曾经踩过的坑：有一次我把Conv+Add+ReLU融合后，发现精度掉了0.3%。排查了半天，原来是Add的输入来自两个不同精度的分支——一个float16，一个int8。融合时没做精度对齐，导致数值误差累积。所以，融合前一定要检查所有输入的数据类型是否一致。
MatMul+Softmax 融合
这个模式在Transformer里太关键了。Attention机制的核心就是：Q * K^T -> Softmax -> * V。其中MatMul和Softmax是紧挨着的。
为什么要融合？因为Softmax需要全局的max和sum，如果分开算，MatMul的结果要写回显存，然后Softmax再读出来。这一写一读，带宽就浪费了。融合后，MatMul的结果直接留在寄存器里，Softmax就地计算。
融合后的计算流程：
// 原始流程：
// 1. MatMul: S = Q * K^T (写回显存)
// 2. Softmax: P = softmax(S) (从显存读S，写回P)
// 融合流程：
// 1. 逐行计算S = Q * K^T
// 2. 在寄存器里计算该行的max和sum
// 3. 直接计算softmax并写回P
// 全程S不落显存我个人建议，在做MatMul+Softmax融合时，一定要考虑数值稳定性。Softmax里有个exp(x - max)的操作，如果MatMul的输出值很大（比如超过20），exp可能会溢出。所以融合时要加上max的在线计算。
LayerNorm 融合
LayerNorm在Transformer里也是标配。它本身的计算包括：求均值、求方差、归一化、缩放平移。这四个步骤如果分开做，要来回读写显存好几次。
融合的思路很简单：把四个步骤合并成一个kernel，一次性读完输入，在寄存器里算完所有中间结果，最后写回输出。
| 步骤 | 原始实现（多次访存） | 融合实现（一次访存） |
|---|---|---|
| 1. 求均值 | 读输入，写mean | 读一次输入，寄存器内完成所有计算，写一次输出 |
| 2. 求方差 | 读输入和mean，写var | |
| 3. 归一化 | 读输入、mean、var，写norm | |
| 4. 缩放平移 | 读norm、gamma、beta，写输出 |
我的建议：LayerNorm融合时，要注意输入数据的布局。如果是[N, H, W, C]这种格式，求均值和方差是在最后一个维度上做的，这时候用向量化加载（比如float4）能大幅提升带宽利用率。我见过有人用标量加载，结果带宽利用率只有20%。
四种融合模式对比
好了，四种模式都讲完了。我画了一张图，帮你理清它们的关系和适用场景。
这四种模式，我建议你从Conv+Bn+ReLU开始练手。为什么？因为它最成熟，公式最固定，踩坑的概率最小。等你把这个搞定了，再去挑战MatMul+Softmax这种需要在线计算的融合。
好了，今天就聊到这儿。记住一句话：融合不是万能的，但不融合是万万不能的。尤其是在边缘设备上，一次显存读写可能就是几微秒的延迟，能省则省。