GPU 架构学习
第二十一课:从低效 Reduction Kernel 到高性能归约
这一课用 Reduction 实战串联原子操作、Shared Memory、分支发散、Warp Shuffle、向量化加载和多级归约。
目录
一、Reduction 为什么不像向量加法那么简单二、串行 Reduction三、树形 Reduction四、版本零:所有线程直接 Atomic Add原子操作为什么能保证正确性为什么这个版本通常很慢原子版本什么时候可能合理五、基本优化思想:先在 Block 内归约六、为什么使用 Shared Memory七、为什么加载后必须同步八、低效版本一:交错寻址归约九、交错寻址的第一个问题:分支发散十、交错寻址的第二个问题:取模运算十一、Shared Memory Bank Conflict十二、优化版本二:顺序寻址归约顺序寻址的优势十三、为什么每一轮都需要同步十四、危险写法:条件 Barrier十五、优化版本三:每个线程先加载两个元素这样做有什么好处Global Memory 访问是否合并十六、进一步优化:Grid-Stride LoopGrid-Stride Loop 的优势Grid-Stride Loop 的权衡十七、为什么归约通常是带宽受限十八、最后一个 Warp 是否还需要 __syncthreads()十九、Warp Shuffle 是什么二十、Warp Shuffle ReductionWarp Shuffle 的优势二十一、一个 Block 内包含多个 Warp 怎么办二十二、现代 Block Reduction 结构二十三、一个优化后的 Block Reduction Kernel二十四、这个 Kernel 中的数据层次二十五、为什么仍然需要第二级归约方法一:启动第二个 Kernel方法二:每个 Block 最后执行一次 Atomic Add方法三:Cooperative Grid Synchronization二十六、为什么第二级 Kernel 通常很小二十七、向量化加载向量化加载的优势向量化加载的要求二十八、向量化并不总是更快二十九、使用多个局部累加器提高 ILP多累加器的代价三十、Block Size 如何选择三十一、Block Size 不必与输入长度完全匹配三十二、不同归约操作的单位元三十三、浮点加法不满足严格结合律三十四、Atomic Reduction 的结果可能不完全确定三十五、树形归约的数值误差三十六、Kahan Summation 的基本思想三十七、使用 FP64 累加的权衡三十八、使用 Nsight Compute 比较各版本三十九、预期的性能变化V0:每线程原子V1:交错 Shared MemoryV2:顺序寻址V3:每线程加载两个元素V4:Warp ShuffleV5/V6:Grid-Stride、向量化、多累加器四十、优化成功的最终表现四十一、为什么直接使用 CUB 往往更合理四十二、自定义 Reduction 仍然有价值的场景四十三、Reduction 与 Softmax 的关系四十四、Reduction 的完整优化路径本课核心结论第二十二课:Softmax 为什么难以高效,以及 FlashAttention 如何减少显存访问导语:这一课用 Reduction 实战串联原子操作、Shared Memory、分支发散、Warp Shuffle、向量化加载和多级归约。
Reduction,归约,是把一组数据通过某种运算合并成一个结果。
例如求和:
其他常见归约包括:
最大值 Max
最小值 Min
逻辑与 / 或
范数
均值
方差中的部分统计量
Softmax 中的最大值和指数和
Reduction 很适合用来学习 GPU 优化,因为它同时涉及:
- 大规模并行读取;
- 线程间通信;
- Shared Memory;
- Barrier;
- Warp Shuffle;
- 原子操作;
- 多级 Kernel;
- 浮点误差;
- 内存带宽限制。
一、Reduction 为什么不像向量加法那么简单
向量加法中,每个输出元素相互独立:
C[i] = A[i] + B[i];
可以直接采用:
一个线程 → 一个输出元素
但 Reduction 最终只有一个输出:
如果每个线程计算一个输入元素,最终仍需要把所有线程的结果合并起来。
因此 Reduction 包含两个阶段:
大量线程并行读取和局部计算
↓
逐步合并成少量结果
↓
最终得到一个结果
越接近最终结果,可用的并行任务越少。
二、串行 Reduction
CPU 风格的串行实现为:
float sum = 0.0f;
for (int i = 0; i < n; ++i) {
sum += input[i];
}
它形成一条很长的依赖链:
sum₁ = sum₀ + x₀
sum₂ = sum₁ + x₁
sum₃ = sum₂ + x₂
...
每次加法都依赖上一次结果,因此难以并行。
GPU 需要把它改造成树形结构。
三、树形 Reduction
假设有8个数:
第一层并行计算:
第二层:
第三层:
示意:

图示:三、树形 Reduction
Reduction 深度从串行的:
降低为:
但需要线程之间交换中间结果。
四、版本零:所有线程直接 Atomic Add
最容易想到的 GPU 实现是:
__global__ void reduce_atomic(
const float* input,
float* output,
int n)
{
int i = blockIdx.x * blockDim.x
+ threadIdx.x;
if (i < n) {
atomicAdd(output, input[i]);
}
}
启动前需要把:
*output = 0.0f;
初始化为0。
每个线程读取一个元素,然后对同一个输出地址执行:
atomicAdd(output, value);
原子操作为什么能保证正确性
普通代码:
*output += input[i];
实际包含:
读取 output
执行加法
写回 output
多个线程可能同时读取旧值,造成更新丢失。
atomicAdd 保证一次读—改—写操作不可被其他线程插入,因此结果不会丢失。
为什么这个版本通常很慢
所有线程竞争同一个地址:
Thread 0 ─┐
Thread 1 ─┤
Thread 2 ─┼→ 同一个 output
Thread 3 ─┤
... ┘
即使原子硬件很强,同一个地址上的更新仍然必须以某种顺序提交。
假设有一百万个输入:
就需要执行一百万次全局原子更新。
这会产生严重的热点竞争。
原子版本什么时候可能合理
对于很小的数据量,原子版本可能并不差,因为它:
- 代码简单;
- 只需要一次 Kernel;
- 不需要 Shared Memory;
- 不需要第二级归约;
- Kernel 启动开销较少。
现代 GPU 的原子操作性能也比早期架构强得多。
因此:
小规模数据
或
多个线程更新大量不同地址
时,Atomic 可能合理。
真正的问题是:
大量线程
+
全部竞争同一个地址
五、基本优化思想:先在 Block 内归约
不要让每个线程都执行全局原子操作。
可以先让一个 Block 内的线程共同求出部分和:
Block 0 → partial_sum[0]
Block 1 → partial_sum[1]
Block 2 → partial_sum[2]
...
然后再对这些部分和进行第二级归约。
假设原来有一百万次全局原子操作,使用256线程的 Block 后,可能只剩约:
个 Block 部分和。
全局写入或原子竞争次数大幅减少。
六、为什么使用 Shared Memory
同一个 Block 的线程需要交换部分和。
线程寄存器默认相互不可见,因此可以使用 Shared Memory:
线程寄存器中的输入值
↓
写入 Shared Memory
↓
Block 内树形归约
↓
Thread 0 得到 Block 部分和
基础代码结构是:
extern __shared__ float sdata[];
int tid = threadIdx.x;
int i = blockIdx.x * blockDim.x + tid;
sdata[tid] = (i < n) ? input[i] : 0.0f;
__syncthreads();
// 在 sdata 中归约
if (tid == 0) {
partial[blockIdx.x] = sdata[0];
}
七、为什么加载后必须同步
每个线程只负责写入:
sdata[tid]
但归约过程中会读取其他线程写入的元素。
例如 Thread 0 可能读取:
sdata[0]
sdata[1]
如果 Thread 1 尚未完成写入,结果就不正确。
所以加载后必须:
__syncthreads();
保证:
所有线程完成 Shared Memory 写入
↓
所有线程才能进入归约阶段
八、低效版本一:交错寻址归约
一种直接写法是:
for (int stride = 1;
stride < blockDim.x;
stride *= 2) {
if (tid % (2 * stride) == 0) {
sdata[tid] += sdata[tid + stride];
}
__syncthreads();
}
假设一个 Block 有8个线程。
第一轮:
Thread 0:s[0] += s[1]
Thread 2:s[2] += s[3]
Thread 4:s[4] += s[5]
Thread 6:s[6] += s[7]
第二轮:
Thread 0:s[0] += s[2]
Thread 4:s[4] += s[6]
第三轮:
Thread 0:s[0] += s[4]
功能正确,但效率较低。
九、交错寻址的第一个问题:分支发散
第一轮中,同一 Warp 内:
偶数线程执行加法
奇数线程不执行
Active Mask 近似为:
1010101010101010...
第二轮:
1000100010001000...
随着归约进行,活跃线程越来越稀疏。
虽然部分线程被屏蔽,但 Warp 仍然作为整体执行指令。
因此产生明显的 SIMT 利用率下降。
十、交错寻址的第二个问题:取模运算
条件:
tid % (2 * stride) == 0
可能需要额外整数运算。
如果除数是2的幂,编译器通常可以转换成位运算,但条件和地址控制仍然比顺序寻址复杂。
更重要的是,访问步长不断增长,可能产生不理想的 Shared Memory Bank 映射。
十一、Shared Memory Bank Conflict
Shared Memory 被划分成多个 Bank。
同一 Warp 中,如果多个线程访问同一个 Bank 的不同地址,访问可能被拆成多次。
交错访问模式中,活跃线程地址间隔不断增大:
0, 2, 4, 6...
0, 4, 8, 12...
0, 8, 16, 24...
某些阶段可能使多个请求映射到相同 Bank,产生 Bank Conflict。
需要注意:
- 具体冲突程度依架构、数据类型和访问宽度而定;
- 所有线程读取同一地址时通常可以广播,不等同于普通冲突;
- Bank 数量和组织不能脱离具体 GPU 判断。
十二、优化版本二:顺序寻址归约
更好的方式是让前半部分线程持续工作:
for (int stride = blockDim.x / 2;
stride > 0;
stride >>= 1) {
if (tid < stride) {
sdata[tid] += sdata[tid + stride];
}
__syncthreads();
}
假设 Block 有8个线程。
第一轮:
Thread 0:s[0] += s[4]
Thread 1:s[1] += s[5]
Thread 2:s[2] += s[6]
Thread 3:s[3] += s[7]
第二轮:
Thread 0:s[0] += s[2]
Thread 1:s[1] += s[3]
第三轮:
Thread 0:s[0] += s[1]
顺序寻址的优势
同一轮中,活跃线程连续:
第一轮:Thread 0~3
第二轮:Thread 0~1
第三轮:Thread 0
相比交错模式:
- 分支掩码更规则;
- Shared Memory 地址连续;
- 更容易避免严重 Bank Conflict;
- 地址计算更简单;
- 编译器更容易优化。
十三、为什么每一轮都需要同步
第一轮后:
sdata[0]
已经被 Thread 0 更新。
第二轮中其他线程可能需要读取第一轮的结果。
必须保证:
本轮所有加法完成
↓
下一轮才能开始
因此每轮结束后需要:
__syncthreads();
但这里有一个性能问题:
归约后半段只剩很少线程工作,整个 Block 仍在反复执行 Barrier。
这正是下一步优化的方向。
十四、危险写法:条件 Barrier
不能写成:
if (tid < stride) {
sdata[tid] += sdata[tid + stride];
__syncthreads();
}
因为只有部分线程进入 Barrier。
__syncthreads() 要求 Block 中所有尚未退出的线程以一致方式到达同步点。
否则可能导致:
- 死锁;
- 未定义行为;
- 部分线程永久等待。
正确形式是:
if (tid < stride) {
...
}
__syncthreads();
十五、优化版本三:每个线程先加载两个元素
前面的实现中:
一个线程读取一个元素
可以改为:
一个线程读取两个元素
并先在寄存器中相加
例如:
int i = blockIdx.x * blockDim.x * 2
+ threadIdx.x;
float value = 0.0f;
if (i < n) {
value += input[i];
}
if (i + blockDim.x < n) {
value += input[i + blockDim.x];
}
sdata[threadIdx.x] = value;
这样做有什么好处
原来一个256线程 Block 只处理256个元素。
现在处理最多:
个元素。
好处包括:
- Block 数量约减半;
- Shared Memory 归约元素数不增加;
- 全局部分和数量减少;
- 每次加载后的计算更多;
- Kernel 管理和同步开销降低。
Global Memory 访问是否合并
第一批读取:
Thread 0 → input[base+0]
Thread 1 → input[base+1]
...
连续。
第二批读取:
Thread 0 → input[base+BLOCK]
Thread 1 → input[base+BLOCK+1]
...
也连续。
因此两组读取都可以形成合并访存。
十六、进一步优化:Grid-Stride Loop
如果数据量很大,不一定让每个线程只处理两个元素。
可以让每个线程循环处理多个位置:
size_t i =
blockIdx.x * blockDim.x * 2
+ threadIdx.x;
size_t grid_stride =
gridDim.x * blockDim.x * 2;
float sum = 0.0f;
while (i < n) {
sum += input[i];
if (i + blockDim.x < n) {
sum += input[i + blockDim.x];
}
i += grid_stride;
}
每个线程先在寄存器中形成自己的局部和,再进行 Block 内归约。
Grid-Stride Loop 的优势
1. 限制 Block 数量
不必为每小段数据都创建一个 Block。
2. 增加每线程工作量
减少部分和数量和全局写回。
3. 提高数据复用和指令效率
线程在循环中连续完成多个加载和加法。
4. 适合 Persistent 风格
可以让每个 SM 上驻留有限数量的 Block,持续处理整个数组。
Grid-Stride Loop 的权衡
如果每线程处理太多元素:
- 依赖链可能变长;
- 可并行 Block 数量减少;
- 单个 Block 执行时间变长;
- 负载不均衡可能增加。
所以需要根据数据规模和 GPU 数量选择 gridDim。
十七、为什么归约通常是带宽受限
对 FP32 求和,每读取一个元素:
- 读取4字节;
- 执行一次加法。
计算强度近似:
最终部分和写入所占比例通常较小。
这个算术强度很低,因此大规模 Reduction 通常属于:
Global Memory 带宽受限
优化目标主要是:
- 让访存连续;
- 减少多余读写;
- 减少同步;
- 减少部分和数量;
- 尽量接近显存带宽上限。
并不是把加法器跑满。
十八、最后一个 Warp 是否还需要 __syncthreads()
当归约只剩32个线程时,所有活跃线程都位于同一个 Warp。
传统代码有时会直接依赖 Warp 的同步执行,去掉 Barrier。
但在现代 CUDA 编程中,不应随意依赖隐式 Lockstep,尤其是线程之间通过 Shared Memory 通信时。
更安全、更高效的方法是:
使用 Warp Shuffle 在寄存器之间直接交换数据。
十九、Warp Shuffle 是什么
Warp Shuffle 允许 Warp 内线程直接交换寄存器值。
例如:
value += __shfl_down_sync(
0xffffffff,
value,
16);
表示每个 Lane 从编号大16的 Lane 取得 value。
假设 Warp 中有32个 Lane:
第一轮,Offset 16:
Lane 0 += Lane 16
Lane 1 += Lane 17
...
Lane 15 += Lane 31
第二轮,Offset 8:
Lane 0 += Lane 8
...
依次:
16 → 8 → 4 → 2 → 1
最后 Lane 0 获得整个 Warp 的和。
二十、Warp Shuffle Reduction
可以写成一个函数:
__device__ __forceinline__
float warp_reduce_sum(float value)
{
for (int offset = 16;
offset > 0;
offset >>= 1) {
value += __shfl_down_sync(
0xffffffff,
value,
offset);
}
return value;
}
前提是:
- Warp 的32个 Lane 都有效并参加;
- 使用完整掩码
0xffffffff; - 所有参与 Lane 以一致方式调用该函数。
Warp Shuffle 的优势
传统 Shared Memory 交换需要:
写 Shared Memory
同步
读 Shared Memory
执行加法
Warp Shuffle 可以:
直接从其他 Lane 的寄存器获取数据
因此减少:
- Shared Memory Load/Store;
- Bank Conflict;
- Block Barrier;
- 地址计算;
- Shared Memory 容量需求。
二十一、一个 Block 内包含多个 Warp 怎么办
假设一个 Block 有256线程:
可以分两级归约。
第一阶段,每个 Warp 独立归约:
Warp 0 → warp_sum[0]
Warp 1 → warp_sum[1]
...
Warp 7 → warp_sum[7]
每个 Warp 的 Lane 0 把结果写入 Shared Memory:
if (lane == 0) {
warp_sums[warp_id] = value;
}
第二阶段,让第一个 Warp 读取8个部分和,再执行一次 Warp Shuffle。
二十二、现代 Block Reduction 结构
Block 中所有线程
↓
每个线程先计算寄存器局部和
↓
每个 Warp 用 Shuffle 归约
↓
各 Warp 的 Lane 0 写 Shared Memory
↓
一次 __syncthreads()
↓
第一个 Warp 读取所有 Warp 部分和
↓
第一个 Warp 再用 Shuffle 归约
↓
Lane 0 得到整个 Block 的和
整个 Block 通常只需要一次 Shared Memory Barrier。
二十三、一个优化后的 Block Reduction Kernel
下面是一个适合学习的版本:
template<int BLOCK_SIZE>
__global__ void reduce_sum_kernel(
const float* __restrict__ input,
float* __restrict__ partial,
size_t n)
{
static_assert(BLOCK_SIZE % 32 == 0,
"BLOCK_SIZE must be a multiple of 32");
constexpr int WARPS =
BLOCK_SIZE / 32;
__shared__ float warp_sums[WARPS];
int tid = threadIdx.x;
int lane = tid & 31;
int warp_id = tid >> 5;
size_t i =
static_cast<size_t>(blockIdx.x)
* BLOCK_SIZE * 2
+ tid;
size_t stride =
static_cast<size_t>(gridDim.x)
* BLOCK_SIZE * 2;
float sum = 0.0f;
// 每个线程处理多个元素。
while (i < n) {
sum += input[i];
size_t second = i + BLOCK_SIZE;
if (second < n) {
sum += input[second];
}
i += stride;
}
// 第一级:Warp 内寄存器归约。
for (int offset = 16;
offset > 0;
offset >>= 1) {
sum += __shfl_down_sync(
0xffffffff,
sum,
offset);
}
// 每个 Warp 的 Lane 0 写入一个部分和。
if (lane == 0) {
warp_sums[warp_id] = sum;
}
__syncthreads();
// 第二级:第一个 Warp 归约所有 Warp 部分和。
if (warp_id == 0) {
sum = (lane < WARPS)
? warp_sums[lane]
: 0.0f;
for (int offset = 16;
offset > 0;
offset >>= 1) {
sum += __shfl_down_sync(
0xffffffff,
sum,
offset);
}
if (lane == 0) {
partial[blockIdx.x] = sum;
}
}
}
二十四、这个 Kernel 中的数据层次
每个线程:
Global Memory 中读取多个元素
↓
寄存器 sum 中局部累加
每个 Warp:
32个线程的寄存器 sum
↓
Warp Shuffle
↓
一个 Warp 部分和
每个 Block:
多个 Warp 部分和
↓
少量 Shared Memory
↓
第一个 Warp 再归约
↓
一个 Block 部分和
最终:
多个 Block 部分和
↓
第二级归约
↓
最终结果
二十五、为什么仍然需要第二级归约
Block 之间不能使用普通:
__syncthreads();
进行同步。
因此每个 Block 得到自己的结果后:
partial[0]
partial[1]
partial[2]
...
还需要把这些部分和继续归约。
常见方法有三种。
方法一:启动第二个 Kernel
第一次:
N 个输入
↓
B 个部分和
第二次:
B 个部分和
↓
更少的部分和
不断递归,直到只剩一个结果。
优点:
- 结构清晰;
- 没有热点原子;
- 容易控制归约顺序;
- 适合大规模数据。
缺点:
- 需要额外 Kernel 启动;
- 需要临时缓冲区;
- 小规模第二阶段启动开销可能较明显。
方法二:每个 Block 最后执行一次 Atomic Add
Block 内归约完成后:
if (threadIdx.x == 0) {
atomicAdd(output, block_sum);
}
原子次数从:
下降为:
这通常比每线程原子快很多。
优点:
- 只需一个 Kernel;
- 不需要额外部分和数组;
- 实现较简单。
缺点:
- Block 数量很大时仍有竞争;
- 浮点加法顺序不确定;
- 结果可能存在微小的运行间差异。
方法三:Cooperative Grid Synchronization
使用 Cooperative Groups 的 Grid 级同步,可以在一个 Kernel 中执行多级归约。
但需要:
- Cooperative Launch 支持;
- 所有 Block 必须能够同时驻留;
- 资源和启动方式受限制;
- 实现复杂度更高。
普通归约通常优先考虑多 Kernel 或每 Block 一次 Atomic。
二十六、为什么第二级 Kernel 通常很小
假设第一阶段有4096个 Block。
第二阶段输入只有:
个部分和。
第三阶段可能只剩:
个结果。
因此后续归约数据量迅速缩小。
总显存流量仍主要来自第一次完整读取原始数组。
二十七、向量化加载
可以尝试使用:
float4
让一个线程一次读取4个 FP32:
float4 v =
reinterpret_cast<const float4*>(input)[index];
sum += v.x + v.y + v.z + v.w;
一次向量读取逻辑上搬运:
向量化加载的优势
它可能减少:
- Load 指令数量;
- 地址计算指令;
- 循环控制;
- 每个字节对应的指令开销。
但它不会自动减少总数据字节数:
标量读取4个 FP32:16 Bytes
float4 读取:仍是16 Bytes
Reduction 仍然通常是带宽受限。
向量化加载的要求
需要保证:
- 地址满足
float4对齐要求; - 输入长度正确处理;
- 尾部不足4个元素时单独处理;
- 强制类型转换符合对象和别名规则;
- Warp 整体访问仍然连续。
CUDA 分配得到的基础地址通常对齐良好,但加入任意偏移后不一定仍满足16字节对齐。
二十八、向量化并不总是更快
如果标量 Load 已经能被编译器高效组合,或 Kernel 已完全受显存带宽限制,向量化带来的提升可能很小。
向量化还可能:
- 增加同时活跃数据;
- 增加寄存器使用;
- 形成更长的局部加法依赖;
- 增加尾部处理复杂度。
所以必须比较实际 Duration,而不能只看 Load 指令减少。
二十九、使用多个局部累加器提高 ILP
如果一个线程连续执行:
sum += input[i0];
sum += input[i1];
sum += input[i2];
sum += input[i3];
形成依赖链:
sum₁ 依赖 sum₀
sum₂ 依赖 sum₁
sum₃ 依赖 sum₂
可以使用多个累加器:
float sum0 = 0.0f;
float sum1 = 0.0f;
float sum2 = 0.0f;
float sum3 = 0.0f;
sum0 += input[i0];
sum1 += input[i1];
sum2 += input[i2];
sum3 += input[i3];
float sum =
(sum0 + sum1)
+ (sum2 + sum3);
这样多个加法相互独立,可以增加 ILP。
多累加器的代价
- 寄存器使用增加;
- Occupancy 可能下降;
- 最后需要额外合并;
- 浮点求和顺序变化。
对于带宽受限 Reduction,多累加器主要用于隐藏加载和加法延迟,而不是提高峰值加法吞吐。
三十、Block Size 如何选择
常见选择:
128 Threads
256 Threads
512 Threads
需要权衡:
- Warp 数量;
- Shared Memory;
- 寄存器;
- Block 数量;
- Barrier 开销;
- 每 Block 处理数据量;
- 尾部效应。
对于现代 Warp Shuffle 归约,256线程常是合理起点:
Shared Memory 只需要保存8个 Warp 部分和。
但最终最优值必须测试。
三十一、Block Size 不必与输入长度完全匹配
通过边界判断:
if (i < n)
可以处理任意输入长度。
例如:
Block Size 为256,每 Block 最多加载512个元素。
需要两个 Block:
Block 0:处理 0~511
Block 1:处理 512~999,并对越界位置填0
归约中常把无效元素视为加法单位元:
对于最大值归约,则单位元应使用足够小的值,例如:
不能统一填0。
三十二、不同归约操作的单位元
| 操作 | 单位元 |
|---|---|
| 求和 | 0 |
| 乘积 | 1 |
| 最大值 | |
| 最小值 | |
| 逻辑与 | true |
| 逻辑或 | false |
边界线程应加载正确的单位元,否则可能改变结果。
例如最大值归约中,若所有输入都是负数,而越界位置填0,最终会错误地得到0。
三十三、浮点加法不满足严格结合律
数学上:
但浮点数存在舍入,因此通常:
Reduction 优化会改变求和顺序,所以 GPU 结果可能与 CPU 串行结果存在微小差异。
这通常不是程序错误,而是浮点运算顺序变化导致的。
三十四、Atomic Reduction 的结果可能不完全确定
多个 Block 对同一地址执行:
atomicAdd
原子操作保证不会丢失更新,但不同 Block 的提交顺序不固定。
例如可能是:
((a+b)+c)+d
也可能是:
(a+(c+d))+b
所以不同运行间最后几个 Bit 可能不同。
原子性保证的是:
每次更新完整执行
不保证:
所有运行使用相同加法顺序
三十五、树形归约的数值误差
与纯串行累加相比,平衡树形归约通常会把数量级相近的部分和逐层合并。
在许多情况下,树形求和的误差增长比一条超长串行链更好。
但它仍然不是精确求和。
数值要求较高时可以考虑:
- 使用 FP64 累加;
- 使用 Kahan Summation;
- 使用分块补偿求和;
- 对数据按数量级排序;
- 使用更严格的确定性算法。
代价通常是更高计算量、更高带宽或更低吞吐率。
三十六、Kahan Summation 的基本思想
普通求和:
sum += x;
会丢失部分小量。
Kahan 求和维护一个误差补偿项:
float y = x - compensation;
float t = sum + y;
compensation = (t - sum) - y;
sum = t;
它可以提高精度,但每个元素需要更多运算和依赖。
对于原本带宽受限的 Reduction,增加这些计算未必按比例降低性能,但依赖链会更加复杂。
是否使用取决于数值要求。
三十七、使用 FP64 累加的权衡
输入可以是 FP32,但局部累加器使用 FP64:
double sum = 0.0;
sum += static_cast<double>(input[i]);
优点:
- 大幅减小累计舍入误差;
- 对大规模求和更稳定。
缺点:
- 寄存器占用增加;
- FP64 加法吞吐率可能较低;
- 数据转换增加;
- 消费级 GPU 的 FP64 性能通常相对较弱。
如果 Kernel 完全受显存带宽限制,FP64 累加可能仍有可接受性能;但必须实际测量。
三十八、使用 Nsight Compute 比较各版本
可以逐步建立以下版本:
V0:每线程 Global Atomic
V1:Shared Memory 交错归约
V2:Shared Memory 顺序归约
V3:每线程加载两个元素
V4:Warp Shuffle
V5:Grid-Stride Loop
V6:向量化加载和多累加器
每一步都比较:
Kernel Duration
DRAM Throughput
L1/L2 Throughput
Atomic 指令和竞争
Shared Memory Wavefronts
Bank Conflict
Barrier Stall
Long/Short Scoreboard
Eligible Warps
Registers Per Thread
Occupancy
三十九、预期的性能变化
V0:每线程原子
可能看到:
原子相关吞吐受限
大量请求竞争同一地址
Kernel Duration 很长
V1:交错 Shared Memory
可能看到:
全局 Atomic 大幅减少
但分支效率较低
Shared Memory 冲突或 Wavefront 较多
Barrier 较多
V2:顺序寻址
可能看到:
Branch Efficiency 改善
Shared Memory 访问更规则
Duration 下降
V3:每线程加载两个元素
可能看到:
Block 数量减少
部分和数量减少
Barrier 和写回开销下降
DRAM 吞吐率提高
V4:Warp Shuffle
可能看到:
Shared Memory 指令减少
Barrier Stall 降低
Short Scoreboard 下降
V5/V6:Grid-Stride、向量化、多累加器
可能看到:
内存指令效率改善
地址和循环指令减少
DRAM 带宽接近硬件上限
达到这一阶段后:
Compute Throughput 仍然可能很低
Memory Throughput 很高
这是正常的,因为 Reduction 本来就是低计算强度算子。
四十、优化成功的最终表现
一个优秀的大规模 FP32 Reduction Kernel 通常表现为:
Global Memory 访问连续且合并
几乎没有 Local Memory Spill
Shared Memory Bank Conflict 很少
Barrier 数量很少
每线程进行适量局部累加
每 Warp 使用 Shuffle
每 Block 只写一个部分和
DRAM 带宽接近可达到上限
并不要求:
FP32 Core 达到100%
Occupancy 达到100%
因为它主要是内存流式算子。
四十一、为什么直接使用 CUB 往往更合理
生产代码中,通常不建议从头手写所有 Reduction。
CUDA 生态中常使用 CUB:
cub::DeviceReduce::Sum(...)
或者 Block 级:
cub::BlockReduce<T, BLOCK_SIZE>
这些库会针对不同架构和数据规模选择优化策略。
优势包括:
- 正确处理边界;
- 使用高效 Warp 原语;
- 支持多种数据类型;
- 支持多种归约操作;
- 随 CUDA 工具链持续优化。
自己实现 Reduction 的主要价值是学习架构和处理特殊融合场景。
四十二、自定义 Reduction 仍然有价值的场景
例如要同时完成:
读取输入
→ 预处理
→ 求和
→ 求平方和
→ 求最大值
如果分别调用多个库 Kernel,需要反复读取输入。
可以在一个 Kernel 中融合:
sum += x;
sum_sq += x * x;
max_v = max(max_v, x);
最终一次读取产生多个统计量。
这样提高计算强度,减少显存流量。
这种融合常见于:
- LayerNorm;
- BatchNorm;
- Softmax;
- 损失统计;
- 特征归一化;
- 在线统计量计算。
四十三、Reduction 与 Softmax 的关系
Softmax 为:
为防止溢出,通常先计算最大值:
再计算:
最后:
其中包含两次 Reduction:
第一次:Max Reduction
第二次:Sum Reduction
高性能 Softmax 的核心就是:
- Warp/Block 归约;
- 减少重复读取;
- 寄存器保存局部数据;
- 融合指数和归一化;
- 控制数值稳定性。
所以掌握 Reduction,是理解 Transformer 算子的基础。
四十四、Reduction 的完整优化路径
可以总结为:
每线程直接 Global Atomic
↓
Block 内使用 Shared Memory
↓
消除交错寻址
↓
顺序树形归约
↓
每线程先处理多个元素
↓
寄存器局部累加
↓
Warp Shuffle
↓
只保存 Warp 部分和
↓
Block 只输出一个结果
↓
多级 Kernel 或每 Block 一次 Atomic
↓
向量化读取、Grid-Stride 和算子融合
本课核心结论
第一,Reduction 的困难在于多个线程最终必须合并成少量结果。
第二,每线程对同一地址执行 Global Atomic 会产生严重热点竞争。
第三,更高效的方法是先在 Block 内归约,只输出一个 Block 部分和。
第四,Shared Memory 允许同一 Block 的线程交换中间结果。
第五,交错寻址会产生不规则活跃掩码、额外地址控制和潜在 Bank Conflict。
第六,顺序寻址让连续线程工作,访问模式更规则。
第七,每轮 Shared Memory 归约后必须同步,但 __syncthreads() 不能只由部分线程执行。
第八,每线程先加载两个或多个元素,可以减少 Block 数量和部分和数量。
第九,Grid-Stride Loop 允许有限数量的 Block 持续处理大数组。
第十,Warp Shuffle 可以直接交换 Lane 寄存器数据,减少 Shared Memory 和 Barrier。
第十一,一个 Block 可以先做 Warp 内归约,再由第一个 Warp 归约各 Warp 部分和。
第十二,Block 之间通常通过第二级 Kernel、每 Block 一次 Atomic 或 Cooperative Grid 继续归约。
第十三,FP32 求和计算强度约为:
因此大规模 Reduction 通常是内存带宽受限。
第十四,向量化 Load 主要减少指令和地址开销,不会减少总读取字节数。
第十五,多个独立累加器可以提高 ILP,但会增加寄存器使用。
第十六,浮点加法不满足严格结合律,不同归约顺序会产生微小数值差异。
第十七,Atomic 保证更新不丢失,但不保证浮点加法顺序固定。
第十八,生产代码通常优先使用 CUB 等成熟库,自定义实现主要用于算子融合和特殊数据流。
下一课将继续分析 Transformer 中最典型的归约算子:
第二十二课:Softmax 为什么难以高效,以及 FlashAttention 如何减少显存访问
主要包括:
数值稳定 Softmax
Max Reduction 与 Sum Reduction
Warp Softmax
Block Softmax
为什么朴素 Attention 要保存 N×N 矩阵
QKᵀ、Softmax 和 PV 的数据流
Online Softmax
Tiling
FlashAttention 的 IO-Aware 思想
为什么它不是近似算法