算子融合：降低推理延迟的实战课程

2. 常见融合模式：Conv+BN+ReLU融合、Conv+Add+ReLU融合、LayerNorm融合、QKV融合

好，咱们直接进入正题。

算子融合这玩意儿，说白了就是把多个小操作合并成一个。为什么要这么干？因为每次调用一个算子，都有额外的开销——内存读写、kernel启动、数据搬运。你想想看，这些开销累加起来，推理速度能快才怪。

我个人习惯把融合模式分成四类。今天咱们一个一个过。

2.1 Conv+BN+ReLU 融合

这是最经典的融合模式，没有之一。我在项目中遇到过无数次，几乎每个CNN模型都能用上。

先说说原理。Conv是卷积，BN是批归一化，ReLU是激活函数。推理时，BN的计算可以合并到Conv的权重里。为什么？因为BN在推理时是线性变换。

BN的公式长这样：

y = gamma * (x - mean) / sqrt(var + eps) + beta

其中gamma、beta是学习到的参数，mean和var是训练时统计的均值方差。推理时这些值都是固定的。

我们把Conv的输出记作 W * x + b，代入BN公式：

y = gamma * (W * x + b - mean) / sqrt(var + eps) + beta
  = (gamma * W / sqrt(var + eps)) * x + (gamma * (b - mean) / sqrt(var + eps) + beta)

看到了吗？新的权重 W' = gamma * W / sqrt(var + eps)，新的偏置 b' = gamma * (b - mean) / sqrt(var + eps) + beta。

然后ReLU直接跟在后面，就是一个 max(0, x)。

融合后的效果：原来需要3次kernel调用，现在只需要1次。显存访问量减少，计算效率提升。

代码实现其实不复杂。我贴一段伪代码，你们感受一下：

def fuse_conv_bn_relu(conv_weight, conv_bias, bn_gamma, bn_beta, bn_mean, bn_var, bn_eps):
    # 计算融合后的权重
    scale = bn_gamma / np.sqrt(bn_var + bn_eps)
    fused_weight = conv_weight * scale.reshape(-1, 1, 1, 1)
    
    # 计算融合后的偏置
    fused_bias = bn_gamma * (conv_bias - bn_mean) / np.sqrt(bn_var + bn_eps) + bn_beta
    
    # 返回融合后的Conv+ReLU
    return fused_weight, fused_bias

小提示：实际部署时，很多框架（比如TensorRT、ONNX Runtime）会自动做这个融合。但如果你自己写推理引擎，这个优化一定要手动做。

2.2 Conv+Add+ReLU 融合

这个模式常见于残差网络（ResNet）。结构是：一个卷积的输出，加上另一个分支的输入，然后过ReLU。

嗯，这里要注意。融合的关键在于，Add操作本身很简单，但单独调用一次kernel也有开销。把它和Conv、ReLU合并，能省掉一次中间结果的读写。

我曾经在一个项目里，把ResNet50的所有Conv+Add+ReLU都做了融合。推理延迟直接降了12%。效果很明显。

实现上，需要把Add的输入作为Conv的额外输入。Conv计算完后，直接加上残差分支的值，再走ReLU。

// 伪代码示意
__global__ void fused_conv_add_relu_kernel(
    float* input, float* weight, float* bias,
    float* residual, float* output,
    int C, int H, int W) {
    
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    if (idx < C * H * W) {
        float conv_out = conv_compute(input, weight, bias, idx);
        float add_out = conv_out + residual[idx];
        output[idx] = max(0.0f, add_out);  // ReLU
    }
}

避坑指南：我曾经遇到过一个问题——残差分支的shape和Conv输出的shape不一致。比如stride=2时，空间尺寸减半。这时候Add之前需要先做下采样或padding。融合时一定要处理好这个对齐。

2.3 LayerNorm 融合

LayerNorm在Transformer里用得很多。它的计算包括：求均值、求方差、归一化、缩放+平移。

这四个步骤，如果分开做，需要4次kernel调用。但我们可以合并成1次。

我习惯这么写：

def fused_layer_norm(x, gamma, beta, eps=1e-5):
    # 一次kernel完成所有操作
    mean = x.mean(dim=-1, keepdim=True)
    var = x.var(dim=-1, keepdim=True, unbiased=False)
    x_norm = (x - mean) / torch.sqrt(var + eps)
    return gamma * x_norm + beta

等等，这看起来还是分开的？嗯，在PyTorch里确实是这样。但如果你用CUDA写，可以这样：

// 一个kernel里同时计算mean、var、归一化、缩放
__global__ void fused_layer_norm_kernel(...) {
    // 1. 加载数据到shared memory
    // 2. 计算mean
    // 3. 计算var
    // 4. 归一化
    // 5. 乘以gamma加beta
    // 6. 写回
}

这样做的好处是，数据只从global memory读一次，写一次。中间结果都留在寄存器或shared memory里。

性能提升：LayerNorm融合后，通常能快2-3倍。特别是当hidden size比较大时（比如768、1024），效果更明显。

2.4 QKV 融合

这个在Transformer的Attention里。Q、K、V三个矩阵，分别由输入乘以三个不同的权重矩阵得到。

传统做法是：

输入 x 乘以 W_q，得到 Q
输入 x 乘以 W_k，得到 K
输入 x 乘以 W_v，得到 V

三次矩阵乘法，三次kernel调用。

融合的做法是：把W_q、W_k、W_v拼成一个大的权重矩阵 W_qkv。一次矩阵乘法，得到 [Q, K, V] 拼接的结果，然后再拆开。

# 融合前
Q = x @ W_q
K = x @ W_k
V = x @ W_v

# 融合后
W_qkv = torch.cat([W_q, W_k, W_v], dim=-1)  # 拼权重
QKV = x @ W_qkv  # 一次矩阵乘
Q, K, V = torch.split(QKV, [d_q, d_k, d_v], dim=-1)  # 拆开

你想想看，一次大矩阵乘法，比三次小矩阵乘法快多了。为什么？因为大矩阵乘法能更好地利用GPU的并行计算能力，而且减少了kernel启动次数。

个人经验：我在部署BERT模型时，QKV融合是必做的优化。配合LayerNorm融合，整个Transformer block的推理延迟能降30%以上。

知识体系总览

下面这张图，我把四种融合模式的核心逻辑画出来了。你们可以对照着看：

好了，这四种融合模式，是算子融合里最基础也最常用的。我个人建议，先从Conv+BN+ReLU开始练手，这个最直观，效果也最明显。然后逐步尝试其他模式。

记住一句话：能融合的算子，尽量别分开。这是降低推理延迟的黄金法则。

总结一下：

Conv+BN+ReLU：把BN参数合并到Conv权重里，一次调用搞定
Conv+Add+ReLU：把残差加法合并到Conv输出后，省掉中间读写
LayerNorm：一个kernel完成均值、方差、归一化、缩放
QKV：三个权重拼成一个，一次矩阵乘出三个结果

下一节，咱们聊聊更高级的融合模式，比如Multi-Head Attention内部的融合。不过那是后话了。先把这四种吃透，你的模型推理速度就能上一个台阶。

📚 课程目录

算子融合：降低推理延迟的实战课程

2. 常见融合模式：Conv+BN+ReLU融合、Conv+Add+ReLU融合、LayerNorm融合、QKV融合

2.1 Conv+BN+ReLU 融合

2.2 Conv+Add+ReLU 融合

2.3 LayerNorm 融合

2.4 QKV 融合

知识体系总览