算子融合：复杂算子融合策略

可视化演示

融合策略总览

主厨第一次改造后厨时，只会做最简单的调整：切菜和焯水挨得近，就把切菜台搬到焯水锅旁边；两个灶台挨着，就合并成一个。这种“相邻算子合并”确实有用，但天花板很低——真正让厨房效率飞升的，是认出“商务套餐”这种反复出现的复杂组合，专门为它设计一条流水线。复杂算子融合策略，就是从这个“套餐思维”出发的。

在 AI 编译器中，计算图（Computation Graph）就像一份菜谱：每个算子（Operator）是菜谱上的一个步骤，算子的输入输出是食材。编译器的基础融合能减少一部分中间传输，但面对多头注意力（Multi-Head Attention, MHA）、“LayerNorm + GELU“这类复合模式时，简单合并就力不从心了——因为它们不是一个接一个的”动作“，而是一个内部有复杂数据流动的”套餐”。

复杂融合策略的核心收益，几乎都集中在减少内存访问（Memory Traffic）上。现代 GPU 算力极强，但数据从显存搬到计算单元的速度远远跟不上计算速度——这就是“访存瓶颈”。如果多个算子各自独立执行，每个都要把中间结果写回显存再读出来；融合成一个 Kernel 后，中间结果可以留在芯片上的寄存器或共享内存里，只读一次输入、写一次输出。

$$\displaystyle \text{访存缩减比} = \frac{T_{\text{fused}}}{T_{\text{orig}}} = \frac{2M}{2NM} = \frac{1}{N}$$

1. 原始访存量：不融合时，N 个算子每个都要把输入读进计算单元、把输出写回显存，读写各算一次。

$$T_{\text{orig}} = 2NM$$

2. 融合后访存量：合成一个 Kernel 后，只需读一次原始输入、写一次最终输出，中间数据留在片上。

$$T_{\text{fused}} = 2M$$

3. 两者相除：用融合后访存量除以原始访存量，得到剩余占比。

$$\frac{T_{\text{fused}}}{T_{\text{orig}}} = \frac{2M}{2NM} = \frac{1}{N}$$

主厨的经验之谈：流水线设计得再好，如果食材还在案板和灶台之间来回浪费，收益就全被运输吃掉了。算子融合的一切技巧，归根结底都是“让数据少跑路”。

模式感知融合

主厨的笔记本上记着几十种“经典套餐结构”：只要看到洋葱、青椒、牛肉三样食材出现在一起，他立刻知道这是招牌菜的前奏，整个流程早就烂熟于心。这种“认出模式、套用模板”的能力，在编译器里叫模式感知融合（Pattern-Aware Fusion）。

编译器拿到一份计算图后，会像主厨扫视食材一样，用“图模式匹配”去图上寻找是否出现了已知的复合结构：把模式库里的模板（如 MHA 的三段式结构、FFN 块里的“线性层-激活-线性层”）与计算图中的子图做同构匹配，命中就标记为可融合区域。要实现这一点，编译器需要一套精心维护的“模式库”（Pattern Library），预置深度学习模型中高频出现的复合模式，并为每种模式附带模板代码和收益预估模型。

但匹配到模式不等于必须融合——编译器还要做匹配代价评估：如果某个模式命中后，访存节省很小、或者中间张量本来就能留在缓存里，那融合反而会引入额外的调度和数值风险。理智的编译器会先算一笔账，再决定要不要动手。

$$\displaystyle \text{Fuse}(P) = \begin{cases} \text{True}, & \text{if } \frac{M_{\text{orig}} - M_{\text{fused}}}{M_{\text{orig}}} \ge \theta \\ \text{False}, & \text{otherwise} \end{cases}$$

1. 计算节省访存：模式 P 融合前的访存量减去融合后的访存量，得到省下的数据搬运量。

$$M_{\text{saved}} = M_{\text{orig}} - M_{\text{fused}}$$

2. 计算节省比例：用节省量除以原始访存量，得到相对收益，消除张量尺寸的影响。

