算子融合：从数学等价到编译器合法性的严格推导

本文围绕四个典型融合模式展开：

1. MatMul → Bias Add → ReLU
2. Conv → Bias → ReLU
3. Add → ReLU
4. Residual Add → LayerNorm

重点回答两个问题：

• 为什么这些算子在数学和数据依赖上可以融合？
• 哪些情况不能融合，或者虽然能融合但不应该融合？

---

1. 算子融合的本质

假设原计算图是：

X
↓
P
↓
T
↓
C
↓
Y

其中：

• P：Producer，生产者；
• C：Consumer，消费者；
• T：中间 Tensor。

数学上：

T=P(X) T=P(X) T=P(X)

Y=C(T) Y=C(T) Y=C(T)

代入：

Y=C(P(X)) Y=C(P(X)) Y=C(P(X))

如果编译器能够在同一个 Kernel 内保持完全相同的计算语义，那么中间结果 T 就不一定需要写回 HBM/DDR。

原始执行：

X
↓
P
↓
T 写 HBM
↓
C 从 HBM 读 T
↓
Y

融合执行：

X
↓
P
↓
Register / Accumulator / SRAM 中的临时值
↓
C
↓
Y

因此：

算子融合不是消除数据依赖，而是在保持数据依赖不变的前提下，消除不必要的 Kernel 边界和中间 Tensor 的物化（materialization）。

---

2. 算子能够融合的基本合法性条件

两个算子能够融合，通常至少要满足以下条件。

2.1 数据依赖必须能够在同一个 Kernel 内保持

例如：

MatMul → Bias → ReLU

依赖顺序仍然必须是：

MatMul
↓
Bias
↓
ReLU

融合并不允许错误重排。

2.2 中间结果不需要被程序外部观察

例如：

%0 = matmul(...)
%1 = relu(%0)
return %1

如果 %0 除了给 relu 使用之外没有其他用途，那么 %0 可以只存在于寄存器或片上 SRAM 中。

但如果：

          ┌→ Consumer A
MatMul → T
          └→ Consumer B

那么直接把 T 完全消掉就没有那么简单，因为两个 Consumer 都需要它。

2.3 算子不能有不可忽略的副作用

纯数学算子：

Add
Mul
MatMul
Conv
ReLU
GELU

一般比较容易融合。

但如果算子涉及：

文件 I/O
Host callback
外部设备状态
不可重放随机状态
特殊原子更新

就不能简单按照纯函数进行重排或消除。

2.4 同一个 Kernel 必须能够完成所需同步

例如：

Add → ReLU

不需要线程间同步。

而：

Residual Add → LayerNorm

需要：

Reduction
+
线程间同步

如果 Reduction 的范围超过单个 CTA/Workgroup 能够同步的范围，就可能不能在一个普通 Kernel 内完成。

2.5 融合后的资源必须可接受

融合可能增加：

Register Pressure
Shared Memory / SRAM
Live Range
同步开销
代码尺寸

所以：

"数学上可以融合"不等于"性能上一定应该融合"。

---

3. MatMul → Bias Add → ReLU

这是典型的 Epilogue Fusion（尾处理融合 / 后处理融合）。

3.1 MatMul

设：

A∈RM×K A\in \mathbb{R}^{M\times K} A∈RM×K

B∈RK×N B\in \mathbb{R}^{K\times N} B∈RK×N

矩阵乘法：

$$

C_{ij}

\sum_{k=0}^{K-1}A_{ik}B_{kj}

其中： * iii：输出行； * jjj：输出列； * kkk：Reduction 维。 ### 3.2 Bias Add 设 Bias： b∈RN b\\in\\mathbb{R}\^{N} b∈RN 对输出列进行广播： ##

D_{ij}

C_{ij}+b_j

将 MatMul 代入： ##

D_{ij}

\sum_kA_{ik}B_{kj}+b_j

### 3.3 ReLU ReLU 定义： ReLU⁡(x)=max⁡(0,x) \\operatorname{ReLU}(x)=\\max(0,x) ReLU(x)=max(0,x) 所以： ##

Y_{ij}

\max(0,D_{ij})

继续代入： ##

