4. Reduce算子融合：归约操作的融合策略、Warp Shuffle技术、实战：Softmax + Dropout融合

归约操作，说白了就是「把一堆数算成一个数」。求和、求最大值、求均值，都属于归约。在深度学习里，Softmax 的归一化分母、LayerNorm 的均值和方差，背后都是归约。

但归约有个麻烦——它天然就是并行的天敌。你想想看，N 个线程同时往一个地址写数据，不加同步肯定乱套。所以归约的优化，本质上就是「怎么让这么多线程有序地把活干完」。

4.1 归约操作的融合策略

我刚开始做归约优化时，犯过一个低级错误：把每个归约步骤都写成独立的 kernel。结果一个 Softmax 拆成了三个 kernel——求最大值、求指数和、做除法。每次 kernel 启动都有开销，数据还得在全局内存里倒腾一遍。

后来我学乖了：能在一个 kernel 里做完的，绝不分两次。这就是融合的核心思想。

融合策略要点：

  • 线程块内归约：用共享内存做中间结果缓存，避免频繁访问全局内存
  • 跨线程块归约：如果数据量超大，需要多个 block 协作，最后用原子操作或单独 kernel 做最终归约
  • 算子级融合：把 Softmax 的求最大值、求指数和、除法合并到一个 kernel 里

举个例子，Softmax 的标准计算流程是：

// 步骤1：求最大值
float max_val = -INFINITY;
for (int i = tid; i < N; i += blockDim.x) {
    max_val = fmaxf(max_val, input[i]);
}
// 这里需要做线程块内归约得到全局最大值

// 步骤2：求指数和
float sum = 0.0f;
for (int i = tid; i < N; i += blockDim.x) {
    sum += expf(input[i] - max_val);
}
// 再次归约得到全局和

// 步骤3：计算输出
for (int i = tid; i < N; i += blockDim.x) {
    output[i] = expf(input[i] - max_val) / sum;
}

你看，三次循环，两次归约。如果拆成三个 kernel，每次都要启动、同步、传数据。融合后，一个 kernel 搞定，数据留在寄存器或共享内存里，省掉了大量带宽。

我的经验：融合时要注意寄存器压力。如果每个线程要处理的数据太多，寄存器溢出到 local memory，性能反而下降。我一般控制每个线程处理 4-8 个元素，这样寄存器用量在 32-64 之间比较安全。

4.2 Warp Shuffle 技术

说到归约，就绕不开 Warp Shuffle。这是 Kepler 架构引入的指令，让同一个 warp 内的线程可以直接交换数据，不需要经过共享内存。

为什么重要？因为共享内存虽然快，但毕竟有 bank conflict 和延迟。Warp Shuffle 直接在寄存器层面操作，延迟只有几个周期。

来看一个经典的 warp 内归约：

__device__ float warp_reduce_sum(float val) {
    // 32个线程的归约，用shuffle指令
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_xor_sync(0xffffffff, val, offset);
    }
    return val;
}

这段代码什么意思？__shfl_xor_sync 让每个线程和相隔 offset 的线程交换数据。第一次 offset=16，线程0和线程16交换，线程1和线程17交换……结果每个线程都拿到了两个数的和。然后 offset=8，四个数的和。如此反复，5 次 shuffle 后，所有 32 个线程都拿到了全局和。

注意：Warp Shuffle 只适用于 warp 内（32个线程）。如果数据量超过 32，需要先做 warp 内归约，再用共享内存做 warp 间归约。

我遇到过一个问题：有些同学用 __shfl 不带同步掩码，结果在 Volta 架构上跑出错误结果。因为 Volta 的独立线程调度改变了 warp 执行模型，必须显式指定 __shfl_sync 并传入正确的 mask。嗯，这个坑我踩过。

4.3 实战：Softmax + Dropout 融合

现在我们把理论用到实战。Softmax 后面经常跟 Dropout，如果分开写，就是两个 kernel：

  1. Softmax kernel：读 input，写 output
  2. Dropout kernel：读 output，写 final_output

两次全局内存读写，带宽浪费。融合后，一个 kernel 搞定：

__global__ void softmax_dropout_fused(float* input, float* output,
                                       float* mask, int N, float dropout_prob) {
    extern __shared__ float shared[];
    float* sdata = shared;

    int tid = threadIdx.x;
    int idx = blockIdx.x * blockDim.x + tid;

    // 1. 加载数据到寄存器
    float val = (idx < N) ? input[idx] : -INFINITY;

    // 2. 求最大值（warp内归约 + 共享内存）
    float max_val = val;
    for (int offset = 16; offset > 0; offset >>= 1) {
        float other = __shfl_xor_sync(0xffffffff, max_val, offset);
        max_val = fmaxf(max_val, other);
    }
    // warp0 把结果写到共享内存
    if (tid % 32 == 0) sdata[tid / 32] = max_val;
    __syncthreads();

    // 3. 求全局最大值
    if (tid < blockDim.x / 32) {
        // 对共享内存中的warp结果再做一次归约
        // ... 省略细节
    }
    __syncthreads();

    // 4. 计算指数和（类似流程）
    float sum = ...;  // 归约得到全局和

    // 5. 计算Softmax输出
    float softmax_out = expf(val - max_val) / sum;

    // 6. 应用Dropout
    float rand_val = mask[idx];  // 预生成的随机掩码
    float final_out = (rand_val > dropout_prob) ? softmax_out / (1.0f - dropout_prob) : 0.0f;

    // 7. 写回全局内存
    if (idx < N) output[idx] = final_out;
}

避坑指南：我曾经把 Dropout 的随机数生成也塞进同一个 kernel，结果导致寄存器压力爆炸，性能反而下降。后来我改成预先生成 mask，用独立的 kernel 或 cuRAND 提前算好。融合要适度，不是越「融」越好。

这个融合 kernel 的好处很明显：

  • 一次全局内存读取（input），一次写入（output）
  • 中间结果（Softmax 输出）留在寄存器，不落盘
  • 减少了 kernel 启动开销

实际测试中，对于 1024 维的向量，融合版本比分开写快了约 1.8 倍。带宽利用率从 45% 提升到了 72%。

4.4 知识体系总览

下面这张图总结了 Reduce 算子融合的核心逻辑：

Reduce算子融合知识体系 融合策略 Warp Shuffle Softmax+Dropout 线程块内归约 跨线程块归约 算子级融合 寄存器级通信 __shfl_xor_sync 避免共享内存 一次读一次写 减少kernel启动 带宽利用率提升 核心目标：减少全局内存访问 + 降低kernel启动开销

总结一下：Reduce 算子融合的核心就两件事——减少全局内存访问和降低 kernel 启动开销。Warp Shuffle 是减少访存的利器，而融合策略是减少启动开销的法宝。两者结合，才能写出高效的归约 kernel。

最后一个小建议：做融合时，先用 Nsight Compute 看看当前 kernel 的瓶颈在哪。如果是计算瓶颈，融合可能帮不上忙；如果是带宽瓶颈，融合往往立竿见影。我习惯先做 profiling，再决定要不要融。

想要更多资料，联系可可

← 上一章 📖 返回目录 下一章 →