$$\rho = \frac{M_{\text{saved}}}{M_{\text{orig}}}$$

3. 阈值判断：只有收益比例不低于预设阈值 θ 时，才启动融合。

$$\rho \ge \theta$$

主厨不会见了食材就下锅——他先算算这单菜能赚多少、备菜要花多久，值不值得动用那套复杂的流水线。模式感知融合的精髓，正是“认得出”与“值得做”缺一不可。

复杂融合Kernel

当主厨决定为招牌套餐设计专用流水线时，他得亲自下场示范每一个动作：握刀的角度、颠锅的力度、火候的切换，每个细节都关系到最终出品速度。编译器手写复杂融合 Kernel，就是这位“主厨亲自示范”的过程——它不再是简单的规则拼接，而是针对目标硬件微架构（芯片内部的寄存器、缓存、执行单元等硬件结构）精心调校的底层代码。

手写融合 Kernel 有几件趁手的“兵器”。第一件是循环重排（Loop Reordering）：把多层循环的嵌套顺序调整，让数据在内存中的访问尽量连续，提升缓存命中率。第二件是寄存器分块（Register Blocking）：把计算切分成小块，让一小块数据在寄存器里反复复用，而不是反复去取内存。第三件是向量化指令（Vectorization）：用 SIMD 指令一次处理多个数据元素，摊薄指令开销。第四件是共享内存利用（Shared Memory）：把中间结果暂存在 GPU 的片上共享内存里，Kernel 内的多个计算阶段直接接力。

衡量这些优化是否有效的核心指标，是算术强度（Arithmetic Intensity）——单位访存对应的计算量。融合 Kernel 的算术强度越高，说明数据在片上被“榨”得越充分，访存瓶颈的影响越小。

$$\displaystyle I = \frac{W}{Q} = \frac{\sum_{\text{op} \in \text{fused}} \text{FLOPs}_{\text{op}}}{M_{\text{read}} + M_{\text{write}}}$$

1. 计算量求和：把融合进来的每个算子的浮点运算次数全部加起来，得到总计算量。

$$W = \sum_{\text{op} \in \text{fused}} \text{FLOPs}_{\text{op}}$$

2. 访存量求和：该 Kernel 启动一次需要从显存读入的数据量与写出数据量之和。

$$Q = M_{\text{read}} + M_{\text{write}}$$

3. 两者相除：算术强度是单位访存对应的计算量，值越大说明访存瓶颈越弱，性能越接近硬件峰值。

$$I = \frac{W}{Q}$$

这位主厨的每个动作背后都有讲究：刀法讲究“一刀多用”，火候讲究“一波带走”，灶台讲究“顺手接力”。手写 Kernel 的艺术，就是把数据流动的每一个环节都安排得明明白白。

多头注意力融合

“商务套餐”是这家餐厅的招牌——Transformer 里的多头注意力（Multi-Head Attention, MHA）就是 AI 界的“商务套餐”。后厨一天要出几百份，每一步都值得单独优化。MHA 的第一个重复动作，就是三个投影：Q、K、V 分别要对输入做一次矩阵乘法。新手会把三份食材分别切三次，主厨则一次切完再分装——这就是 QKV 投影整合。

原来三个独立的线性层，合并成一个大的矩阵乘：把三个权重矩阵拼在一起，输入只读一次，一次算完再按列切回 Q、K、V。这省下了两次读输入的访存量，还减少两次 Kernel 启动开销。接着，注意力核心计算（查询与键的点积、softmax、再乘值向量）可以在一个 Kernel 内接力完成，中间的注意力分数矩阵（形状随序列长度平方增长）不再写回显存——这是融合中收益最大的一步。最后，输出投影与残差连接也顺势并进同一个 Kernel，整个 MHA 从“五个算子五次访存”变成“一个 Kernel 一次访存”。

$$\displaystyle XW_{\text{all}} = X \cdot [W_Q \mid W_K \mid W_V], \quad [Q \mid K \mid V] = \text{split}(XW_{\text{all}})$$

