常见融合模式：逐元素融合、归约融合、矩阵乘法融合、卷积+激活函数融合、BatchNorm融合
聊到算子融合，很多人第一反应是「把几个小算子拼成一个」。嗯，方向没错，但具体怎么拼、拼哪些，这里头门道不少。我这些年调过不少模型，踩过坑也攒了些经验，今天就把最常见的五种融合模式掰开揉碎了讲讲。
逐元素融合：最基础的「搭积木」
逐元素操作，说白了就是每个元素独立计算，互不干扰。比如加法、乘法、ReLU、Sigmoid 这些。它们的特点是计算密度低，访存开销大——你想想看，每次从显存读数据，算完又写回去，来回折腾好几次，带宽全浪费了。
核心思路：把多个逐元素操作合并成一个 kernel 执行。比如 `y = a + b` 再 `z = relu(y)`，直接合成一个 kernel 算完。
实战案例：我在项目中遇到过 Transformer 的 LayerNorm 加残差连接，里面有好几个逐元素操作。一开始分开写，训练速度慢得离谱。后来我把 `x + sublayer(x)` 和 LayerNorm 的归一化、缩放、平移全塞进一个 kernel，吞吐直接提升了 15%。
// 伪代码：逐元素融合示例
// 融合前
float* y = elementwise_add(a, b);
float* z = relu(y);
// 融合后
float* z = fused_add_relu(a, b); // 一个 kernel 搞定小技巧：逐元素融合对带宽敏感型算子效果最好。如果你的模型里有很多 `add + relu`、`mul + add` 这种组合，优先考虑融合。
归约融合：减少中间结果的读写
归约操作，比如 sum、mean、max，它们会把一堆数据压缩成一个标量。这类算子的特点是：计算量不大，但需要全局通信。如果和前面的逐元素操作分开做，中间结果得写回显存，再读出来归约，效率极低。
怎么做？把前面的逐元素操作和归约操作合并。比如 `sum(relu(x))`，直接在一个 kernel 里先算 ReLU，再累加，中间结果根本不用落盘。
我曾经踩过的坑：有一次做 BatchNorm 的推理加速，我把 `mean` 和 `var` 的计算分开写，结果发现显存带宽被中间结果占满了。后来改成融合 kernel，显存占用直接降了 30%。记住，归约融合的核心是「减少中间张量的物化」。
| 融合模式 | 典型场景 | 收益 |
|---|---|---|
| 逐元素融合 | add + relu, mul + add | 减少 kernel launch 开销 |
| 归约融合 | sum + relu, mean + square | 减少中间结果读写 |
矩阵乘法融合：计算密集型算子的「组合拳」
矩阵乘法（GEMM）是深度学习的心脏。但很多时候，GEMM 前后会跟着 bias 加法、激活函数。如果分开做，GEMM 算完的数据得写回显存，再读出来做后续操作——这简直是暴殄天物。
融合方式：把 bias 加法和激活函数直接嵌入 GEMM 的 kernel 里。比如 `C = A * B + bias`，然后 `relu(C)`，合成一个 `fused_gemm_relu`。
我个人的习惯：在写自定义算子时，只要 GEMM 后面跟着逐元素操作，我一定优先考虑融合。NVIDIA 的 cuBLAS 和 cuDNN 其实已经内置了这类融合，但如果你用 PyTorch 的 `torch.nn.functional.linear`，它默认是不融合的。我建议用 `torch.jit.script` 或者手动写 CUDA kernel 来实现。
// 伪代码：矩阵乘法融合
// 融合前
C = gemm(A, B);
C = add_bias(C, bias);
C = relu(C);
// 融合后
C = fused_gemm_bias_relu(A, B, bias); // 一个 kernel卷积+激活函数融合：CNN 里的「黄金搭档」
卷积操作后面跟激活函数，几乎是 CNN 的标准配置。比如 Conv2D + ReLU。分开写的话，卷积输出的 feature map 得先存下来，再读出来做 ReLU。这中间的数据量可不小——尤其是当通道数多的时候。
融合原理：在卷积的每个输出元素计算完成后，立即应用激活函数，不写回显存。这样既省了带宽，又减少了 kernel 启动次数。
避坑指南：不是所有激活函数都适合融合。像 ReLU 这种简单的分段函数，融合起来毫无压力。但像 Swish（`x * sigmoid(x)`）这种复杂函数，融合后可能会影响卷积的计算效率。我建议先做 profiling，看看瓶颈在哪。
BatchNorm 融合：推理时的「杀手锏」
BatchNorm 在训练时是独立算子，但在推理时，它的参数是固定的。这时候我们可以把 BatchNorm 的缩放和平移操作「吸收」到前面的卷积或全连接层里。
怎么做？假设卷积层输出是 `y = W * x + b`，BatchNorm 是 `z = gamma * (y - mean) / sqrt(var + eps) + beta`。把这两个公式合并，得到新的权重 `W'` 和偏置 `b'`，推理时直接用新参数算，BatchNorm 就消失了。
// 伪代码：BatchNorm 融合
// 原始卷积 + BatchNorm
y = conv(x, W, b);
z = batchnorm(y, gamma, beta, mean, var);
// 融合后
W' = gamma * W / sqrt(var + eps);
b' = gamma * (b - mean) / sqrt(var + eps) + beta;
z = conv(x, W', b'); // 直接算，没有 BatchNorm 了注意：BatchNorm 融合只适用于推理阶段。训练时 mean 和 var 是动态的，不能提前融合。另外，如果你的模型用了 GroupNorm 或 LayerNorm，融合方式会不一样，别搞混了。
知识体系图：五种融合模式一览
这五种模式，说白了就是「减少中间环节」。你想想看，GPU 最怕的就是「等数据」——等从显存读、等写回显存。融合就是把这些等待时间压缩到最低。我个人建议，刚开始做算子融合时，先从逐元素融合和 BatchNorm 融合入手，这两个最容易上手，收益也最明显。
嗯，今天就聊到这儿。记住一点：融合不是万能的，但该融合时不融合，那你的训练吞吐就永远上不去。