4、算子融合（Kernel Fusion）：将多个小算子合并为一个，减少Kernel Launch开销

算子融合，说白了就是把几个小算子捏成一个。

为什么要这么做？因为每次启动一个Kernel，GPU都要做一堆准备工作——分配寄存器、加载参数、同步线程。这些开销虽然单个看起来不大，但架不住次数多啊。我见过一个项目，光Kernel Launch就占了总时间的30%以上，你说亏不亏？

4.1 为什么需要算子融合？

先看一个最简单的例子。假设你要做：

y = x + 1
z = y * 2

如果不融合，GPU要启动两个Kernel：

Kernel 1：计算加法，写回显存
Kernel 2：从显存读y，计算乘法，写回z

这里有个巨大的浪费——y在显存里转了一圈，其实完全没必要。融合之后，一个Kernel搞定：

z = (x + 1) * 2

嗯，就这么简单。但实际项目中，融合的收益远不止减少一次读写。

核心收益有三点：

减少Kernel Launch次数——每次启动都有几十微秒的固定开销
减少显存带宽占用——中间结果不用写回显存
提高计算密度——每个线程做更多计算，利用率更高

4.2 常见的融合模式

我在项目中总结了几种最常见的融合模式，你八成也会遇到：

融合模式	典型场景	收益
Element-wise + Element-wise	ReLU + Add, Mul + Add	减少带宽，减少Launch
Reduce + Element-wise	Softmax, LayerNorm	减少中间结果读写
Conv + Activation	Conv2D + ReLU	减少显存访问
MatMul + Bias + Activation	全连接层	减少三次Launch为一次

我个人习惯把融合分成两类：水平融合和垂直融合。

水平融合：把多个独立的同类型算子合并，比如两个Element-wise算子
垂直融合：把前后依赖的算子合并，比如Conv + ReLU

4.3 实战：手写一个融合Kernel

拿最经典的LayerNorm来举例。LayerNorm的计算流程是：

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

如果不融合，需要5个Kernel：计算mean、计算var、减法、除法、乘加。每个Kernel都要读一遍x，写一遍中间结果。

融合之后，一个Kernel搞定：

__global__ void fused_layernorm_kernel(
    float* x, float* y, float* gamma, float* beta,
    int rows, int cols, float eps) {
    
    // 每个线程处理一行
    int row = blockIdx.x;
    
    // 第一步：计算mean
    float sum = 0.0f;
    for (int i = threadIdx.x; i < cols; i += blockDim.x) {
        sum += x[row * cols + i];
    }
    // 规约求和
    __syncthreads();
    float mean = sum / cols;
    
    // 第二步：计算var
    float var_sum = 0.0f;
    for (int i = threadIdx.x; i < cols; i += blockDim.x) {
        float diff = x[row * cols + i] - mean;
        var_sum += diff * diff;
    }
    // 规约求方差
    __syncthreads();
    float var = var_sum / cols;
    float inv_std = rsqrtf(var + eps);
    
    // 第三步：归一化 + 缩放平移
    for (int i = threadIdx.x; i < cols; i += blockDim.x) {
        float normalized = (x[row * cols + i] - mean) * inv_std;
        y[row * cols + i] = normalized * gamma[i] + beta[i];
    }
}

我的经验：写融合Kernel时，最容易被忽略的是同步问题。上面代码里用了两次__syncthreads()，一次在计算mean之后，一次在计算var之后。我曾经漏掉第二次同步，结果方差算出来全是错的，排查了整整一个下午。

4.4 融合的边界条件

是不是所有算子都能融合？当然不是。有几种情况我建议你别硬融：

计算模式差异太大——比如一个算子需要大量共享内存，另一个不需要，硬融反而降低性能
寄存器压力过大——融合后寄存器不够用，编译器会 spill 到 local memory，得不偿失
分支逻辑复杂——不同线程走不同分支，融合后 warp divergence 严重

避坑指南：我曾经把一个Conv和一个复杂的自定义激活函数融合在一起，结果因为激活函数里有很多if-else分支，导致warp内线程严重分化，性能反而下降了40%。后来我改成只融合简单的ReLU和GELU，效果就好多了。

4.5 融合的自动化工具

手动写融合Kernel毕竟费劲。现在主流框架都提供了自动融合的能力：

PyTorch JIT：用torch.jit.script可以自动融合一些简单算子
TensorRT：在推理时自动做图优化和算子融合
XLA：Google的编译器，能把多个算子编译成一个高效的融合Kernel
Triton：我最近比较喜欢用的，写融合Kernel比CUDA方便很多

但说实话，自动融合的效果往往不如手写的好。尤其是遇到复杂的融合模式，编译器不敢做太激进的优化。我的建议是：先用自动工具跑一遍，然后profile看看哪些地方没融合好，再手动优化。

4.6 融合的性能评估

怎么判断融合有没有效果？别只看端到端时间，要看这几个指标：

Kernel Launch次数——融合后应该明显减少
显存带宽利用率——融合后应该更高
计算吞吐——融合后应该提升
寄存器使用量——这个要小心，别超了

我记得有一次优化一个Transformer模型，把LayerNorm、Add和Residual融合在一起，Kernel Launch从12次降到了4次，端到端速度提升了1.8倍。嗯，这就是融合的魅力。

上面这张图很直观地展示了融合前后的区别。左边要启动3次Kernel，中间结果在显存里倒腾；右边一次搞定，数据全在寄存器里流转。你说哪个快？

一个小技巧：如果你不确定某个融合是否有效，可以用NVIDIA Nsight Compute跑一下。重点关注L1/L2 cache命中率和显存带宽利用率。融合后这两个指标通常会有明显改善。

算子融合这件事，说白了就是用空间换时间——多写几行代码，换来几倍的性能提升。这笔账，怎么算都划算。

📚 课程目录

4、算子融合（Kernel Fusion）：将多个小算子合并为一个，减少Kernel Launch开销

4.1 为什么需要算子融合？

4.2 常见的融合模式

4.3 实战：手写一个融合Kernel

4.4 融合的边界条件

4.5 融合的自动化工具

4.6 融合的性能评估