\boxed{
Y_{ij}

\max

\left(

0,

\sum_k A_{ik}B_{kj}+b_j

\right)

}

这说明三个算子可以组成一个完全确定的复合函数。 ### 3.4 为什么适合 Kernel Fusion GEMM 内部本来就存在 accumulator： ```cpp acc = 0; for (k) acc += A[i][k] * B[k][j]; ``` 当 K Reduction 完成： ```text acc = C[i,j] ``` 传统执行： ```text acc ↓ store C 到 HBM ↓ Bias Kernel 再读 C ↓ 写 D ↓ ReLU Kernel 再读 D ↓ 写 Y ``` 融合以后： ```cpp acc = 0; for (k) acc += A[i][k] * B[k][j]; acc += bias[j]; acc = max(acc, 0); Y[i][j] = acc; ``` 数据路径变成： ```text A/B ↓ MMA / MAC ↓ Accumulator ↓ + Bias ↓ ReLU ↓ 最终 Store ``` 中间 `C`、`D` 都不必完整写回 HBM。 ### 3.5 为什么 ReLU 不能一般性折叠进 MatMul 权重 MatMul + Bias 是仿射变换： AB+b AB+b AB+b 但 ReLU： max⁡(0,x) \\max(0,x) max(0,x) 是非线性函数。 一般不存在固定矩阵 W′W'W′、Bias b′b'b′，使得： ReLU⁡(AB+b)=AW′+b′ \\operatorname{ReLU}(AB+b)=AW'+b' ReLU(AB+b)=AW′+b′ 对任意输入都成立。 因此： > ReLU 不能一般性地"折叠进权重"，但可以作为同一 Kernel 的 Epilogue 融合执行。 *** ** * ** *** ## 4. Conv → Bias → ReLU 它和 GEMM 的结构非常相似。 ### 4.1 Conv 以二维卷积为例。 输入： Xn,c,h,w X_{n,c,h,w} Xn,c,h,w 权重： Wo,c,r,s W_{o,c,r,s} Wo,c,r,s 输出： ##

C_{n,o,p,q}

\sum_{c,r,s}

X_{n,c,;pS_h+rD_h-P_h,;qS_w+sD_w-P_w}

W_{o,c,r,s}

这里： * nnn：batch； * ooo：output channel； * ccc：input channel； * p,qp,qp,q：输出空间位置； * r,sr,sr,s：卷积核位置。 ### 4.2 Bias Bias 通常沿 output channel： ##

D_{n,o,p,q}

C_{n,o,p,q}+b_o

代入： ##

D_{n,o,p,q}

\sum_{c,r,s}XW+b_o

### 4.3 ReLU ##

Y_{n,o,p,q}

\max(0,D_{n,o,p,q})

因此： ##

\boxed{
Y_{n,o,p,q}

\max

\left(

0,

\sum_{c,r,s}

X_{...}W_{o,c,r,s}+b_o

\right)

}

### 4.4 为什么能融合 Conv 的一个输出元素也是 Reduction： ```text 很多 X × W ↓ Accumulator ↓ 一个 C[n,o,p,q] ``` 因此可以在输出仍然位于 accumulator 时： ```text Conv Accumulator ↓ + Bias ↓ ReLU ↓ Store ``` 所以这是典型的 **Conv Epilogue Fusion**。 *** ** * ** *** ## 5. Add → ReLU 这是最简单的 **Elementwise Fusion（逐元素融合 / 元素级融合）**。 ### 5.1 Add Zi=Ai+Bi Z_i=A_i+B_i Zi=Ai+Bi ### 5.2 ReLU ##

Y_i

\max(0,Z_i)

