GPU并行计算（CUDA） —— Softmax算子逐步优化：从基础实现到online softmax

侦探不想感冒～

405人浏览 · 2026-05-26 14:29:56

侦探不想感冒～ · 2026-05-26 14:29:56 发布

前言

在分类任务中，模型输出的logits需要通过Softmax函数转换为概率分布。Softmax确保所有类别的预测概率之和为1，使得这些概率值具有实际意义。这一过程对于多分类问题至关重要，并且在深度学习模型中广泛应用。
本文将对softmax算子的多个版本进行分析并逐步优化。

Softmax公式

$\text{softmax}(\mathbf{z})_i = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}}$
直接计算指数函数可能会遇到数值稳定性问题。
当指数函数的输入非常大时, 会导致结果趋向于无穷大（Inf），这可能会导致计算结果出现非数字（NaN）错误。为了避免这种情况，一个常见的解决方案是在计算指数之前，从所有输入值中减去最大值。这样做可以确保指数函数的输入不会过大，从而提高数值计算的稳定性。

Safe Softmax

$\text{softmax}(\mathbf{z})_i = \frac{e^{-\max(\mathbf{z})} \cdot e^{z_i}}{e^{-\max(\mathbf{z})} \cdot \sum_{j=1}^{K} e^{z_j}}=\frac{e^{z_i - \max(\mathbf{z})}}{\sum_{j=1}^{K} e^{z_j - \max(\mathbf{z})}}$
我们可以在 naive softmax 公式的基础上，让分子和分母同时乘以 $e^{-\max(\mathbf{z})}$ 。由于分子和分母乘以相同的非零因子，分式的值不变。合并指数项后，每个元素都会减去当前行的最大值，此时指数函数的输入范围被限制在 $(-\infty, 0]$ ，因此指数结果位于 $(0, 1]$ ，从而避免了指数溢出问题。

1.CPU实现Softmax

通常来讲, cpu实现safe softmax需要有三个步骤(三次遍历):

求当前行的最大值: ${\max(\mathbf{z})}$
求每个元素的exp值: $e^{z_j-\max(\mathbf{z})}$ 和当前行的exp总和: $Sum=\sum_{j=1}^{K} e^{z_j - \max(\mathbf{z})}$
求出当前行每个元素的softmax值 $e^{z_j-\max(\mathbf{z})}/Sum$

代码如下:

void mySoftmax_cpu(float* out, const float* in, int N, int C) {
    // CPU N: batch_size C: vector_len
    for(int i = 0; i < N; i++){
        // 锁定当前循环处理的行
        const float *in_row = in + i * C;
        float *out_row = out + i * C;
        // 1.求当前行最大值
        float maxval = -INFINITY;
        for(int j = 0; j < C; j++){
            if(in_row[j] > maxval){
                maxval = in_row[j];
            }
        }
        // 2.计算每个元素的exp值和该行的exp和
        float sum = 0.0f;
        for(int j = 0; j < C; j++){
            out_row[j] = expf(in_row[j] - maxval);
            sum += out_row[j];
        }
        // 3.计算每个位置的softmax结果
        for(int j = 0; j < C; j++){
            out_row[j] /= sum;
        }
    }
}

2.Softmax_v1 (Reduce Softmax)

我们可以发现，三次遍历中的前两次遍历一个是求最大值，一个是求和，这正好可以利用GPU的并行性进行规约(Reduce)操作，以大幅提升性能。

详细Reduce讲解见之前的博客：
GPU并行计算（CUDA） – 归约（Reduce）算子深度优化：从基础实现到极致性能

这里使用该博客的reduce_v3，使用两次warp_shuffle操作对一个线程块内的数据进行规约，减少对shared memory的访问。与reduce_v3不同的是，这里我们让一个线程块处理一行数据，而不是用多个线程块规约一行数据。因此，块内规约的结果就是最终结果，我们这里只需要调用一次核函数，而不是两次。
步骤如下:

Reduce求最大值
Reduce求exp总和
求每个元素的exp值，再求softmax值

除了Reduce操作与cpu实现不同，为了方便前两步Reduce操作，这里把求每个元素的exp值由第二步换成了第三步。
softmax_v1代码如下：