1. 拼接权重矩阵：把三个形状相同的投影权重横向拼接成一个大矩阵。

$$W_{\text{all}} = [W_Q \mid W_K \mid W_V]$$

2. 一次大矩阵乘：输入 X 只需读取一次，与拼接后的权重做一次大矩阵乘。

$$XW_{\text{all}}$$

3. 切分结果：按列将结果切回三块，得到 Q、K、V 供多头注意力使用。

$$[Q \mid K \mid V] = \text{split}(XW_{\text{all}})$$

主厨对招牌套餐的每一道工序都了如指掌：三样食材一起备、一锅出、一次装盘。MHA 融合的思路如出一辙——把整个注意力流程压缩成一条流水线，让计算密不透风地在芯片内部完成。

归一层激活组合

后厨还有一对“黄金搭档”：焯水后的青菜必须立刻过凉水锁色，然后随手淋上耗油——这两个动作在每道青菜类的菜里都会出现。Transformer 的世界里也有这么一对“黄金搭档”：LayerNorm 和 GELU。它们经常肩并肩出现（比如前馈网络 FFN 里的“LayerNorm → 线性层 → GELU”结构），单独跑的话，每步都要把中间结果写回显存，来回折腾。

LayerNorm 融合的核心，是把“求均值、求方差、归一化”这三趟遍历合并成一趟。原本一个 LayerNorm 算子内部要读多次张量，融合后可以只读一次，把统计量边算边存。再把紧随其后的 GELU 吸收进来：在归一化循环的最后一步，直接对归一化结果施加 GELU 激活，归一化后的中间张量根本不用落盘。这就是组合模式识别的价值——看到 LayerNorm 后面跟着 GELU，编译器自动套用这对黄金搭档的融合模板。

数值稳定性是融合时必须盯紧的细节。LayerNorm 在原版实现里用“先减均值再算方差”的两遍算法保证精度；融合成一遍遍历后，如果直接累加平方和，大数据量下可能损失精度，所以要采用数值稳定的在线统计算法，或者保留必要的补偿项。

$$\displaystyle \text{LN}(x_i) = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}}, \quad \mu = \frac{1}{n}\sum_{i=1}^{n} x_i, \quad \sigma^2 = \frac{1}{n}\sum_{i=1}^{n}(x_i - \mu)^2$$

1. 求均值：对输入的所有元素求和并除以元素个数，得到样本均值。

$$\mu = \frac{1}{n}\sum_{i=1}^{n} x_i$$

2. 求方差：计算每个元素与均值之差的平方平均值，得到方差。

$$\sigma^2 = \frac{1}{n}\sum_{i=1}^{n}(x_i - \mu)^2$$

3. 归一化：每个元素减去均值、除以标准差，ε 用于防止除零并保持数值稳定。

$$\hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}}$$

4. 顺带激活：在归一化同一轮循环里紧接计算 GELU，中间结果不落回内存。

$$y_i = \text{GELU}(\hat{x}_i)$$

主厨把“焯水-过凉-淋耗油”合并成一个连贯动作，手起勺落之间，青菜已经在盘子里冒着热气。LayerNorm 与 GELU 的融合，正是把这对黄金搭档安排在同一趟循环里，让数据在寄存器里“一次入锅、直接出锅”。

Transformer加速

流水线改造完成，后厨终于到了验收时刻。主厨站在出菜口，掐着秒表看订单周转时间：改造前一份套餐从下单到出餐要 15 分钟，改造后只要 6 分钟——这就是“端到端收益”。Transformer 模型加速的验收，同样要看端到端的推理延迟，而不是单个 Kernel 的威风。

在真实部署中，复杂融合策略给推理加速带来了三个层面的效果。首先是推理延迟的显著下降：以 NVIDIA TensorRT、TVM 等引擎为例，对 BERT、Llama 这类 Transformer 模型应用 MHA、LayerNorm+GELU 融合后，推理耗时通常能降低 2~3 倍。其次是内存带宽压力的释放：融合后显存访存量大幅下降，带宽占用率降低，给同卡并发更多请求留出了空间，吞吐量随之上了一个台阶。最后是端到端收益的“滚雪球”：融合省下的显存带宽和中间张量存储，让同一块 GPU 能塞下更大的批次或更长的上下文，进而摊薄其他开销。

