CUDA 实现online softmax算子

一、Online Softmax 引入

传统 Safe Softmax 需要三遍扫描输入数据，而 Online Softmax 通过巧妙的数学递推，在一遍扫描中同时完成 max 和 sum 的计算，将内存读取次数从 3 次降到 2 次。

Online Softmax 的关键洞察是：max 和 sum 可以在一遍扫描中同时维护，代价只是在发现新的最大值时”修正”之前累积的 sum。

这就是 Online Softmax 的精髓——用乘法修正代替重新计算。

数学推导：

设已经处理了前 k 个元素，当前维护的状态为：

$m_{k} = m a x (x_{0}, x_{1}, . . ., x_{k - 1})$ （前 k 个元素的最大值）
$d_{k} = \sum_{j = 0}^{k - 1} e^{x_{j} - m_{k}}$ （基于当前 max 的指数和）

当处理第 k+1 个元素 $x_{k}$ 时：

更新最大值： $m_{k + 1} = m a x (m_{k}, x_{k})$
修正已有的 sum 并加入新元素： $d_{k + 1} = d_{k} \cdot e^{m_{k} - m_{k + 1}} + e^{x_{k} - m_{k + 1}}$

推导过程：
$d_{k + 1} = \sum_{j = 0}^{k} e^{x_{j} - m_{k + 1}} = e^{x_{k} - m_{k + 1}} + \sum_{j = 0}^{k - 1} e^{x_{j} - m_{k + 1}} = e^{x_{k} - m_{k + 1}} + \sum_{j = 0}^{k - 1} e^{x_{j} - m_{k}} \cdot e^{x_{j} - m_{k + 1}} = d_{k} \cdot e^{m_{k} - m_{k + 1}} + e^{x_{k} - m_{k + 1}}$

关键在于 $e^{m_{k} - m_{k + 1}}$ 这个修正因子。如果新元素不是新的最大值（ $m_{k + 1} = m_{k}$ ），则 $e^{m_{k} - m_{k + 1}} = e^{0} = 1$ ，修正退化为简单的累加——没有任何额外开销。

二、单线程 Online Softmax

最直接的实现：每个线程处理一行，串行地维护 (m,d) 状态，最后再做一遍归一化。

// V0: 单线程 Online Softmax（一个线程处理一行）
__global__ void online_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;

    // 第 1 遍：Online 同时计算 max 和 sum
    float m = -INFINITY;
    float d = 0.0f;

    for (int i = 0; i < N; i++) {
        float xi = x[i];
        float m_new = fmaxf(m, xi);
        d = d * expf(m - m_new) + expf(xi - m_new);
        m = m_new;
    }

    // 第 2 遍：归一化
    float inv_d = 1.0f / d;
    for (int i = 0; i < N; i++) {
        y[i] = expf(x[i] - m) * inv_d;
    }
}

和 Safe Softmax V0 类似，单线程处理整行存在严重的并行度不足和非合并访存问题。但相比三遍扫描的 V0，这里只读了 2 遍输入数据。

注意：V0 的 Online 递推中每步都需要 expf 调用（用于修正因子），这增加了计算负担。但由于 Softmax 本身是 Memory-Bound 的，这些额外计算在高性能版本中会被内存访问延迟完全隐藏。

三、Block 级并行 + Warp Shuffle 合并规约

Online Softmax 的递推是串行的——每一步的 (m,d)(m,d) 依赖上一步。如何并行化？

关键观察：可以先让每个线程独立处理一部分元素，各自维护局部的 (mlocal,dlocal)，最后将这些局部结果合并。

两个局部结果 (m1,d1) 和 (m2,d2) 的合并公式：

m_{m e r g e d} = m a x (m_{1}, m_{2}) ， d_{m e r g e d} = d_{1} \cdot e^{m_{1} - m_{m e r g e d}} + d_{2} \cdot e^{m_{2} - m_{m e r g e d}}