代入： Yi=max⁡(0,Ai+Bi) \\boxed{ Y_i=\\max(0,A_i+B_i) } Yi=max(0,Ai+Bi) 每个位置 iii 完全独立。 ```cpp float a = A[i]; float b = B[i]; float z = a + b; float y = max(z, 0); Y[i] = y; ``` 中间 `z` 只存在于寄存器。 因此没有： ```text 线程间 Reduction 线程间通信 Barrier ``` 所以它是最容易融合的一类。 *** ** * ** *** ## 6. Residual Add → LayerNorm 这是比前三者复杂得多的 **Reduction Fusion（归约融合）**。 ### 6.1 Residual Add 假设一个 token 的 hidden vector 为： Xn=\[Xn,0,Xn,1,...,Xn,D−1\] X_n=\[X_{n,0},X_{n,1},...,X_{n,D-1}\] Xn=\[Xn,0,Xn,1,...,Xn,D−1\] Residual branch： Rn=\[Rn,0,Rn,1,...,Rn,D−1\] R_n=\[R_{n,0},R_{n,1},...,R_{n,D-1}\] Rn=\[Rn,0,Rn,1,...,Rn,D−1\] Residual Add： ##

Z_{n,j}

X_{n,j}+R_{n,j}

注意： > "Residual"虽然来源于残差 F(x)=H(x)−xF(x)=H(x)-xF(x)=H(x)−x，实际残差连接的计算形式是 H(x)=x+F(x)H(x)=x+F(x)H(x)=x+F(x)。 ### 6.2 LayerNorm 的均值 LayerNorm 沿 hidden dimension DDD 求均值： ##

\mu_n

\frac{1}{D}

\sum_{j=0}^{D-1}

Z_{n,j}

把 Residual Add 代入： ##

\boxed{
\mu_n

\frac1D

\sum_j

(X_{n,j}+R_{n,j})

}

这里： ```text D 个 Z 值 ↓ Sum Reduction ↓ 1 个 sum ↓ 除以 D ↓ μ_n ``` ### 6.3 LayerNorm 的方差 ##

\sigma_n^2

\frac1D

\sum_j

(Z_{n,j}-\mu_n)^2

代入： ##

\boxed{
\sigma_n^2

\frac1D

\sum_j

(X_{n,j}+R_{n,j}-\mu_n)^2

}

LayerNorm 不是把数据"变成正态分布"，而是对这一组实际数值做标准化。 ### 6.4 Normalize 标准化： ##

\hat Z_{n,j}

\frac{

Z_{n,j}-\mu_n

}{

\sqrt{\sigma_n^2+\epsilon}

}

不是"除以方差"，而是除以标准差： σ2+ϵ \\sqrt{\\sigma\^2+\\epsilon} σ2+ϵ 代入 Residual： ##

\hat Z_{n,j}

\frac{

X_{n,j}+R_{n,j}-\mu_n

}{

\sqrt{\sigma_n^2+\epsilon}

}

### 6.5 Scale + Shift LayerNorm 最终还有可学习参数： ##

Y_{n,j}

\gamma_j\hat Z_{n,j}+\beta_j

所以完整公式： ##

\boxed{
Y_{n,j}

\gamma_j

\frac{

X_{n,j}+R_{n,j}-\mu_n

}{

\sqrt{\sigma_n^2+\epsilon}

}

+\beta_j

}

其中： ##

\mu_n

\frac1D\sum_j(X_{n,j}+R_{n,j})

##

\sigma_n^2

\frac1D

\sum_j

(X_{n,j}+R_{n,j}-\mu_n)^2

