CUDA 实现softmax算子

一、Softmax 引入

Softmax 的作用就是从一堆数据中选出每个数据对于总体的重要程度——数字大的占比大，数字小的占比小，并且所有占比加起来等于 1。也就是说 Softmax 是把一组任意实数压缩成了一个概率分布。

数学定义：对于输入向量 $X = (x_{0}, x_{1}, . . ., x_{N - 1})$ ，Softmax 输出为：

S o f t m a x (x_{i}) = \frac{e^{x_{i}}}{\sum_{j = 0}^{N - 1} e^{x_{j}}}

Softmax 的公式涉及指数运算 $e^{x_{i}}$ ，这在浮点数表示中极易溢出：

FP32：最大可表示值约为 $3.4 \times 10^{38}$ ，对应 $e^{88.7}$ 。当 $x_{i} > 88.7$ 时， $e^{x_{i}}$ 就变成 +inf
FP16：最大值约为 $6.55 \times 10^{4}$ ，对应 $e^{11.09}$ 。当 $x_{i} > 11$ 时即溢出
下溢同样危险：当 $x_{i}$ 很小（如 −100−100）时， $e^{x i} \approx 0$ ，分母可能下溢为 0，导致 NaN

在实际的 Attention 计算中， $Q K^{T}$ 的值可以轻松超过 88（尤其是在大 $d_{k}$ 且未做缩放的情况下），朴素实现必然崩溃。

不过可以通过对输入向量的每个元素减去同一个常数 c 的方式来避免溢出，同时还不会改变 Softmax 的输出：

\frac{e^{x_{i} - c}}{\sum_{j} e^{x_{j} - c}} = \frac{e^{x_{i}} \cdot e^{- c}}{\sum_{j} e^{x_{j}} \cdot e^{- c}} = \frac{e^{x_{i}}}{\sum_{j} e^{x_{j}}}

因为每个数减去一个最大值后，它的值必然是小于等于 0 的，而 $e^{x}$ 中 x 小于等于 0 是，值域就变成了 (0, 1]

二、Softmax 朴素实现

每个线程独立处理输入矩阵的一行，串行完成该行的 max、exp-sum、normalize 三步。行与行之间完全独立，天然并行。

__global__ void softmax_v0(float* input, float* output, int M, int N) {
    int row = blockIdx.x * blockDim.x + threadIdx.x;
    if (row >= M) return;

    float* x = input  + row * N;
    float* y = output + row * N;

    // Pass 1: 求最大值
    float max_val = -INFINITY;
    for (int i = 0; i < N; i++) {
        max_val = fmaxf(max_val, x[i]);
    }

    // Pass 2: 指数求和
    float sum = 0.0f;
    for (int i = 0; i < N; i++) {
        sum += expf(x[i] - max_val);
    }

    // Pass 3: 归一化
    float inv_sum = 1.0f / sum;
    for (int i = 0; i < N; i++) {
        y[i] = expf(x[i] - max_val) * inv_sum;
    }
}

性能缺陷：

并行度不足：每行只有 1 个线程，N=4096 个元素被串行处理。GPU 的数千个 CUDA Core 绝大部分在空闲
非合并访存：同一 Warp 内的 32 个线程处理 32 个不同行，它们在内循环中访问的地址相差 N×4 字节（一整行的跨度），造成严重的非合并访存
冗余计算：expf(x[i] - max_val) 在 Pass 2 和 Pass 3 中重复计算了两次

三、Safe Softmax 三遍扫描

上一个版本的核心问题是并行度不够——一行 N 个元素只用 1 个线程处理。解决方案是让一个 Block 内的多个线程协作处理同一行。每个线程负责一行中的一段数据，通过 Shared Memory 做并行规约求出 max 和 sum。

这等于把 Softmax 分解为两个 Reduce 操作（max-reduce + sum-reduce）加一个 element-wise 操作（normalize）。