// Warp 内 Online 合并规约
__device__ void warpReduceOnline(float& m, float& d) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        float m2 = __shfl_down_sync(0xffffffff, m, offset);
        float d2 = __shfl_down_sync(0xffffffff, d, offset);

        float m_new = fmaxf(m, m2);
        d = d * expf(m - m_new) + d2 * expf(m2 - m_new);
        m = m_new;
    }
}

// V1: Block 级并行 + Warp Shuffle 两级合并规约
__global__ void online_softmax_v1(float* input, float* output, int M, int N) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid % 32;
    int wid  = tid / 32;

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

    // 第 1 遍：每线程 Online 处理一段
    float local_m = -INFINITY;
    float local_d = 0.0f;

    for (int i = tid; i < N; i += blockDim.x) {
        float xi = x[i];
        float m_new = fmaxf(local_m, xi);
        local_d = local_d * expf(local_m - m_new) + expf(xi - m_new);
        local_m = m_new;
    }

    // 第一级：Warp 内合并
    warpReduceOnline(local_m, local_d);

    // Warp 间通过 Shared Memory 交换
    __shared__ float warp_m[32];
    __shared__ float warp_d[32];

    if (lane == 0) {
        warp_m[wid] = local_m;
        warp_d[wid] = local_d;
    }
    __syncthreads();

    // 第二级：Warp 0 做最终合并
    int num_warps = blockDim.x / 32;
    if (wid == 0) {
        local_m = (lane < num_warps) ? warp_m[lane] : -INFINITY;
        local_d = (lane < num_warps) ? warp_d[lane] : 0.0f;
        warpReduceOnline(local_m, local_d);
    }

    // 广播最终结果
    __shared__ float final_m, final_d;
    if (tid == 0) {
        final_m = local_m;
        final_d = local_d;
    }
    __syncthreads();

    float row_max = final_m;
    float inv_sum = 1.0f / final_d;

    // 第 2 遍：归一化输出
    for (int i = tid; i < N; i += blockDim.x) {
        y[i] = expf(x[i] - row_max) * inv_sum;
    }
}

合并顺序会影响浮点精度（加法不严格满足结合律），但对于 Softmax 这种对微小误差不敏感的场景，树形规约的精度完全足够。

Online 版本的规约操作比普通 sum-reduce 复杂（每步需要 2 次 expf），但节省的内存带宽远大于多出的计算量。对于 Memory-Bound 操作，减少一次全局读取的收益是决定性的。

四、向量化加载 float4

和 Safe Softmax 的优化路径一致，使用 float4 向量化加载减少指令数。Online 递推阶段，每次加载 4 个元素，逐个更新 (m,d) 状态。

// V2: float4 向量化加载 + Online Softmax
__global__ void online_softmax_v2(float* input, float* output, int M, int N) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid % 32;
    int wid  = tid / 32;

    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;

    // 第 1 遍：向量化 Online 递推
    float local_m = -INFINITY;
    float local_d = 0.0f;

    for (int i = tid; i < N4; i += blockDim.x) {
        float4 data = x4[i];

        // 对 4 个元素依次更新 (m, d)
        float vals[4] = {data.x, data.y, data.z, data.w};
        for (int k = 0; k < 4; k++) {
            float m_new = fmaxf(local_m, vals[k]);
            local_d = local_d * expf(local_m - m_new) + expf(vals[k] - m_new);
            local_m = m_new;
        }
    }
    // 尾部处理
    for (int i = N4 * 4 + tid; i < N; i += blockDim.x) {
        float xi = x[i];
        float m_new = fmaxf(local_m, xi);
        local_d = local_d * expf(local_m - m_new) + expf(xi - m_new);
        local_m = m_new;
    }

    // 两级 Warp Shuffle 合并（与 V1 相同）
    warpReduceOnline(local_m, local_d);

    __shared__ float warp_m[32];
    __shared__ float warp_d[32];
    if (lane == 0) {
        warp_m[wid] = local_m;
        warp_d[wid] = local_d;
    }
    __syncthreads();

    int num_warps = blockDim.x / 32;
    if (wid == 0) {
        local_m = (lane < num_warps) ? warp_m[lane] : -INFINITY;
        local_d = (lane < num_warps) ? warp_d[lane] : 0.0f;
        warpReduceOnline(local_m, local_d);
    }

    __shared__ float final_m, final_d;
    if (tid == 0) {
        final_m = local_m;
        final_d = local_d;
    }
    __syncthreads();

    float row_max = final_m;
    float inv_sum = 1.0f / final_d;

    // 第 2 遍：向量化归一化
    for (int i = tid; i < N4; i += blockDim.x) {
        float4 data = x4[i];
        float4 result;
        result.x = expf(data.x - row_max) * inv_sum;
        result.y = expf(data.y - row_max) * inv_sum;
        result.z = expf(data.z - row_max) * inv_sum;
        result.w = expf(data.w - row_max) * inv_sum;
        y4[i] = result;
    }
    for (int i = N4 * 4 + tid; i < N; i += blockDim.x) {
        y[i] = expf(x[i] - row_max) * inv_sum;
    }
}