__global__ void mySoftmax_v1(float* out, const float* in, int N, int C){
    // 每个block处理一行, 定位到该block所处理的行
    unsigned int tid = threadIdx.x;
    const float *in_row = in + blockIdx.x * C;
    float *out_row = out + blockIdx.x * C;
    // 求最大值
    float max = reduce_max(in_row, C);    
    // 通过 Shared Memory 广播最大值给所有线程
    __shared__ float s_max;
    if (threadIdx.x == 0) s_max = max;
    __syncthreads();    // 确保0号线程写完, 其他线程再读
    max = s_max;    // 每个线程得到该行全局最大值

    // 求和
    float sum = reduce_sum(in_row, C, max);
    __shared__ float s_sum;
    if(threadIdx.x == 0) s_sum = sum;
    __syncthreads();
    sum = s_sum;

    for(int i = tid; i < C; i += blockDim.x){
        out_row[i] = expf(in_row[i] - max) / sum;
    }
}

求expSum和max的规约算子代码如下：

constexpr int warpSize = 32;
__device__ float block_reduce_sum(float val){
    unsigned int tid = threadIdx.x;
    int lane = tid % warpSize;
    int warpId = tid / warpSize;
    val += __shfl_down_sync(0xffffffff, val, 16);
    val += __shfl_down_sync(0xffffffff, val, 8);
    val += __shfl_down_sync(0xffffffff, val, 4);
    val += __shfl_down_sync(0xffffffff, val, 2);
    val += __shfl_down_sync(0xffffffff, val, 1);

    // 为方便二次warpShuf, 一个块内最多32个warp
    __shared__ float warpSum[32]; 
    if(lane == 0){
        warpSum[warpId] = val;
    }
    __syncthreads();
    if(warpId == 0){
        val = (tid < blockDim.x / warpSize) ? warpSum[tid] : 0.0f;
        val += __shfl_down_sync(0xffffffff, val, 16);
        val += __shfl_down_sync(0xffffffff, val, 8);
        val += __shfl_down_sync(0xffffffff, val, 4);
        val += __shfl_down_sync(0xffffffff, val, 2);
        val += __shfl_down_sync(0xffffffff, val, 1);
    }
    return val;
}

__device__ float block_reduce_max(float val){
    unsigned int tid = threadIdx.x;
    int lane = tid % warpSize;
    int warpId = tid / warpSize;
    val = fmaxf(val,  __shfl_down_sync(0xffffffff, val, 16));
    val = fmaxf(val,  __shfl_down_sync(0xffffffff, val, 8));
    val = fmaxf(val,  __shfl_down_sync(0xffffffff, val, 4));
    val = fmaxf(val,  __shfl_down_sync(0xffffffff, val, 2));
    val = fmaxf(val,  __shfl_down_sync(0xffffffff, val, 1));

    // 为方便二次warpShuf, 一个块内最多32个warp
    __shared__ float warpMax[32]; 
    if(lane == 0){
        warpMax[warpId] = val;
    }
    __syncthreads();
    if(warpId == 0){
        val = (tid < blockDim.x / warpSize) ? warpMax[tid] : -INFINITY;
        val = fmaxf(val,  __shfl_down_sync(0xffffffff, val, 16));
        val = fmaxf(val,  __shfl_down_sync(0xffffffff, val, 8));
        val = fmaxf(val,  __shfl_down_sync(0xffffffff, val, 4));
        val = fmaxf(val,  __shfl_down_sync(0xffffffff, val, 2));
        val = fmaxf(val,  __shfl_down_sync(0xffffffff, val, 1));
    }
    return val;
}

__device__ float reduce_max(const float *in_row, int C){
    unsigned tid = threadIdx.x;
    float max = -INFINITY;
    for(int i = tid; i < C; i += blockDim.x){
        max = fmaxf(max, in_row[i]);
    }
    max = block_reduce_max(max);
    return max;
}

__device__ float reduce_sum(const float *in_row, int C, float max){
    unsigned tid = threadIdx.x;
    float sum = 0.0f;
    for(int i = tid; i < C; i += blockDim.x){
        sum += expf(in_row[i] - max);
    }
    sum = block_reduce_sum(sum);
    return sum;
}

3.Softmax_v2 (Online Softmax)

前文讲到，为了防止Softmax指数爆炸，我们引入了 ${-\max(\mathbf{z})}$ 。这导致了计算累加和依赖于最大值，我们必须先遍历一次找到最大值，再遍历一次计算累加和，最后再遍历一次计算最终结果。计算一行的softmax需要访存3K次。在计算吞吐量极高的GPU上，频繁的读取显存操作就像在高速公路上每开一公里就停下来交一次费，性能瓶颈卡在显存带宽上。
既然遍历三次太慢，我们能不能将计算累加和和最大值合并为一次遍历？