// V1: 一个 Block 处理一行，Block 内并行规约
__global__ void softmax_v1(float* input, float* output, int M, int N) {
    extern __shared__ float smem[];  // 用于规约

    int row = blockIdx.x;
    int tid = threadIdx.x;

    float* x = input  + row * N;
    float* y = output + row * N;

    // Pass 1: 并行求最大值
    float max_val = -INFINITY;
    for (int i = tid; i < N; i += blockDim.x) {
        max_val = fmaxf(max_val, x[i]);
    }
    smem[tid] = max_val;
    __syncthreads();

    // Shared Memory 规约求全局最大值
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            smem[tid] = fmaxf(smem[tid], smem[tid + s]);
        }
        __syncthreads();
    }
    max_val = smem[0];
    __syncthreads();

    // Pass 2: 并行求指数和
    float sum = 0.0f;
    for (int i = tid; i < N; i += blockDim.x) {
        sum += expf(x[i] - max_val);
    }
    smem[tid] = sum;
    __syncthreads();

    // Shared Memory 规约求总和
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            smem[tid] += smem[tid + s];
        }
        __syncthreads();
    }
    sum = smem[0];
    __syncthreads();

    // Pass 3: 归一化写出
    float inv_sum = 1.0f / sum;
    for (int i = tid; i < N; i += blockDim.x) {
        y[i] = expf(x[i] - max_val) * inv_sum;
    }
}

提升来源：

并行度大幅提升：256 个线程协作处理一行，每线程只需处理 N/256=16 个元素
合并访存：同一 Block 内的相邻线程访问行内的相邻地址（stride=1），形成完美的合并访存

仍存在的问题：

Shared Memory 规约需要多轮 __syncthreads()
三遍扫描意味着输入数据被读取了 3 次

四、Warp Shuffle 规约优化

上个版本的规约完全依赖 Shared Memory，存在两个问题：一是数据需要”寄存器 → SMEM → 寄存器”往返，路径偏长；二是每一轮规约都需要 __syncthreads()，在线程数已经较少的后期阶段同步开销占比偏高。

从 Reduce 优化的经验来看，Warp Shuffle（__shfl_down_sync）可以在寄存器间直接交换数据，延迟更低、无 Bank Conflict、Warp 内天然同步。我们将规约改为”Warp 内 Shuffle + Warp 间 Shared Memory”的两级结构：先在每个 Warp 内用 Shuffle 规约出局部结果，再把每个 Warp 的结果写到一个 32 元素的 SMEM 数组，由第一个 Warp 做最终规约。

// Warp 内规约：Max
__device__ float warpReduceMax(float val) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
    }
    return val;
}

// Warp 内规约：Sum
__device__ float warpReduceSum(float val) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

// 两级规约：Block 内 Max
__device__ float blockReduceMaxShuffle(float val) {
    __shared__ float warp_max[32];

    int lane = threadIdx.x % 32;
    int wid  = threadIdx.x / 32;

    val = warpReduceMax(val);

    if (lane == 0) warp_max[wid] = val;
    __syncthreads();

    int num_warps = blockDim.x / 32;
    val = (lane < num_warps) ? warp_max[lane] : -INFINITY;
    if (wid == 0) val = warpReduceMax(val);

    // 广播结果给所有线程
    __shared__ float block_result;
    if (threadIdx.x == 0) block_result = val;
    __syncthreads();
    return block_result;
}

// 两级规约：Block 内 Sum
__device__ float blockReduceSumShuffle(float val) {
    __shared__ float warp_sum[32];

    int lane = threadIdx.x % 32;
    int wid  = threadIdx.x / 32;

    val = warpReduceSum(val);

    if (lane == 0) warp_sum[wid] = val;
    __syncthreads();

    int num_warps = blockDim.x / 32;
    val = (lane < num_warps) ? warp_sum[lane] : 0.0f;
    if (wid == 0) val = warpReduceSum(val);

    __shared__ float block_result;
    if (threadIdx.x == 0) block_result = val;
    __syncthreads();
    return block_result;
}