虽然 float4 一次加载 4 个元素，但 Online 递推仍然逐元素更新 (m,d)。这看似没有收益，但实际上：

内存层面：float4 加载用更少的指令传输相同的数据量，减少了指令调度的瓶颈
计算层面：4 个 expf 调用可以流水线化，编译器能更好地排布指令
整体效果：循环迭代减少 4 倍，循环开销（分支判断、地址计算）大幅降低

五、寄存器缓存消除第二遍读取

如果每行元素不多（N/blockDim≤N/blockDim≤ 几十），每个线程需要处理的元素数量有限，可以将第 1 遍读入的数据暂存在寄存器中，第 2 遍直接从寄存器取值做归一化，避免第二次全局内存读取。

这就是真正的一遍（One-Pass）Softmax——全局内存只读一次、写一次。

// V3: 寄存器缓存实现 One-Pass Softmax
// 适用于 N / blockDim.x 较小的场景（每线程处理元素数不超过 MAX_ELEMS）
#define MAX_ELEMS_PER_THREAD 32

__global__ void online_softmax_v3(float* input, float* output, int M, int N) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid % 32;
    int wid  = tid / 32;

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

    // 寄存器数组缓存输入数据
    float reg_cache[MAX_ELEMS_PER_THREAD];

    // 第 1 遍（也是唯一一遍全局读取）：读数据 + Online 递推
    float local_m = -INFINITY;
    float local_d = 0.0f;
    int count = 0;

    for (int i = tid; i < N; i += blockDim.x) {
        float xi = x[i];
        reg_cache[count] = xi;  // 缓存到寄存器
        count++;

        float m_new = fmaxf(local_m, xi);
        local_d = local_d * expf(local_m - m_new) + expf(xi - m_new);
        local_m = m_new;
    }

    // 两级 Warp Shuffle 合并
    warpReduceOnline(local_m, local_d);

    __shared__ float warp_m[32];
    __shared__ float warp_d[32];
    if (lane == 0) {
        warp_m[wid] = local_m;
        warp_d[wid] = local_d;
    }
    __syncthreads();

    int num_warps = blockDim.x / 32;
    if (wid == 0) {
        local_m = (lane < num_warps) ? warp_m[lane] : -INFINITY;
        local_d = (lane < num_warps) ? warp_d[lane] : 0.0f;
        warpReduceOnline(local_m, local_d);
    }

    __shared__ float final_m, final_d;
    if (tid == 0) {
        final_m = local_m;
        final_d = local_d;
    }
    __syncthreads();

    float row_max = final_m;
    float inv_sum = 1.0f / final_d;

    // 从寄存器缓存中读取数据并写出（无需再读全局内存）
    int idx = 0;
    for (int i = tid; i < N; i += blockDim.x) {
        y[i] = expf(reg_cache[idx] - row_max) * inv_sum;
        idx++;
    }
}

六、多行并行+Grid Stride