### 6.6 为什么 Residual Add 可以和 LayerNorm 融合 原始计算： ```text X ─┐ ├→ Add → Z 写 HBM R ─┘ ↓ LayerNorm 再读 Z ↓ Y ``` 但 LayerNorm 真正需要的是： ```text Z[j] = X[j] + R[j] ``` 因此完全可以在 LayerNorm Kernel 内边计算 `X[j] + R[j]`，边参与 `sum / sumsq` Reduction。 例如第一阶段： ```cpp for (j) { z = X[j] + R[j]; sum += z; sumsq += z * z; } ``` 得到均值和方差。实际高质量实现通常会考虑数值稳定性，例如使用 Welford 算法。 ### 6.7 为什么这种融合更复杂 对于： ```text Add → ReLU ``` 算出 `Z[0]` 后就可以立刻得到 `Y[0]`。 但 LayerNorm 不行，因为： Yn,0 Y_{n,0} Yn,0 需要： μn,σn2 \\mu_n,\\sigma_n\^2 μn,σn2 而这两个量依赖整行： ```text Z[0] Z[1] ... Z[D-1] ``` 所以需要： ```text Z[0] ─┐ Z[1] ─┤ Z[2] ─┼→ Reduction → μ, σ² ... ─┤ Z[D] ─┘ ↓ Normalize 所有元素 ``` 因此需要： ```text 线程间 Reduction + 同步 + 寄存器 / SRAM / Shared Memory ``` *** ** * ** *** ## 7. 四种情况的统一分类 | 计算模式 | 典型类型 | 核心特征 | |--------------------------|--------------------|-------------------------| | Add → ReLU | Elementwise Fusion | 每个元素独立 | | MatMul → Bias → ReLU | Epilogue Fusion | 主计算后利用 accumulator 做后处理 | | Conv → Bias → ReLU | Epilogue Fusion | Reduction 主计算后做逐元素后处理 | | Residual Add → LayerNorm | Reduction Fusion | 需要跨元素 Reduction 和同步 | *** ** * ** *** ## 8. 哪些算子通常容易融合 ### 8.1 Elementwise → Elementwise 例如： ```text Add → ReLU Mul → Add Add → GELU Bias → SiLU Scale → Clamp ``` 原因： Yi=f(g(Xi)) Y_i=f(g(X_i)) Yi=f(g(Xi)) 每个元素独立。 ### 8.2 GEMM / Conv → Elementwise Epilogue 例如： ```text MatMul → Bias MatMul → Bias → ReLU MatMul → Bias → GELU Conv → Bias → ReLU ``` 原因： ```text Reduction 主计算 ↓ 输出已经在 accumulator ↓ 继续做逐元素后处理 ``` ### 8.3 Producer → Reduction 例如： ```text Residual Add → LayerNorm Add → ReduceSum Sub → ReduceMax ``` 可以融合，但实现复杂，需要： ```text Reduction 同步 片上缓存 ``` ### 8.4 Prologue Fusion 例如： ```text Dequantize → MatMul Cast → MatMul Scale → MatMul ``` 可以边加载数据边完成前处理，不必把转换后的完整 Tensor 写入 HBM。 *** ** * ** *** ## 9. 哪些情况不能直接融合 需要区分： 1. **语义上不能直接融合** 2. **语义上可以，但实现或性能上不值得融合** *** ** * ** *** ## 10. 情况一：需要全局同步，但单 Kernel 无法提供 假设： ```text Huge Tensor ↓ Global Reduce ↓ 后续所有元素都依赖 Global Reduce 的结果 ``` 例如 Reduction 范围跨很多 CTA： ```text CTA0 ─┐ CTA1 ─┤ CTA2 ─┼→ Global Sum CTA3 ─┘ ↓ 后续计算 ``` 普通 GPU Kernel 内的 block-level barrier 只能同步一个 CTA 内线程。 如果没有合适的 Grid-level synchronization，就必须拆成多个 Kernel： ```text Kernel 1： 各 CTA 算 partial reduction ↓ 写 partial result Kernel 2： 做 global reduction ↓ 得到最终统计值 Kernel 3： 消费统计值 ``` *** ** * ** *** ## 11. 情况二：中间 Tensor 被多个 Consumer 使用 例如： ```text ┌→ ReLU MatMul → C ──┤ └→ AnotherOp ``` 如果把 MatMul + ReLU 融合，并完全不保存 C，那么 AnotherOp 就失去输入。 可选方案： ```text A. 仍然 materialize C B. 重复计算 MatMul C. 把更多 Consumer 一起融合 ``` 但都会增加成本。 *** ** * ** *** ## 12. 情况三：不可忽略的副作用 例如： ```text 文件 I/O 全局状态更新 Host callback 特殊设备操作 ``` 这些不能当作普通纯函数随意删除、重排或融合。 *** ** * ** *** ## 13. 情况四：内存别名导致依赖不明确 例如： ```cpp *p = value; x = *q; ``` 如果： ```text p == q ``` 那么存在 Memory Dependency。 编译器如果不能证明： ```text p != q ``` 就不能随便交换相关操作。 这依赖 Alias Analysis。 *** ** * ** *** ## 14. 情况五：严格浮点语义不允许改变 Reduction 顺序 浮点加法不满足严格结合律： (a+b)+c≠a+(b+c) (a+b)+c \\neq a+(b+c) (a+b)+c=a+(b+c) Reduction Fusion 为了并行，可能把： ```text a + b + c + d ``` 改成树形： ```text (a+b) + (c+d) ``` 如果编译模式要求严格保持原运算顺序，那么某些 Reduction 重排就不合法。 所以： ```text fast-math ``` 与： ```text strict floating-point ``` 的优化空间不同。 *** ** * ** *** ## 15. 情况六：训练时中间结果必须给 Backward 使用 Forward： ```text X ↓ Op ↓ T ↓ NextOp ``` Backward 可能需要 `T`。 如果 Forward Fusion 完全不保存 T，那么 Backward 就拿不到它。 可选： ```text 保存 T ``` 或者： ```text Rematerialization / Recomputation ``` 即反向时重新计算。 因此： > Forward Fusion 和 Activation Saving 是典型的计算---内存权衡。 *** ** * ** *** ## 16. 可以融合但不应该融合的典型情况 ### 16.1 Register Pressure 过高 融合后大量中间值同时存活： ```text Live Range ↑ ↓ Register Pressure ↑ ↓ Occupancy ↓ ``` 严重时： ```text Register Spill ↓ Local Memory ``` 性能可能更差。 ### 16.2 Shared Memory / SRAM 超限 例如： ```text Producer Tile + Consumer Tile + Double Buffer + Reduction Buffer ``` 全部融合后： ```text Shared Memory > Hardware Limit ``` Kernel 可能无法 Launch，或者只能严重降低并发度。 ### 16.3 Occupancy 太低 例如： ```text 融合前： 8 CTA / SM 融合后： 1 CTA / SM ``` 虽然减少 HBM Traffic，但 TLP 大幅下降。 ### 16.4 重复计算成本大于内存收益 如果融合要求 Producer 重复执行： ```text Add ``` 通常问题不大。 但如果 Producer 是： ```text 大型 MatMul 大型 Conv ``` 重新计算可能得不偿失。 ### 16.5 Layout 不兼容 例如 Producer 最优 Layout 和 Consumer 最优 Layout 不一致。 融合后如果需要昂贵的： ```text Transpose Shuffle Shared Memory Exchange ``` 融合未必更快。 ### 16.6 不同设备执行 例如： ```text Op A → GPU Op B → NPU ``` 两个算子属于不同设备，通常不能融合成一个硬件 Kernel。 *** ** * ** *** ## 17. Fusion 与 Folding 的区别 ### Fusion 数学运算仍然存在，只是放进一个 Kernel。 例如： ```text MatMul + Bias + ReLU ``` 仍然执行这三个运算，只是不再产生多个 Kernel 和中间 HBM Tensor。 ### Folding 直接改变参数，使某个算子数学上消失。 典型： ```text Conv → BatchNorm ``` Conv： z=Wx+b z=Wx+b z=Wx+b BatchNorm： y=γz−μσ2+ϵ+β y= \\gamma \\frac{z-\\mu}{\\sqrt{\\sigma\^2+\\epsilon}} +\\beta y=γσ2+ϵ z−μ+β 定义： α=γσ2+ϵ \\alpha= \\frac{\\gamma}{\\sqrt{\\sigma\^2+\\epsilon}} α=σ2+ϵ γ 得到： y=α(Wx+b−μ)+β y= \\alpha(Wx+b-\\mu)+\\beta y=α(Wx+b−μ)+β 定义： W′=αW W'=\\alpha W W′=αW b′=α(b−μ)+β b'=\\alpha(b-\\mu)+\\beta b′=α(b−μ)+β 于是： y=W′x+b′ \\boxed{ y=W'x+b' } y=W′x+b′ 因此： ```text Conv + BN ``` 可以变成： ```text New Conv ``` BN 运算本身被消掉。 *** ** * ** *** ## 18. 算子融合常见分类 ### Elementwise Fusion 中文： > **逐元素融合 / 元素级融合** 例如： ```text Add → ReLU ``` ### Epilogue Fusion 中文： > **尾处理融合 / 后处理融合** 例如： ```text MatMul → Bias → ReLU Conv → Bias → GELU ``` ### Prologue Fusion 中文： > **前处理融合** 例如： ```text Dequant → MatMul ``` ### Reduction Fusion 中文： > **归约融合** 例如： ```text Residual Add → LayerNorm ``` ### Producer-Consumer Fusion 中文： > **生产者---消费者融合** 例如： ```text Conv → ReLU ``` ### Vertical Fusion 中文： > **纵向融合 / 链式融合** 沿依赖链： ```text A → B → C ``` 进行融合。 ### Horizontal Fusion 中文： > **横向融合 / 并行融合** 多个原本独立的 Kernel 合并执行，以减少 Launch 开销。 *** ** * ** *** ## 19. 实用的 Fusion 判断流程 看到： ```text Op A ↓ Op B ``` 可以按下面顺序判断。 ### 第一步：数据依赖 问： ```text B 到底需要 A 的哪些数据？ ``` 如果只是一个元素，最容易融合。 如果需要整个 Tile，难度上升。 如果需要全局 Tensor 的统计结果，可能需要多阶段 Kernel。 ### 第二步：中间 Tensor 是否有其他 Consumer ```text A → B ``` 最好。 如果： ```text A → B → C ``` 融合难度明显上升。 ### 第三步：有没有 Reduction ```text 没有 Reduction → 通常容易 有 CTA 内 Reduction → 可以，但需要同步 有跨 CTA Global Reduction → 可能需要多个 Kernel ``` ### 第四步：能否留在片上 判断： ```text Register Shared Memory SRAM Accumulator ``` 是否能保存所需中间数据。 ### 第五步：资源代价 检查： ```text Register Pressure Shared Memory Occupancy Code Size Synchronization ``` ### 第六步：Cost Model 比较： ```text 减少的 HBM Traffic + 减少的 Kernel Launch ``` 与： ```text 增加的 Register + SRAM + 同步 + 重复计算 + Occupancy 损失 ``` 只有收益大于成本，才应该执行融合。 *** ** * ** *** ## 20. 四个例子的最终统一数学表达 ### MatMul + Bias + ReLU ##