// V2: Warp Shuffle 两级规约
__global__ void softmax_v2(float* input, float* output, int M, int N) {
    int row = blockIdx.x;
    int tid = threadIdx.x;

    float* x = input  + row * N;
    float* y = output + row * N;

    // Pass 1: 求行最大值
    float local_max = -INFINITY;
    for (int i = tid; i < N; i += blockDim.x) {
        local_max = fmaxf(local_max, x[i]);
    }
    float max_val = blockReduceMaxShuffle(local_max);

    // Pass 2: 求指数和
    float local_sum = 0.0f;
    for (int i = tid; i < N; i += blockDim.x) {
        local_sum += expf(x[i] - max_val);
    }
    float sum = blockReduceSumShuffle(local_sum);

    // Pass 3: 归一化
    float inv_sum = 1.0f / sum;
    for (int i = tid; i < N; i += blockDim.x) {
        y[i] = expf(x[i] - max_val) * inv_sum;
    }
}

Warp Shuffle 规约的数据路径是寄存器到寄存器，相比 Shared Memory 规约，省去了之间到 SMEM 的路径，因此速度快了很多，并且不需要 __syncthreads()，因为 warp 内是天然同步的。

五、向量化加载 float4

V2 的瓶颈已经转移到全局内存加载阶段。每次循环中，每个线程只加载 1 个 float（4 字节），指令调度开销相对于数据吞吐比例偏高。使用 float4 向量化加载，每条指令搬运 16 字节（4 个 float），可以：

减少加载指令总数（减少 4 倍的循环迭代）
提升内存事务的利用效率
增加指令级并行（ILP）

// V3: float4 向量化加载
__global__ void softmax_v3(float* input, float* output, int M, int N) {
    int row = blockIdx.x;
    int tid = threadIdx.x;

    float* x = input  + row * N;
    float* y = output + row * N;

    // 向量化指针
    float4* x4 = reinterpret_cast<float4*>(x);
    float4* y4 = reinterpret_cast<float4*>(y);
    int N4 = N / 4;  // float4 元素数

    // Pass 1: 向量化求最大值
    float local_max = -INFINITY;
    for (int i = tid; i < N4; i += blockDim.x) {
        float4 data = x4[i];
        local_max = fmaxf(local_max, fmaxf(fmaxf(data.x, data.y),
                                            fmaxf(data.z, data.w)));
    }
    // 处理尾部元素（N 不是 4 的倍数时）
    for (int i = N4 * 4 + tid; i < N; i += blockDim.x) {
        local_max = fmaxf(local_max, x[i]);
    }
    float max_val = blockReduceMaxShuffle(local_max);

    // Pass 2: 向量化求指数和
    float local_sum = 0.0f;
    for (int i = tid; i < N4; i += blockDim.x) {
        float4 data = x4[i];
        local_sum += expf(data.x - max_val) + expf(data.y - max_val)
                   + expf(data.z - max_val) + expf(data.w - max_val);
    }
    for (int i = N4 * 4 + tid; i < N; i += blockDim.x) {
        local_sum += expf(x[i] - max_val);
    }
    float sum = blockReduceSumShuffle(local_sum);

    // Pass 3: 向量化归一化
    float inv_sum = 1.0f / sum;
    for (int i = tid; i < N4; i += blockDim.x) {
        float4 data = x4[i];
        float4 result;
        result.x = expf(data.x - max_val) * inv_sum;
        result.y = expf(data.y - max_val) * inv_sum;
        result.z = expf(data.z - max_val) * inv_sum;
        result.w = expf(data.w - max_val) * inv_sum;
        y4[i] = result;
    }
    for (int i = N4 * 4 + tid; i < N; i += blockDim.x) {
        y[i] = expf(x[i] - max_val) * inv_sum;
    }
}

