Softmax
Softmax
LeetGPU 题面
warp 级 max 和 sum
/* Warp-level reduction operation */
template<int WARP_SIZE = 32>__device__ __forceinline__ float warp_max(float val) {#pragma unroll for (int mask = WARP_SIZE >> 1; mask >= 1; mask >>= 1) { val = fmaxf(val, __shfl_xor_sync(0xffffffff, val, mask)); } return val;}修饰词 & 内联标识
__forceinline__
其中，__forceinline__ 是 CUDA 函数修饰符，强制编译器内联该函数，无视启发式阈值。
- 节省函数调用开销，包括寄存器保存、恢复和指令跳转等操作。
- 代价：代码膨胀，寄存器压力增加，可能降低占用率。
- 常用于小函数、device 函数、被
#pragma unroll循环体内调用的函数。
#pragma unroll
#pragma unroll 是 CUDA 编译器指令，告诉编译器完全展开紧跟其后的 for 循环。
像上述的代码中所示的那样， WARP_SIZE = 32， 所以 mask 的取值序列是 16-8-4-2-1，编译器会将循环展开为：
val = fmaxf(val, __shfl_xor_sync(0xffffffff, val, 16));val = fmaxf(val, __shfl_xor_sync(0xffffffff, val, 8));val = fmaxf(val, __shfl_xor_sync(0xffffffff, val, 4));val = fmaxf(val, __shfl_xor_sync(0xffffffff, val, 2));val = fmaxf(val, __shfl_xor_sync(0xffffffff, val, 1));之所以我们需要用到它，是因为它可以：
- 消除循环控制开销
- 更多指令级并行（ILP）
- 编译器优化机会
::: tip 指令级并行：处理器同时执行多条不相干的指令，而不是串行等待。 :::
__shfl_xor_sync
warp_max 用来在一个 warp 内求所有线程的最大值，并将最大值广播到 warp 中每个线程。
fmaxf 函数是求浮点数中的最大值。
__shfl_xor_sync 函数是 CUDA warp shuffle 指令，让同一个 warp 内的线程直接交换寄存器值，不走共享内存，不经过显存。
三个参数的含义
| 参数 | 含义 |
|---|---|
0xffffffff | 参与掩码，32 位全 1 = warp 内 32 个线程全部参与 |
val | 当前线程要交换的寄存器值 |
mask | XOR 偏移量，目标线程 ID = 当前线程 ID XOR mask |
工作机制
以 mask = 16 为例：
线程 0 ←→ XOR 16 → 线程 16线程 1 ←→ XOR 16 → 线程 17线程 2 ←→ XOR 16 → 线程 18...线程 15 ←→ XOR 16 → 线程 31__shfl_xor_sync 返回配对线程的 val 值。当前线程拿到对方的 val，对方也拿到当前线程的 val。
为什么用 XOR 而不是固定配对
XOR 天然形成蝶形交换（butterfly shuffle），这是经典并行归约的通信模式：
轮次 mask=16: 线程0←→线程16, 线程1←→线程17, ... (跨半区交换)轮次 mask=8 : 线程0←→线程8, 线程1←→线程9, ... (跨1/4区交换)轮次 mask=4 : 线程0←→线程4, 线程1←→线程5, ... (跨1/8区交换)轮次 mask=2 : 线程0←→线程2, 线程1←→线程3, ...轮次 mask=1 : 线程0←→线程1, 线程2←→线程3, ...每轮每个线程拿到配对线程的值和自己的值做 fmaxf，5 轮后所有 32 个线程都持有全 warp 的最大值。
和共享内存方案对比
| 对比项 | __shfl_xor_sync | 共享内存归约 |
|---|---|---|
| 延迟 | ~1 cycle（寄存器级） | ~20+ cycles（共享内存访问） |
需要 __syncthreads() | 不需要 | 需要 |
| 额外内存 | 无 | 需要分配 shared memory |
这就是为什么 warp-level reduction 几乎总是用 shuffle 指令而不是共享内存——寄存器交换是目前 GPU 上最快的线程间通信方式。
核心机制： XOR 保证每轮交换不出重复配对，蝶形把局部最大值逐层往外扩散，log₂N 轮覆盖全 warp。没有共享内存，没有同步屏障，纯寄存器交换，一个 cycle 完成。
三个场景不可以用该方法
- 跨 warp 不行：shuffle 只在同一个 warp 内通信。
- 通信模式固定：特定场景。
- 线程分歧会出错：某些线程走了不同分支（if/else)。
| 指令 | 行为 | 典型场景 |
|---|---|---|
__shfl_sync | 从指定线程广播 | 广播常量、配置 |
__shfl_up_sync | 拿上方线程的值 | 前缀和、流式处理 |
__shfl_down_sync | 拿下方线程的值 | 前缀和（反向）、数据移位 |
__shfl_xor_sync | XOR 蝶形配对 | 归约（求和、求最大值） |
求和也是一样
template<int WARP_SIZE = 32>__device__ __forceinline__ float warp_sum(float val) {#pragma unroll for (int mask = WARP_SIZE >> 1; mask >= 1; mask >>= 1) { // warp 内共享内存 寄存器直传 32 个线程 val += __shfl_xor_sync(0xffffffff, val, mask); } return val;}Block 级 max 和 sum
template<int BLOCK_SIZE = 256>__device__ __forceinline__ float block_max(float val) { constexpr int WARP_SIZE = 32; constexpr int NUM_WARPS = (BLOCK_SIZE + WARP_SIZE - 1) / WARP_SIZE;
const int warp_id = threadIdx.x / WARP_SIZE; const int lane_id = threadIdx.x % WARP_SIZE;
static __shared__ float shared_mem[NUM_WARPS];
const float warp_max_val = warp_max<WARP_SIZE>(val);
if (lane_id == 0) { shared_mem[warp_id] = warp_max_val; }
__syncthreads();
const float thread_val = (lane_id < NUM_WARPS) ? shared_mem[lane_id] : -FLT_MAX; const float total_max = warp_max<WARP_SIZE>(thread_val);
return total_max;}如上述代码所示，先求出一个 BLOCK 内的 WARP 数量，然后将线程 ID 拆分成 “哪个 warp” 和 “warp 里哪个”。
为什么只让 lane 0 写？ warp 内 32 线程持有的 warp_max_val 都相同（归约结果广播了）。32 个线程写同一个值到同一个位置 → 浪费带宽，且可能写冲突。lane_id == 0 防重复写入。
接下来干什么？ shared_mem 现在存的是每个 warp 的局部最大值，下一步再对这些值做一次归约，得到全局最大值。这是经典的两级归约：warp 内部归约 → 跨 warp 归约。
__syncthreads()
屏障指令。block 内所有线程都到达这条线，才一起放行。
为什么需要它： 共享内存跨 warp 读写，写方（各个 warp 的 lane 0 写 shared_mem）和读方（warp 0 的 lane 0-7 读 shared_mem）之间没有硬件保证顺序。不加屏障 → 可能读到旧值。
warp 0: 写 shared_mem[0] ─┐warp 1: 写 shared_mem[1] ─┤ 全到 → 同时放行warp 2: 写 shared_mem[2] ─┤warp 3: 写 shared_mem[3] ─┘ ↓ 开始读 shared_mem，值全写完，正确。第二级归约
__syncthreads() 之后，shared_mem 存的是每个 warp 的局部最大值。接下来用 lane_id < NUM_WARPS 做边界检查：
lane_id < NUM_WARPS→ 取shared_mem[lane_id]（有效 warp 最大值）- 其余 lane → 填
-FLT_MAX（负无穷，参与比较但永不获胜）
然后对 32 线程的 thread_val 再跑一次 warp_max，有效值留下，占位值被淘汰，得出全局最大值。
为何复用 warp_max： NUM_WARPS ≤ 32（256 线程 / 32 = 8），一级归约结果正好装进一个 warp。用 -FLT_MAX 填充多余 lane，同一个函数搞定两级归约。
求和也是一样
/* Block-level reduction operations */
template<int BLOCK_SIZE = 256>__device__ __forceinline__ float block_sum(float val) { constexpr int WARP_SIZE = 32; constexpr int NUM_WARPS = (BLOCK_SIZE + WARP_SIZE - 1) / WARP_SIZE;
const int warp_id = threadIdx.x / WARP_SIZE; const int lane_id = threadIdx.x % WARP_SIZE;
// warp 间 共享内存 同一个 block 内的所有 warp static __shared__ float shared_mem[NUM_WARPS];
const float warp_sum_val = warp_sum<WARP_SIZE>(val);
if (lane_id == 0) { shared_mem[warp_id] = warp_sum_val; }
__syncthreads();
const float thread_val = (lane_id < NUM_WARPS) ? shared_mem[lane_id] : 0.0f; const float total_sum = warp_sum<WARP_SIZE>(thread_val);
return total_sum;}softmax
最后是将最大值和求和值拿来做 softmax 的运算。
/* Softmax kernel implementation */
template<int BLOCK_SIZE = 256>__global__ void softmax_kernel(const float* input, float* output, const int N) { const int tid = threadIdx.x;
// 找到所有线程中的最大值 float max_val = -FLT_MAX; for (int i = tid; i < N; i += blockDim.x) { max_val = fmaxf(max_val, input[i]); } max_val = block_max<BLOCK_SIZE>(max_val);
// 计算所有指数函数的和 float sum_val = 0.0f; for (int i = tid; i < N; i += blockDim.x) { const float exp_val = expf(input[i] - max_val); sum_val += exp_val; } sum_val = block_sum<BLOCK_SIZE>(sum_val);
// 计算最终的归一化值 for (int i = tid; i < N; i += blockDim.x) { output[i] = expf(input[i] - max_val) / sum_val; }}分层接力
线程扫数据 → 局部 maxwarp 归约 → warp max（第一级）共享内存 → 传给 warp 0warp 归约 → block max（第二级，全局值）::: tip 每一级往上是吞吐压缩，256 → 8 → 1，O(log₂N) 而不是 O(N)。 :::
HOST
/* HOST */extern "C" void solve(const float* input, float* output, const int N) { constexpr int THREADS_PER_BLOCK = 1024; softmax_kernel<THREADS_PER_BLOCK><<<1, THREADS_PER_BLOCK>>>(input, output, N); cudaDeviceSynchronize();}