Online Softmax由NVIDIA的研究人员在2018年提出。思想是：流式地计算累加和，通过维护两个状态变量来逐步更新。

$m_i$ : 当前遍历到的所有元素的最大值。
$d_i$ : 基于当前遍历到的所有元素和 $m_i$ 计算的分母(累加和).

公式如下：
$m_k = \max(m_{k-1}, x_k) , \quad ( m_0 = -\infty )$
$d_k = d_{k-1} \cdot e^{m_{k-1} - m_k} + e^{x_k - m_k} , \quad ( d_0 = 0 )$

下图为Online Softmax推导过程：
在这里插入图片描述
我们可以发现，每个step的 $d_i$ 的表达式都可以转换为带有 $d_{k-1}$ 的式子求出，这意味着我们可以根据之前step的局部最大值和局部累加和逐步计算当前step的累加和，而不需要先求出全局最大值，再计算累加和。
代码如下：

__device__ float warp_reduce_max(float val){
    val = fmaxf(val, __shfl_down_sync(0xffffffff, val, 16));
    val = fmaxf(val, __shfl_down_sync(0xffffffff, val, 8));
    val = fmaxf(val, __shfl_down_sync(0xffffffff, val, 4));
    val = fmaxf(val, __shfl_down_sync(0xffffffff, val, 2));
    val = fmaxf(val, __shfl_down_sync(0xffffffff, val, 1));
    return __shfl_sync(0xffffffff, val, 0);     // val在warp内广播
}

__device__ float warp_reduce_sum(float val){
    val +=  __shfl_down_sync(0xffffffff, val, 16);
    val +=  __shfl_down_sync(0xffffffff, val, 8);
    val +=  __shfl_down_sync(0xffffffff, val, 4);
    val +=  __shfl_down_sync(0xffffffff, val, 2);
    val +=  __shfl_down_sync(0xffffffff, val, 1);
    return __shfl_sync(0xffffffff, val, 0); 
}
__global__ void mySoftmax_v2_online_warpReduce(float* out, const float* in, int N, int C){
    // 一个block处理一行，一个block1个warp，一个线程处理多个数据，online与reduce结合，先online后reduce
    unsigned int tid = threadIdx.x;     // tid == laneId
    const float *in_row = in + blockIdx.x * C;
    float *out_row = out + blockIdx.x * C;

    float local_max = -INFINITY;
    float local_sum = 0.0f;
    // 32个线程（一个warp）并行计算online softmax
    for(int i = tid; i < C; i += blockDim.x){     
        float x = in_row[i];
        float new_max = fmaxf(local_max, x);
        local_sum = local_sum * expf(local_max - new_max) + expf(x - new_max);
        local_max = new_max;
    }
    // warp内进行reduce softmax
    float g_max = warp_reduce_max(local_max);
    local_sum = local_sum * expf(local_max - g_max);   // 修正局部sum
    float g_sum = warp_reduce_sum(local_sum);

    for(int i = tid; i < C; i += blockDim.x){
        out_row[i] = expf(in_row[i] - g_max) / g_sum;
    }
}

这里依旧是一个线程块处理一行数据，一个线程块设置32个线程，即一个warp. 让32个线程并行地计算online softmax处理属于自己的数据，得到局部和和局部最大值，然后修正局部和，利用warp_shfl进行warp内规约。
该版本可以看作 online softmax 与 warp reduce 的混合实现。相较于单线程 online softmax，它将一行数据的处理从单线程串行扩展为 warp 内 32 个 lane 并行，因此具有更高的并行性；相较于传统 safe softmax / block reduce softmax，它虽然增加了一部分 expf 计算，但将输入 in 的全局内存读取次数由约 3K 次降低到约 2K 次，即从 3K loads + K stores 变为 2K loads + K stores.

总结

本文从 naive softmax 出发，首先分析了直接计算指数可能带来的数值溢出问题，并引入 safe softmax 通过减去最大值提升数值稳定性。随后，本文将 softmax 的最大值计算和指数和计算映射到 GPU Reduce 操作中，提高了行内并行度。进一步地，本文介绍了 online softmax 的思想，将最大值更新与分母累加融合到一次遍历中，从而减少对输入数据的全局内存读取次数。最后实现的 online + warp reduce 混合版本，在保持数值稳定性的同时，兼顾了 online softmax 的低访存优势和 warp 级并行计算能力，适合中小规模向量长度下的 softmax 优化。总体来看，softmax 优化的核心是在数值稳定性、并行度、访存次数和计算开销之间进行权衡。