CUDA算子优化-Soft_max
高性能计算系列文章目录：
- C++ 资源管理
- C++ 并发
- CUDA 编程简明教程
- CUDA 核函数线程配置
- 程序性能指标与 Roofline 模型
- CUDA 算子优化 - Reduction
- CUDA 算子优化 - GEMM
- CUDA 算子优化 - Transpose
系列文章相关完整代码已上传至 Github，欢迎 Review ➜ CudaNote
注：原文中的数学公式以公式对象呈现，剪藏过程中丢失。文中"其中是输入向量""设输入总数据量为"等句子以及性能对比表中的具体数值存在缺失，完整公式请结合上下文与原作者的 CudaNote 仓库 对照阅读。
1. 前言
Softmax 函数作为深度学习中的核心激活函数，在神经网络推理和训练过程中占据重要地位。本文将展开对 Softmax 函数在 CUDA 平台上的多种优化策略进行讨论。包括内存访问优化、并行化策略等技术。
Softmax 函数将一个实数向量映射为概率分布，在实际工程中，其数学定义为：
其中是输入向量，是向量维度，，减去最大值的操作是为了防止数据溢出。
如下图，从 Softmax 的数学定义中不难看出，其计算主要包括：最大值归约、指数和归约和归一化三个步骤，基于此，实现出 cpu 和 GPU naive 版本的 Softmax，并逐步展开优化。
CPU 版本
void soft_max_host(float *input, float *output, size_t size)
{
// 1. Calculate the maximum value
float max_val = input[0];
for (int i = 1; i < size; ++i)
{
if (input[i] > max_val)
{
max_val = input[i];
}
}
// 2. Calculate the sum of exponents
float sum = 0.0f;
for (int i = 0; i < size; ++i)
{
output[i] = expf(input[i] - max_val);
sum += output[i];
}
// 3.Normalize
for (int i = 0; i < size; ++i)
{
output[i] /= sum;
}
}
GPU naive 版本
__global__ void soft_max_naive_kernel_v1(const float *input, float *output, size_t size)
{
int tid = threadIdx.x;
// Only thread 0 performs all computations
if (tid == 0) {
// Phase 1: Find global maximum
float global_max = -FLT_MAX;
for (int i = 0; i < size; i++) {
global_max = fmaxf(global_max, input[i]);
}
// Phase 2: Compute exponential sum
float global_sum = 0.0f;
for (int i = 0; i < size; i++) {
global_sum += expf(input[i] - global_max);
}
// Phase 3: Compute final softmax values
for (int i = 0; i < size; i++) {
output[i] = expf(input[i] - global_max) / global_sum;
}
}
}
可以看得出，上述计算过程中，CPU 版本每个步骤都在循环，效率低下。而 GPU 版本 soft_max_naive_kernel_v1 只在一个线程中运行，和 CPU 是一样的逻辑。因此，基于 GPU naive 版本，我们先进行归约操作，然后进行访存优化。
2. Softmax 归约优化
对于归约的详细内容参考上一篇文章 CUDA 算子优化-Reduction，本次对 Softmax 的归约，我们不再通过线程 id 约束来确定归约范围和步长，直接采用 CUDA 的 __shfl_down_sync 来实现。
__shfl_down_sync(mask, var, offset) 是一个warp 内线程间数据交换的指令：
- mask：指定参与操作的线程（0xFFFFFFFF 表示所有 32 个线程）
- var：要交换的变量
- offset：向下交换的偏移量
__shfl_down_sync 计算效率更高，其所有操作在寄存器间完成，零共享内存开销，提供指令级并行。在进行具体实现时，我们需要考虑，当处理大规模数据时，GPU 的线程无法完全覆盖整个数据，这个时候就需要对数据进行分块处理。对于 Softmax 的归约优化，将其定义为分层归约。
设输入总数据量为，总线程数为，GPU Warp 大小，则共有 Warp 数量为，线程索引：。针对 softmax 的计算步骤，分别讨论分层归约的几个部分。
2.1. 数据分块策略
不管是求最大值还是求和，在分层归约中，数据分块策略是相同的，定义每个线程获得的数据块：
其中，用显式公式表示为：
这种数据映射方式确保每个线程处理大致相同数量的元素，虽然采用跨步访问不如连续访问高效，但避免了 bank conflict，同时充分利用了所有线程的计算能力，适用于任意大小的输入数组。
先对数据分块，然后再分块内计算最大值或求和，通过这一操作，可以让数据的局部最大值或局部和刚好覆盖整个核函数的线程，在此基础上，分层进行 Warp 归约和 Block 归约，获得最终的全局最大值和全局和。，时：
- 线程 0 处理：
- 线程 1 处理：
- 线程 2 处理：
- 线程 3 处理：
2.2. 最大值归约
2.3. 线程局部最大值
每个线程在其负责的分块子集上计算局部最大值：
用数学公式表达为：
这一步确保每个数据点都参与最大值计算。
2.4. Warp 内归约
经过线程局部最大值计算，每个线程获得数据分块内的最大值，即通过分块操作，数据最大值刚好覆盖线程配置，此时，通过 warp 内，获取每个 warp 内最大值。设每个 warp 包含 32 个线程，其中，线程 ID：，warp ID：，lane ID 为：。
Warp 内的归约利用 __shfl_down_sync 实现树状结构：
对于偏移量，执行：
其中上标表示第次迭代后的值。经过次迭代后，warp 的最大值为：
且对于所有，都有。
2.5. Block 级归约
线程块内共有个 warps，收集所有 warp 的最大值：
再次使用 warp 级归约（由前个线程执行）：
2.6. 分块函数
整个归约过程可以表示为嵌套的最大值运算：
这个表达式清晰地展示了从元素级到线程级、Warp 级、最终到 Block 级的层层归约过程。
2.7. 指数和计算归约
2.8. 线程局部指数和计算
使用稳定的最大值，同样的数据分块策略，每个线程在其负责的分块子集上计算局部最大值：
2.9. Warp 级求和归约
采用与最大值归约类似的树状结构：
对于最终获得 Warp 级和：
2.10. Block 级求和归约
收集所有 Warp 的和：
再次应用 Warp 级求和：
整个归约过程可以表示为嵌套的指数和运算：
2.11. Softmax 归一化
同样的数据分块策略，每个线程为其负责的每个元素计算最终结果：
对于所有
2.12. 多 Block 扩展
2.13. 全局归约的需求
当极大时，需要多个 Block 协作。同样的分层归约思路，设共有个 Blocks，每个 Block 计算：
- 局部最大值：
- 局部指数和：
2.14. 全局最大值计算
2.15. 全局指数和计算
需要基于统一的重新计算指数和：
具体示例：当：不能直接求和，因为每个是基于不同的计算的。
2.16. 全局关系的数学表达
__device__ float warp_reduce_max(float val)
{
for (int offset = 16; offset > 0; offset /= 2) {
val = fmaxf(val, __shfl_down_sync(0xFFFFFFFF, val, offset));
}
return val;
}
__device__ float warp_reduce_sum(float val)
{
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xFFFFFFFF, val, offset);
}
return val;
}
__global__ void soft_max_reduce_kernel(const float *input, float *output, size_t size)
{
int tid = threadIdx.x;
int warp_id = tid / 32;
int lane_id = tid % 32;
extern __shared__ float shared_mem[];
float *warp_max = shared_mem;
float *warp_sum = shared_mem + (blockDim.x / 32);
float thread_max = -INFINITY;
for (int i = tid; i < size; i += blockDim.x) {
if (input[i] > thread_max) thread_max = input[i];
}
float warp_max_val = warp_reduce_max(thread_max);
if (lane_id == 0) warp_max[warp_id] = warp_max_val;
__syncthreads();
float global_max = -INFINITY;
if (tid < (blockDim.x / 32)) global_max = warp_max[tid];
global_max = warp_reduce_max(global_max);
if (tid == 0) warp_max[0] = global_max;
__syncthreads();
global_max = warp_max[0];
float thread_sum = 0.0f;
for (int i = tid; i < size; i += blockDim.x) {
thread_sum += expf(input[i] - global_max);
}
float warp_sum_val = warp_reduce_sum(thread_sum);
if (lane_id == 0) warp_sum[warp_id] = warp_sum_val;
__syncthreads();
float global_sum = 0.0f;
if (tid < (blockDim.x / 32)) global_sum = warp_sum[tid];
global_sum = warp_reduce_sum(global_sum);
if (tid == 0) warp_sum[0] = global_sum;
__syncthreads();
global_sum = warp_sum[0];
for (int i = tid; i < size; i += blockDim.x) {
output[i] = expf(input[i] - global_max) / global_sum;
}
}
Reduce 版本 soft_max_reduce_kernel 通过归约的方式实现了并行化内存访问，所有线程以分块跨步方式协同访问数据。若数据规模为，完成一次计算的过程中，输入数据读取进行 2 次遍历：次 float 加载 × 4 字节 =字节，输出数据写入进行 1 次遍历：次 float 存储 × 4 字节 =字节，此外，最大值与指数和进行次共享内存访问：warp 最大值存储：次 float 写入 × 4 字节 =字节，warp 和存储：次 float 写入 × 4 字节 =字节。注：
在最大值查找阶段，进行次比较操作，指数和计算阶段，进行次指数计算 + 次加法，归一化阶段，进行次指数计算 + 次乘法，同时，归约阶段，warp 最大值归约进行次比较操作（log₂32=5 步），warp 和归约进行次加法操作，代码实现：
2.17. 总内存访问量
与 Naive 版本相比，由于并行化带来的内存访问重叠和带宽利用率提升，性能远优于 Naive 版本。
3. float4 向量化访存优化
Reduce 版本虽然实现了并行化，但对全局内存的访问模式难以充分利用宽内存总线的优势。全局内存每个内存事务按照 4 字节传输，因此，我们在分层归约的基础上，对 Softmax 核函数进行 float4 向量化访存优化，向量化访存和全局内存事务相关内容可参考本系列文章 CUDA 算子优化-Transpose，本文不做展开，直接进行实现和性能分析。
3.1. 总计算量
__global__ void soft_max_vectorized_kernel(const float *input, float *output, size_t size)
{
int tid = threadIdx.x;
int warp_id = tid / 32;
int lane_id = tid % 32;
extern __shared__ float shared_mem[];
float *warp_max = shared_mem;
float *warp_sum = shared_mem + (blockDim.x / 32);
// 使用float向量化加载优化
float thread_max = -INFINITY;
const int vector_size = 4;
const int num_vectors = (size + vector_size - 1) / vector_size;
const float4 *input4 = reinterpret_cast<const float4 *>(input);
for (int i = tid; i < num_vectors; i += blockDim.x)
{
float4 val4;
int base_idx = i * vector_size;
if (base_idx + 3 < size)
{
val4 = input4[i];
}
else
{
// 处理边界情况
val4.x = (base_idx < size) ? input[base_idx] : -INFINITY;
val4.y = (base_idx + 1 < size) ? input[base_idx + 1] : -INFINITY;
val4.z = (base_idx + 2 < size) ? input[base_idx + 2] : -INFINITY;
val4.w = (base_idx + 3 < size) ? input[base_idx + 3] : -INFINITY;
}
thread_max = max(thread_max, val4.x);
thread_max = max(thread_max, val4.y);
thread_max = max(thread_max, val4.z);
thread_max = max(thread_max, val4.w);
}
float warp_max_val = warp_reduce_max(thread_max);
if (lane_id == 0)
{
warp_max[warp_id] = warp_max_val;
}
__syncthreads();
float global_max = -INFINITY;
if (tid < (blockDim.x / 32))
{
global_max = warp_max[tid];
}
global_max = warp_reduce_max(global_max);
if (tid == 0)
{
warp_max[0] = global_max;
}
__syncthreads();
global_max = warp_max[0];
// 使用float4向量化加载指数和计算
float thread_sum = 0.0f;
for (int i = tid; i < num_vectors; i += blockDim.x)
{
float4 val4;
int base_idx = i * vector_size;
if (base_idx + 3 < size)
{
val4 = input4[i];
}
else
{
val4.x = (base_idx < size) ? input[base_idx] : 0.0f;
val4.y = (base_idx + 1 < size) ? input[base_idx + 1] : 0.0f;
val4.z = (base_idx + 2 < size) ? input[base_idx + 2] : 0.0f;
val4.w = (base_idx + 3 < size) ? input[base_idx + 3] : 0.0f;
}
// 计算指数和
thread_sum += expf(val4.x - global_max);
thread_sum += expf(val4.y - global_max);
thread_sum += expf(val4.z - global_max);
thread_sum += expf(val4.w - global_max);
}
float warp_sum_val = warp_reduce_sum(thread_sum);
if (lane_id == 0)
{
warp_sum[warp_id] = warp_sum_val;
}
__syncthreads();
float global_sum = 0.0f;
if (tid < (blockDim.x / 32))
{
global_sum = warp_sum[tid];
}
global_sum = warp_reduce_sum(global_sum);
if (tid == 0)
{
warp_sum[0] = global_sum;
}
__syncthreads();
global_sum = warp_sum[0];
// 使用float4向量化存储优化输出
float4 *output4 = reinterpret_cast<float4 *>(output);
float inv_global_sum = 1.0f / global_sum;
for (int i = tid; i < num_vectors; i += blockDim.x)
{
float4 val4;
float4 result4;
int base_idx = i * vector_size;
if (base_idx + 3 < size)
{
val4 = input4[i];
}
else
{
val4.x = (base_idx < size) ? input[base_idx] : 0.0f;
val4.y = (base_idx + 1 < size) ? input[base_idx + 1] : 0.0f;
val4.z = (base_idx + 2 < size) ? input[base_idx + 2] : 0.0f;
val4.w = (base_idx + 3 < size) ? input[base_idx + 3] : 0.0f;
}
result4.x = expf(val4.x - global_max) * inv_global_sum;
result4.y = expf(val4.y - global_max) * inv_global_sum;
result4.z = expf(val4.z - global_max) * inv_global_sum;
result4.w = expf(val4.w - global_max) * inv_global_sum;
if (base_idx + 3 < size)
{
output4[i] = result4;
}
else
{
if (base_idx < size)
output[base_idx] = result4.x;
if (base_idx + 1 < size)
output[base_idx + 1] = result4.y;
if (base_idx + 2 < size)
output[base_idx + 2] = result4.z;
if (base_idx + 3 < size)
output[base_idx + 3] = result4.w;
}
}
}
采用 float4 向量化访存，对于个数据元素，将其分组为个向量，在输入数据读取时进行两次遍历 (最大值查找：次 float4 加载 →字节，次 float4 加载 →字节)，输出数据写入进行一次遍历 (结果存储：次 float4 存储 →字节)，此外，最大值与指数和进行次共享内存访问（warp 最大值存储：次 float 写入 →字节，warp 和存储：次 float 写入 →字节），计算强度：
在最值查找阶段，每个向量 4 次比较操作：；计算指数和阶段，每个向量 4 次指数计算：和 3 次加法（累加 4 个值）：；归一化阶段，每个向量 4 次指数计算：，4 次乘法：和次加法（可选）：；同时，warp 归约操作开销：。代码实现：
3.2. 总内存访问字节数
将 Reduce 版本和向量化版本进行对比：
| 性能指标 | Reduce 版本 | 向量化版本 | 变化幅度 |
|---|---|---|---|
| 内存访问量 | bytes | bytes | 相同 |
| 计算量 | FLOP | FLOP | +37.5% |
| 计算强度 | 0.333 FLOP/byte | 0.458 FLOP/byte | +37.5% |
| 内存事务数 | 次 | 次 | -75% |
| 带宽需求 | 高 | 中等 | 优化 |
向量化并没有减少总的内存访问字节数，但计算量增加 37.5%，计算强度同步提升 37.5%，同时，通过减少 75% 的内存事务次数，大幅提升了带宽利用率。