但端到端加速不是简单相加，而是受 Amdahl 定律支配——只有被融合的那部分代码的耗时占比越大，整体加速才越明显。

$$\displaystyle S_{\text{overall}} = \frac{1}{(1-p) + \frac{p}{S_{\text{fused}}}}$$

1. 统计融合占比：复合模式占用的耗时比例越高，融合的端到端收益越大。

$$p = \frac{T_{\text{pattern}}}{T_{\text{total}}}$$

2. 测量局部加速：测量融合前后该部分 Kernel 的耗时比，得到局部加速倍数。

$$S_{\text{fused}} = \frac{T_{\text{fused}}^{\text{before}}}{T_{\text{fused}}^{\text{after}}}$$

3. 代入 Amdahl 定律：不可融合部分不受益，可融合部分按倍数加速，得到整体加速倍数。

$$S_{\text{overall}} = \frac{1}{(1-p) + \frac{p}{S_{\text{fused}}}}$$

主厨验收时最关心的是顾客等菜时间，而不是某个灶台的火有多旺。Transformer 加速的终极指标也一样——融合带来的收益必须落到端到端延迟、吞吐和资源占用上，才算真正把后厨改造成了“出菜流水线”。

AI 应用场景

Python 演示

"""复杂算子融合策略演示: 模式感知融合 + 访存收益评估""" OPS = ["QKV", "Attn", "Out", "Add", "LN", "Linear", "GELU", "Linear", "Add", "LN"] # 模式库: 预置 Transformer 中高频复合模式RULES = [    ("MHSA", ["QKV", "Attn", "Out"]),    ("NormAct", ["LN", "GELU"]),    ("FFN", ["Linear", "GELU", "Linear"]),]  def fuse_once(ops, rules):    """贪心查找最左侧命中的模式并替换为融合算子"""    for i in range(len(ops)):        for name, pat in rules:            if ops[i:i + len(pat)] == pat:                return ops[:i] + [name] + ops[i + len(pat):], name    return ops, None  def traffic(ops):    """估算访存系数: 融合算子计 1 次, 普通算子读写各计 1 次"""    return sum(1 if op in ("MHSA", "NormAct", "FFN") else 2 for op in ops)  def main():    ops = list(OPS)    orig = traffic(ops)    print("原始算子序列:", ops)    print("未融合访存系数:", orig)    print("\n--- 模式感知融合过程 ---")    while True:        ops, name = fuse_once(ops, RULES)        if name is None:            break        print("命中模式 %-6s -> %s" % (name, ops))    print("\n融合后算子序列:", ops)    print("融合后访存系数:", traffic(ops))    print("访存缩减: %.1f%%" % (100 * (1 - traffic(ops) / orig)))  if __name__ == "__main__":    main()

1. RULES = [("MHSA", ["QKV", "Attn", "Out"]), ...]模式库：预置高频复合模式，像主厨的招牌套餐模板。
2. def fuse_once(ops, rules):定义融合入口函数，一次只应用一条规则。
3. for i in range(len(ops)):从左到右扫描算子序列，寻找可命中的起点。
4. if ops[i:i + len(pat)] == pat:切片与模式模板做精确比对，判断是否命中。
5. return ops[:i] + [name] + ops[i + len(pat):], name命中后把整段模式替换为融合算子。
6. def traffic(ops):定义访存估算函数。
7. return sum(1 if op in ("MHSA", "NormAct", "FFN") else 2 for op in ops)融合算子计 1 次访存，普通算子读写各计 1 次。
8. orig = traffic(ops)记录未融合时的访存基数，用于计算缩减比例。
9. print("访存缩减: %.1f%%" % (100 * (1 - traffic(ops) / orig)))输出端到端的访存缩减百分比。

常见误区

面试题