4、算子融合（Fusion）：让计算更高效的核心技巧

算子融合，说白了就是把多个连续的小算子合并成一个大的算子。为什么要这么做？因为每次调用一个算子，都要经历「启动内核→读写数据→返回结果」这个过程。你想想看，如果能把几个算子捏在一起，中间那些临时数据的读写不就省掉了吗？

我在做模型部署时遇到过这样一个场景：一个简单的 Conv+BN+ReLU 组合，分开跑要 3 次显存读写，融合后只需要 1 次。对于大模型来说，这可不是小数目。

4.1 Conv+BN+ReLU 融合

这是最经典的融合模式，几乎每个做推理优化的工程师都会遇到。卷积后面接 BN 再接 ReLU，在训练时是标准配置，但推理时完全可以合并。

核心思路是这样的：BN 层在推理时其实就是一个线性变换——y = γ * (x - μ) / σ + β。而 Conv 也是线性操作。两个线性操作叠加，本质上还是一个线性操作。我们完全可以把 BN 的参数「吸收」到 Conv 的权重和偏置里。

融合公式推导（简化版）：

设 Conv 权重为 W，偏置为 b。BN 的缩放因子为 γ，平移为 β，均值为 μ，方差为 σ²。

融合后的权重：W' = (γ / σ) * W

融合后的偏置：b' = (γ / σ) * (b - μ) + β

然后直接接 ReLU：output = max(0, Conv(x, W', b'))

# 伪代码：Conv+BN+ReLU 融合
def fuse_conv_bn_relu(conv_weight, conv_bias, bn_gamma, bn_beta, bn_mean, bn_var):
    # 计算融合系数
    scale = bn_gamma / torch.sqrt(bn_var + 1e-5)
    
    # 融合权重
    fused_weight = conv_weight * scale.view(-1, 1, 1, 1)
    
    # 融合偏置
    fused_bias = scale * (conv_bias - bn_mean) + bn_beta
    
    # 返回融合后的 Conv + ReLU
    return fused_weight, fused_bias

# 使用方式
fused_w, fused_b = fuse_conv_bn_relu(...)
output = torch.relu(F.conv2d(input, fused_w, fused_b))

我的经验： 融合时要注意 BN 的 eps 参数。不同框架的 eps 默认值不一样（PyTorch 是 1e-5，TensorFlow 是 1e-3），直接搬公式会出精度问题。我曾经因为这个排查了整整一天。

4.2 MatMul+Add 融合

这个模式在 Transformer 里太常见了。QKV 投影、FFN 的第一层，都是 MatMul 后面跟一个偏置加法。很多框架把这叫做「FullyConnected + BiasAdd」融合。

其实实现很简单：把偏置直接加到 MatMul 的结果上。但这里有个性能陷阱——访存模式。

分开做的话，MatMul 写一次结果到显存，Add 再读一次、写一次。融合后，MatMul 在写结果时直接把偏置加上，省掉一次读写。对于大矩阵来说，这能省下几十微秒。

// CUDA 伪代码：融合的 MatMul + Bias
__global__ void fused_matmul_bias_kernel(
    float* output, 
    const float* input, 
    const float* weight, 
    const float* bias,
    int M, int N, int K
) {
    // 正常的 MatMul 计算
    float sum = 0.0f;
    for (int k = 0; k < K; k++) {
        sum += input[row * K + k] * weight[k * N + col];
    }
    // 直接加上偏置，不用额外 kernel
    output[row * N + col] = sum + bias[col];
}

注意： 不是所有 MatMul+Add 都能无脑融合。如果 Add 的另一个输入不是 bias 而是另一个变量（比如残差连接），那就不能简单融合，需要做「Add 下沉」或者「算子重排」。这个我们后面会讲到。

4.3 LayerNorm 融合

LayerNorm 的计算包含：计算均值 → 计算方差 → 归一化 → 缩放平移。这四个步骤如果分开写，要来回读写 4 次数据。

融合的思路是：一次遍历完成所有计算。也就是读一次输入，在寄存器里算完均值、方差、归一化、缩放，然后一次性写回结果。

# 未融合的 LayerNorm（4 次访存）
def layernorm_naive(x, gamma, beta):
    mean = x.mean(dim=-1, keepdim=True)      # 读一次
    var = ((x - mean) ** 2).mean(dim=-1)     # 读一次，写一次
    x_norm = (x - mean) / torch.sqrt(var + eps)  # 读两次，写一次
    return gamma * x_norm + beta             # 读一次，写一次

# 融合后的 LayerNorm（1 次访存）
def layernorm_fused(x, gamma, beta):
    # 实际实现中，用 CUDA kernel 一次搞定
    # 这里只是示意逻辑
    mean, var = compute_mean_var_once(x)     # 一次遍历算均值和方差
    x_norm = (x - mean) / torch.sqrt(var + eps)
    return gamma * x_norm + beta

性能数据（来自我的实测）：

实现方式	显存读写量	耗时（相对）
未融合（4 步）	4 × 输入大小	1.0x
部分融合（2 步）	2 × 输入大小	0.6x
完全融合（1 步）	1 × 输入大小	0.35x

4.4 融合后的内存访问优化

融合算子只是第一步。真正让性能起飞的是融合后的内存访问模式优化。我见过不少团队，算子融合做得很好，但性能提升不明显——问题就出在访存上。

核心原则就两条：

连续访问：尽量让每个线程访问连续的内存地址，利用缓存行
减少全局内存访问：能用寄存器就别用 shared memory，能用 shared memory 就别用 global memory

举个例子，融合后的 Conv+BN+ReLU，如果按 naive 方式写，每个线程还是独立计算，那缓存利用率很低。更好的做法是：

用 tiling 技术把输入分块
每个 block 加载一块数据到 shared memory
在 shared memory 里完成 Conv、BN、ReLU 的全部计算
一次性写回 global memory

避坑指南： 我曾经在融合 LayerNorm 时，为了省显存把中间结果全部放在寄存器里。结果寄存器溢出，性能反而下降了 30%。后来我改用 shared memory 做中间缓存，才把性能提上来。嗯，这里要注意：寄存器不是越多越好，够用就行。

知识体系总览

下面这张图总结了算子融合的核心逻辑和优化路径：

这张图把整个知识体系串起来了。左边是三种核心融合模式，中间是共同目标，右边是具体的优化策略。你可以把它当作一个检查清单——做算子融合时，对照着看看自己漏了哪一步。

总结一下：

Conv+BN+ReLU 融合：参数吸收，线性合并，最成熟也最有效
MatMul+Add 融合：偏置直接嵌入，注意区分 bias 和残差
LayerNorm 融合：一次遍历完成四步，减少 75% 的访存
融合后的访存优化：连续访问、tiling、合理使用 shared memory

我个人习惯在做融合优化时，先用 profiler 看看每个算子的访存占比。如果访存占比超过 60%，那融合大概率能带来明显收益。如果计算占比本来就高，融合效果可能没那么显著——这时候就该考虑算子替换或者算法层面的改进了。

想要更多资料，联系可可

📚 课程目录

4、算子融合（Fusion）：让计算更高效的核心技巧

4.1 Conv+BN+ReLU 融合

4.2 MatMul+Add 融合

4.3 LayerNorm 融合

4.4 融合后的内存访问优化

知识体系总览