\boxed{
Y_{ij}

\operatorname{ReLU}

\left(

\sum_kA_{ik}B_{kj}+b_j

\right)

}

类型： > Epilogue Fusion ### Conv + Bias + ReLU ##

\boxed{
Y_{n,o,p,q}

\operatorname{ReLU}

\left(

\sum_{c,r,s}XW+b_o

\right)

}

类型： > Epilogue Fusion ### Add + ReLU ##

\boxed{
Y_i

\operatorname{ReLU}(A_i+B_i)

}

类型： > Elementwise Fusion ### Residual Add + LayerNorm Zn,j=Xn,j+Rn,j Z_{n,j}=X_{n,j}+R_{n,j} Zn,j=Xn,j+Rn,j ##

\mu_n

\frac1D

\sum_j Z_{n,j}

##

\sigma_n^2

\frac1D

\sum_j

(Z_{n,j}-\mu_n)^2

最终： ##

\boxed{
Y_{n,j}

\gamma_j

\frac{

X_{n,j}+R_{n,j}-\mu_n

}{

\sqrt{\sigma_n^2+\epsilon}

}

+\beta_j

}

$$

类型：

Reduction Fusion

---

21. 最核心结论

算子融合的核心不是：

两个 Op 名字看起来能不能放一起

而是：

Data Dependency
        ↓
Iteration Domain
        ↓
Reduction Dependency
        ↓
Memory Lifetime
        ↓
Synchronization Scope
        ↓
Hardware Resource
        ↓
Cost Model

最终决定：

能不能融合
+
值不值得融合

可以浓缩成一句话：

只要 Producer 的输出能够在不改变程序语义的情况下，直接作为 Consumer 的片上输入，并且同步范围、内存容量、寄存器压力和执行资源都允许，就可以进行 Kernel Fusion；是否真正执行融合，则还要由 Cost Model 判断性能收益。