算子融合技术：让显存“瘦身”的魔法
算子融合，说白了就是把多个连续的小算子合并成一个大的算子。你想想看，原本要分三步走的路，现在一步到位，省时又省力。
我刚开始接触深度学习推理优化时，总觉得算子越细越好，方便调试。后来被显存逼得没办法，才真正体会到融合的威力。嗯，这里要注意，融合不是简单的代码拼接，而是从计算图和内存访问模式上做文章。
为什么需要算子融合？
每个算子执行时，都要把中间结果写回显存，下一个算子再读出来。这个读写过程，就是显存带宽的“黑洞”。
举个例子：一个简单的 Conv + ReLU，如果不融合，Conv 的输出要完整存到显存里，ReLU 再读出来处理。融合后，Conv 算出一个元素，ReLU 立刻处理，中间结果根本不用落盘。
我在项目中遇到过一个大模型，光中间激活值就占了 80% 的显存。融合几个关键算子后，显存直接砍半。那种感觉，就像给程序做了个“瘦身手术”。
常见融合模式
Conv + BN 融合
这是最经典的融合模式之一。BN 层在推理时其实就是一个线性变换：y = gamma * (x - mean) / sqrt(var + eps) + beta。
既然 Conv 也是线性变换（不考虑激活函数），那这两个线性变换完全可以合并成一个。我习惯在模型导出时就把这个融合做好，省得推理时多一次 kernel launch。
// 伪代码：Conv + BN 融合
// 原始：conv_out = Conv(x, w, b)
// bn_out = BN(conv_out, gamma, beta, mean, var)
// 融合后：fused_w = w * gamma / sqrt(var + eps)
// fused_b = (b - mean) * gamma / sqrt(var + eps) + beta
// fused_out = Conv(x, fused_w, fused_b)你看，融合后的 Conv 权重和偏置都变了，但计算结果完全一致。这就是所谓的“数学等价变换”。
显存收益分析：
不融合时，Conv 的输出需要占用 N * C * H * W * 4 字节（FP32）。融合后，这个中间张量直接消失。对于一个 256 通道、64x64 的特征图，一次融合就能省下 16MB 显存。
MatMul + Add 融合
Transformer 里到处都是这种模式：QKV 投影、FFN 层。MatMul 算完，紧接着就是 bias add。
我曾经调试过一个推理速度问题，发现 MatMul 和 Add 之间居然有显存拷贝。后来一查，是框架没做融合。手动改成融合 kernel 后，延迟降了 15%。
// 不融合：两个 kernel launch
// matmul_out = MatMul(A, B)
// add_out = Add(matmul_out, bias)
// 融合后：一个 kernel launch
// fused_out = MatMulAdd(A, B, bias)融合的好处不止省显存。你想想看，少了一次 kernel launch，就少了 GPU 调度开销。对于小 batch 场景，这个开销占比其实不小。
融合带来的显存收益
我整理了一个表格，方便你直观感受融合的效果：
| 融合模式 | 不融合显存占用 | 融合后显存占用 | 节省比例 |
|---|---|---|---|
| Conv + ReLU | 2 倍特征图大小 | 1 倍特征图大小 | 50% |
| Conv + BN + ReLU | 3 倍特征图大小 | 1 倍特征图大小 | 66.7% |
| MatMul + Add | 2 倍输出矩阵大小 | 1 倍输出矩阵大小 | 50% |
| LayerNorm + Add | 2 倍隐藏层大小 | 1 倍隐藏层大小 | 50% |
避坑指南：
我曾经把 Conv + BN + ReLU 一股脑全融合了，结果发现精度掉了 0.1%。查了半天，原来是 BN 的 eps 参数在融合时被忽略了。记住，融合时要保留 eps，否则数值误差会累积。
融合的代价
凡事都有代价。算子融合虽然省显存，但也会带来一些限制：
- 灵活性降低：融合后的算子难以单独调试。我建议在开发阶段先不融合，等模型稳定后再做优化。
- 编译时间增加：融合 kernel 的生成需要额外的编译时间。对于快速迭代的项目，可以先关掉融合。
- 硬件适配：不是所有硬件都支持融合算子。我在某款边缘设备上就遇到过融合 kernel 跑不起来的情况。
说白了，算子融合是一种“空间换时间”的优化策略。你牺牲了一点灵活性，换来了显存和速度的双重收益。对于显存紧张的场景，这招几乎是必选项。
我的个人习惯：
我会在模型导出时做一次静态融合，把 Conv+BN、MatMul+Add 这些固定模式先合并。然后在推理引擎里再做一次动态融合，根据实际输入 shape 决定是否融合。这样既保证了精度，又最大化利用了显存。
这张图很直观：不融合时，每个算子之间都要存中间结果。融合后，中间结果直接“消失”了。说白了，融合就是在减少显存里的“临时工”。
想要更多资料，联系可可