算子融合：减少通信开销实战
融合策略基础：水平融合、垂直融合与粒度选择
大家好，我是你们的老朋友。今天咱们聊聊算子融合里最基础、也最核心的三个策略：水平融合、垂直融合，还有融合的粒度选择。说白了，就是怎么把多个小算子“捏”成一个，减少数据搬来搬去的开销。
我刚开始做分布式训练的时候，总觉得算子越多越灵活。后来发现，GPU算得飞快，但数据在显存和计算单元之间来回倒腾，反而成了瓶颈。你想想看，一个简单的加法后面跟个ReLU，如果分开算，得把中间结果写回显存，再读出来。融合之后，数据直接在寄存器里就处理完了，省了多少事！
核心思想：算子融合的本质，是用计算换通信。多算一点，少搬一点，整体就快了。
一、水平融合：同类型算子合并
水平融合，也叫“同质融合”。就是把多个相同类型的算子，比如一连串的加法、一连串的乘法，或者一连串的激活函数，合并成一个更大的算子。
举个例子，你有一个模型，里面有好几个独立的加法操作：
# 水平融合前
c1 = a + b
c2 = d + e
c3 = f + g
# 水平融合后
# 可以合并成一个向量化加法
c1, c2, c3 = vector_add([a, d, f], [b, e, g])我个人习惯把这种融合叫做“打包处理”。为什么有效？因为GPU擅长做大规模并行计算。你把多个独立的加法打包成一个向量化操作，GPU就能一次性启动更多的计算单元，利用率自然就上去了。
我在项目中遇到过这样一个场景：一个Transformer模型里，有十几个独立的LayerNorm操作。每个LayerNorm都要启动一次kernel，来回读写显存。后来我把它们水平融合成一个batch LayerNorm，显存带宽利用率从30%直接飙到了85%。
小技巧：水平融合特别适合那种“形状相同、操作相同”的算子。比如多个形状一样的矩阵乘法、多个独立的激活函数。如果形状差异太大，强行融合反而会引入padding开销，得不偿失。
二、垂直融合：连续算子合并
垂直融合，也叫“流水线融合”。就是把计算图中前后相连的算子合并成一个。比如 Conv + BN + ReLU，这是最经典的垂直融合案例。
为什么会这样？因为Conv的输出直接喂给BN，BN的输出又直接喂给ReLU。如果不融合，每一步都要把中间结果写回显存。融合之后，数据流就像一条流水线，中间结果只存在于寄存器或L1缓存里。
# 垂直融合前
conv_out = conv2d(input, weight)
bn_out = batch_norm(conv_out)
relu_out = relu(bn_out)
# 垂直融合后
# 合并成一个 fused_conv_bn_relu kernel
output = fused_conv_bn_relu(input, weight, bn_params)嗯，这里要注意：垂直融合不是简单的代码合并。它需要你手动把BN的数学公式“折叠”进Conv的权重里。说白了，就是提前把BN的缩放因子和偏移量算好，直接写到Conv的权重和偏置里。这样推理的时候，BN那一步就完全消失了。
避坑指南：我曾经在训练阶段也尝试做垂直融合，结果发现梯度计算变得极其复杂。后来我学乖了：垂直融合主要用在推理阶段。训练阶段还是老老实实保持算子独立，方便反向传播。
垂直融合的另一个经典例子是 Attention 内部的 QKV 投影合并。原本是三个独立的线性层：Q = X * W_q, K = X * W_k, V = X * W_v。融合之后，可以变成一个大的线性层：QKV = X * [W_q, W_k, W_v]，然后通过split操作分开。这样就把三次显存读写变成了一次。
三、融合的粒度选择：粗还是细？
这是最考验经验的地方。融合粒度，说白了就是“融合到什么程度”。
太细了，比如只融合两个相邻的加法，收益微乎其微。太粗了，比如把整个ResNet的50个层全部融合成一个超级算子，那编译器和硬件根本吃不消，寄存器不够用，反而会 spill 到显存，得不偿失。
我一般遵循三个原则：
- 寄存器容量原则：融合后的算子，中间变量的总大小不能超过GPU寄存器的容量。比如一个SM有64K寄存器，你融合后需要80K的中间变量，那就得拆开。
- 计算密度原则：融合后的算子，计算量要远大于数据搬移量。如果融合后计算量只增加了10%，但数据搬移量减少了50%，那就值得做。反之，如果计算量暴增10倍，数据搬移只减少20%，那就不划算。
- 编译器友好原则：不要挑战编译器的极限。有些框架的编译器（比如XLA、TVM）对超大算子支持不好，编译时间会从秒级变成小时级。我建议一次融合不超过5-8个连续算子。
我的经验：对于大多数模型，融合粒度控制在“一个block内能处理完”是最合适的。比如一个卷积层+BN+ReLU，或者一个矩阵乘法+偏置+激活函数。这种粒度既能显著减少通信开销，又不会让编译器崩溃。
四、三种策略的对比
| 策略 | 适用场景 | 收益来源 | 风险 |
|---|---|---|---|
| 水平融合 | 多个同类型、独立算子 | 提高并行度、减少kernel启动次数 | 形状不一致时引入padding开销 |
| 垂直融合 | 前后相连的算子链 | 减少中间结果读写显存 | 训练时梯度计算复杂 |
| 粒度选择 | 所有融合场景 | 平衡计算与通信 | 过粗导致寄存器溢出 |
五、实战中的选择逻辑
你可能会问：那我到底该用哪种？
我的做法是三步走：
- 先看计算图：用profiler跑一遍，找出哪些算子的kernel启动时间占比高。如果启动时间超过总时间的20%，优先考虑水平融合。
- 再看数据流：如果某个算子的输出只被下一个算子使用，而且中间结果很大（比如特征图），那就做垂直融合。我遇到过最夸张的例子，一个特征图有256MB，每次读写都要花2ms，垂直融合后直接省掉了这2ms。
- 最后调粒度：从最小的融合单元开始（比如2个算子），逐步扩大。每次扩大后都跑一遍benchmark，直到性能不再提升为止。
一个小建议：不要试图一次性把所有算子都融合了。我见过有人把整个BERT的12层Transformer全部融合成一个算子，结果编译花了3小时，跑起来反而更慢了。融合是个精细活，慢慢来比较快。
六、知识体系总览
下面这张图，是我自己总结的算子融合策略选择流程。你可以把它当作一个决策树来用：
这张图的核心逻辑很简单：先判断瓶颈在哪，再选策略，最后调粒度。每一步都要用benchmark说话，别靠猜。
好了，关于融合策略的基础就聊到这。水平融合、垂直融合、粒度选择，这三板斧用好了，通信开销能降一半以上。下一节咱们会深入具体的实现细节，包括如何手写一个 fused kernel，以及如何在主流框架里配置自动融合。
想要更多资料，联系可可