4. 算子融合与Kernel合并：Element-wise算子融合、Reduce与Scan的融合策略、Fused Kernel设计模式、减少Kernel Launch开销

说到算子融合，我脑子里第一个蹦出来的场景是几年前调一个Transformer模型。那时候显存动不动就爆，Kernel Launch次数多到让人头皮发麻。后来我花了一周时间，把十几个Element-wise算子揉成一个，显存直接降了40%。嗯，今天我们就聊聊这个。

4.1 为什么需要算子融合？

你想想看，GPU执行一个Kernel，光启动开销就要几十微秒。如果你的模型里有几百个小的Element-wise算子，每个都单独Launch一次，那光启动时间就够你喝一壶的。更别说每个Kernel还要读写显存，带宽就这么被白白浪费了。

算子融合的核心思想很简单：把多个小算子合并成一个大Kernel。好处有两个：

• 减少Kernel Launch次数：一次启动，干完所有活
• 减少显存带宽消耗：中间结果不用写回显存，直接在寄存器或共享内存里传递

4.2 Element-wise算子融合

Element-wise算子是最容易融合的。说白了，就是每个元素独立计算，不依赖其他元素。比如：

// 不融合：三个Kernel
y = relu(x)    // Kernel 1
z = y + 1      // Kernel 2
w = z * 0.5    // Kernel 3

// 融合后：一个Kernel
w = (relu(x) + 1) * 0.5  // 一个Kernel搞定

我在项目中遇到过最典型的场景是LayerNorm的前向计算。LayerNorm里有mean、variance、normalize、scale、bias五个步骤，如果每个都单独写Kernel，显存占用和启动开销都很大。我当时的做法是把它们全部融合到一个Kernel里：

__global__ void fused_layernorm(float* input, float* output, 
                                 float* gamma, float* beta, int N) {
    // 一个Kernel里完成所有操作
    float sum = 0.0f;
    for (int i = 0; i < N; i++) {
        sum += input[i];
    }
    float mean = sum / N;
    
    float var = 0.0f;
    for (int i = 0; i < N; i++) {
        var += (input[i] - mean) * (input[i] - mean);
    }
    float std = sqrt(var / N + 1e-5f);
    
    for (int i = 0; i < N; i++) {
        output[i] = gamma[i] * (input[i] - mean) / std + beta[i];
    }
}

4.3 Reduce与Scan的融合策略

Reduce和Scan这类算子比较特殊，它们有数据依赖关系。Reduce需要全局规约，Scan需要前缀和计算。融合起来要小心。

Reduce融合：最常见的是把多个Reduce合并成一个。比如你要同时计算mean和variance，可以一次遍历完成：

// 不融合：两次遍历
float mean = reduce_sum(x) / N;        // Kernel 1
float var = reduce_sum((x-mean)^2) / N; // Kernel 2

// 融合：一次遍历
float sum = 0.0f, sum_sq = 0.0f;
for (int i = 0; i < N; i++) {
    sum += x[i];
    sum_sq += x[i] * x[i];
}
float mean = sum / N;
float var = sum_sq / N - mean * mean;

Scan融合：这个稍微复杂点。我曾经在实现Softmax时，把max-reduce、sum-reduce和element-wise操作融合在一起。Softmax需要先找最大值（reduce），然后计算exp和sum（reduce），最后做除法（element-wise）。这三个步骤完全可以融合：

__global__ void fused_softmax(float* input, float* output, int N) {
    // 第一步：找最大值
    float max_val = -FLT_MAX;
    for (int i = 0; i < N; i++) {
        max_val = fmaxf(max_val, input[i]);
    }
    
    // 第二步：计算exp和sum
    float sum = 0.0f;
    for (int i = 0; i < N; i++) {
        sum += expf(input[i] - max_val);
    }
    
    // 第三步：归一化
    for (int i = 0; i < N; i++) {
        output[i] = expf(input[i] - max_val) / sum;
    }
}

4.4 Fused Kernel设计模式

做多了融合，我总结出几种常见的设计模式：

模式名称	适用场景	典型例子
链式融合	连续的Element-wise操作	ReLU + Add + Mul
Reduce-Elementwise融合	Reduce后接Element-wise	Softmax、LayerNorm
Scan-Elementwise融合	Scan后接Element-wise	Cumsum + Scale
多输出融合	一个输入产生多个输出	同时计算mean和var

我个人比较常用的是链式融合和Reduce-Elementwise融合。前者实现简单，效果立竿见影；后者虽然复杂点，但收益更大。

4.5 减少Kernel Launch开销

Kernel Launch开销到底有多大？我测过，在A100上，一个空Kernel的Launch时间大约是5-10微秒。如果你的模型有1000个Kernel，光启动时间就5-10毫秒。对于推理场景，这已经是个不小的开销了。

减少Kernel Launch开销的方法：

• 算子融合：把多个小Kernel合并成一个大Kernel
• 使用CUDA Graph：把多个Kernel Launch打包成一个Graph，一次提交多次执行
• 异步Launch：利用CUDA Stream并行Launch多个Kernel
• 持久化Kernel：让Kernel一直运行在GPU上，通过队列接收任务

4.6 本章小结

算子融合说白了就是用时间换空间，再用空间换时间。融合减少了显存访问和Kernel Launch，但可能增加寄存器和共享内存的使用。做融合时要权衡利弊，不能盲目。

我个人的经验是：先做Element-wise融合，再做Reduce融合，最后考虑Scan融合。每一步都要实测验证，不要想当然。毕竟，GPU优化这件事，数据说了算。

想要更多资料，联系可可