每个线程的循环次数从 N/blockDim 降到 N/(4×blockDim)，减少了 4 倍的循环迭代开销和指令调度次数。同时 float4 加载能更好地隐藏内存延迟。

⚠️ 注意：使用 float4 要求输入数组 16 字节对齐。在实际工程中，通常在内存分配时使用 cudaMalloc（保证 256 字节对齐），并确保 N 是 4 的倍数（不满足时用尾部处理兜底）。

六、两遍融合 Kernel

回顾前面的计算流程：三遍扫描读取输入 3 次。对于 (4096,4096) 的矩阵，总读取量为 3×4096×4096×4=1923×4096×4096×4=192 MB。如果能把 Pass 2（exp-sum）和 Pass 3（normalize）合并为一遍，就能减少到 2 次读取（128 MB），读取量减少 1/3，理论加速比为 1.5 倍。

合并的关键观察：Pass 3 计算 $e^{x_{i} - m} / d$ 时需要用到 d（指数和），而 d 在 Pass 2 结束后才确定。如果我们不存储 Pass 2 的中间结果，只保存 m 和 d，然后在第二遍同时计算 $e^{x_{i} - m}$ 和除以 d，就实现了融合。

事实上，纯粹的两遍方案需要 Online Softmax 算法（下篇文章的主题）。但我们可以做一个折中——将 Pass 2 和 Pass 3 部分融合：在第 2 遍中，每个线程将自己负责的 $e^{x_{i} - m}$ 暂存到 Shared Memory，然后求出全局 d 后立即从 SMEM 读取并写出归一化结果，避免第 3 次全局内存读取。

// V4: 两遍 Kernel（Pass 1 求 max，Pass 2 融合 exp-sum + normalize）
// 需要额外 Shared Memory 暂存 exp 值
__global__ void softmax_v4(float* input, float* output, int M, int N) {
    extern __shared__ float smem[];  // 大小 = N（暂存 exp 值）

    int row = blockIdx.x;
    int tid = threadIdx.x;

    float* x = input  + row * N;
    float* y = output + row * N;

    float4* x4 = reinterpret_cast<float4*>(x);
    int N4 = N / 4;

    // Pass 1: 求行最大值（向量化 + Warp Shuffle）
    float local_max = -INFINITY;
    for (int i = tid; i < N4; i += blockDim.x) {
        float4 data = x4[i];
        local_max = fmaxf(local_max, fmaxf(fmaxf(data.x, data.y),
                                            fmaxf(data.z, data.w)));
    }
    for (int i = N4 * 4 + tid; i < N; i += blockDim.x) {
        local_max = fmaxf(local_max, x[i]);
    }
    float max_val = blockReduceMaxShuffle(local_max);

    // Pass 2: 计算 exp 并暂存，同时累加求 sum
    float local_sum = 0.0f;
    for (int i = tid; i < N; i += blockDim.x) {
        float exp_val = expf(x[i] - max_val);
        smem[i] = exp_val;  // 暂存到 Shared Memory
        local_sum += exp_val;
    }
    float sum = blockReduceSumShuffle(local_sum);

    // 直接从 Shared Memory 读取 exp 值并归一化写出
    float inv_sum = 1.0f / sum;
    for (int i = tid; i < N; i += blockDim.x) {
        y[i] = smem[i] * inv_sum;
    }
}

注意：V4 需要 N×4 字节的 Shared Memory 来暂存 exp 值。当 N=4096 时需要 16KB，在 A100（每 SM 164KB Shared Memory）上完全够用。但当 N 很大（如 32768 以上）时，Shared Memory 装不下整行数据，此方案就不适用了。

C++ > CUDA

#C++ #CUDA

CUDA 实现softmax算子

http://example.com/2026/07/24/CUDA-实现softmax算子/

作者

Yu xin

发布于

2026年7月24日

许可协议

CUDA 实现online softmax算子上一篇

CUDA Streams 下一篇