前面的版本都是”一个 Block 处理一行”、“启动 M 个 Block”。当 M 非常大（如几万行）时，会启动同等数量的 Block，带来两个问题：一是 Grid 调度开销线性增加；二是每个 Block 只跑一行就退出，Block 内的 Shared Memory 缓冲、寄存器状态都无法跨行复用。

解决方案是用 Grid Stride 让固定数量的 Block（通常按 SM 数的若干倍设置）循环覆盖所有行，并继续叠加前面的所有优化（向量化加载 + Warp Shuffle 合并规约）。V4 在算法层面与 V2 完全相同（两遍扫描），区别只在于 Block 与行的映射关系——V2 是 1:1，V4 是 1:多。

Grid Stride 并不能解决”M 很小（如 batch=1）、N 很大”导致 Block 数不足的问题——那种场景需要让多个 Block 协作处理同一行（行内分段规约 + 两段式 Kernel 或原子合并）。

// V4: Grid Stride 多行处理 + 所有优化集成
__global__ void online_softmax_v4(float* input, float* output, int M, int N) {
    int tid = threadIdx.x;
    int lane = tid % 32;
    int wid  = tid / 32;

    // Grid Stride 遍历所有行
    for (int row = blockIdx.x; row < M; row += gridDim.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;

        // 第 1 遍：向量化 Online 递推
        float local_m = -INFINITY;
        float local_d = 0.0f;

        for (int i = tid; i < N4; i += blockDim.x) {
            float4 data = x4[i];
            float vals[4] = {data.x, data.y, data.z, data.w};
            for (int k = 0; k < 4; k++) {
                float m_new = fmaxf(local_m, vals[k]);
                local_d = local_d * expf(local_m - m_new) + expf(vals[k] - m_new);
                local_m = m_new;
            }
        }
        for (int i = N4 * 4 + tid; i < N; i += blockDim.x) {
            float xi = x[i];
            float m_new = fmaxf(local_m, xi);
            local_d = local_d * expf(local_m - m_new) + expf(xi - m_new);
            local_m = m_new;
        }

        // 两级 Warp Shuffle 合并
        warpReduceOnline(local_m, local_d);

        __shared__ float warp_m[32];
        __shared__ float warp_d[32];
        if (lane == 0) {
            warp_m[wid] = local_m;
            warp_d[wid] = local_d;
        }
        __syncthreads();

        int num_warps = blockDim.x / 32;
        if (wid == 0) {
            local_m = (lane < num_warps) ? warp_m[lane] : -INFINITY;
            local_d = (lane < num_warps) ? warp_d[lane] : 0.0f;
            warpReduceOnline(local_m, local_d);
        }

        __shared__ float final_m, final_d;
        if (tid == 0) {
            final_m = local_m;
            final_d = local_d;
        }
        __syncthreads();

        float row_max = final_m;
        float inv_sum = 1.0f / final_d;

        // 第 2 遍：向量化归一化
        for (int i = tid; i < N4; i += blockDim.x) {
            float4 data = x4[i];
            float4 result;
            result.x = expf(data.x - row_max) * inv_sum;
            result.y = expf(data.y - row_max) * inv_sum;
            result.z = expf(data.z - row_max) * inv_sum;
            result.w = expf(data.w - row_max) * inv_sum;
            y4[i] = result;
        }
        for (int i = N4 * 4 + tid; i < N; i += blockDim.x) {
            y[i] = expf(x[i] - row_max) * inv_sum;
        }

        __syncthreads();  // 确保 Shared Memory 在下一行之前被重置
    }
}

当 M 很大时，Grid Stride 让 Block 循环处理多行，好处是：

固定 Grid 大小避免过大的 Grid 调度开销
Block 内的 Shared Memory 和寄存器状态可以跨行复用
GPU 始终保持满载

当 M 较小时，grid_size = M，退化为 V2 的行为——每个 Block 处理一行。

C++ > CUDA

#C++ #CUDA

CUDA 实现online softmax算子

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

作者

Yu xin

发布于

2026年7月24日

许可协议

Linux设备树上一篇

CUDA 实现softmax算子下一篇