CUDA 算子工程：手写 FlashAttention v2 之路

第 21 章性能陷阱与反模式

作者杨艺韬 · 3,132 字 · 发布于 2026-05-09 · 更新于 2026-09-25

这一章是一份清单：十二个在 LLM 算子里反复出现的性能陷阱，每一个都配症状、成因和修法。它的用法不是通读，是在 ncu 报告看不懂的时候回来查。

21.1 陷阱 1：寄存器 spill

症状：kernel 性能远低于预期，ncu 显示 local memory throughput 高、Stall Long Scoreboard 高。

原因：每线程可用的物理寄存器最多 255 个；超过之后编译器把多余的变量放进 local memory——名字叫 local，实际是每线程私有的一块全局内存。它会经过 L1/L2，所以未必每次都打到 HBM，但寄存器操作数是随指令直接读出的，而 local memory 哪怕命中 L1 也要几十个周期（Luo et al. 2024 在 H800 上实测 L1 约 41 cycles），一旦 miss 到 HBM 就是几百个周期（同文实测约 479 cycles）。更糟的是 spill 会同时吃掉 L1 容量，把本来该留在 L1 的数据挤出去。

诊断：

nvcc -arch=sm_90a -Xptxas -v -c my_kernel.cu
# 输出（节选，nvcc 13.4 编译一个 __launch_bounds__(256, 4) 的 kernel 所得，不同 CUDA 版本格式与数值可能不同）:
# ptxas info    : Function properties for _Z9lb_kernelPKfPfi
#     1040 bytes stack frame, 1060 bytes spill stores, 2096 bytes spill loads
# ptxas info    : Used 64 registers, used 0 barriers, 1040 bytes cumulative stack size

spill stores/loads != 0 就是问题。

修复：

减少局部变量数量（合并、复用）。
降低 #pragma unroll 程度。
用 __launch_bounds__(256, 4) 提示编译器降低寄存器使用（参数：每 block 最多 256 thread，希望每 SM 至少驻留 4 个 block；编译器据此把每线程寄存器压到 65536 / (256 × 4) = 64 个以内——上面那段输出的 Used 64 registers 就是这么来的。压得太狠反而会制造 spill，要拿 -Xptxas -v 核对）。
把不常用的状态存到 SMEM 而不是寄存器。

21.2 陷阱 2：SMEM Bank Conflict

症状：SMEM 访问慢，ncu 显示 Bank Conflicts > 0。

原因：第 4 章讲过——一个 warp 内多个线程访问同一个 bank 的不同地址。

诊断：ncu 的 metric l1tex__data_bank_conflicts_pipe_lsu_mem_shared_op_ld.sum（store 侧对应 ..._op_st.sum）。

修复：

+1 padding：__shared__ float smem[32][33] 而不是 [32][32]。
Swizzled layout：用 row XOR col 函数。
Vectorized access：用 float4 让每个线程一次读 16 字节（连续 4 个 bank），warp 的请求被硬件分成几个阶段、每阶段 8 个线程处理，连续访问下无冲突，访存指令数也降到 1/4。注意它和 +1 padding 不兼容（padding 会破坏 16 字节对齐）。
TMA + swizzle：Hopper 上让硬件自动 swizzle。

21.3 陷阱 3：Warp Divergence

症状：算力利用率低，ncu 显示 smsp__sass_average_branch_targets_threads_uniform.pct 远低于 100%。

原因：warp 内 32 个线程走不同分支，必须串行执行不同分支。

反例：

// 反例: 数据相关分支
if (arr[tid] > 0) {
    do_something_a();
} else {
    do_something_b();
}

如果 arr 值随机，几乎每个 warp 里都同时有线程走 a、有线程走 b，整个 warp 要把两条路径先后都执行一遍（另一半线程被 mask 掉），耗时约为两条路径之和。

修复：

数据预排序：让相邻线程的数据相同分支。
Mask 化：把分支变成数学（y = mask * a + (1-mask) * b）。两条路径照样都要算，省掉的只是分支本身；很短的分支编译器通常已自动谓词化，不必手改。
Warp-uniform 分支：让分支条件是 warp_id 而不是 thread_id（整 warp 走同一边）。
接受代价：如果分支不可避免，至少让"重分支"是少数（常见路径快）。

21.4 陷阱 4：L2 Thrashing

症状：L2 miss 高，但单看其中任何一份热数据都放得进 L2。

原因：多个数据流互相驱逐对方。一个常见的例子：跑一个 kernel 同时访问 weight（固定）和 KV cache（每请求不同），两者在 L2 上"打架"。

诊断：ncu 的 lts__t_sector_op_read_hit_rate.pct（L2 读命中率），或直接看 lts__t_sectors_op_read_lookup_miss.sum。

修复：

L2 Persistence：把反复读的小块热数据（cos/sin 表、小 embedding、MoE 路由表之类）锁在 L2。GB 级、每步只读一遍的权重锁进去收益有限。
Streaming load：用 __ldcs（PTX ld.global.cs，sm_90a 上编译成 LDG.E.EF）读冷数据。它并不绕过 L2，而是以 evict-first 策略缓存，让这些只用一次的数据最先被逐出，少挤占热数据。
Tile 重新设计：让一次 kernel 内访问的数据集中在 L2 容量内。

// 用 L2 persistence 把一块反复读的热数据锁在 L2
// 先划出 L2 的持久化区（set-aside），否则下面的窗口没有空间可用
cudaDeviceProp prop;
cudaGetDeviceProperties(&prop, 0);
cudaDeviceSetLimit(cudaLimitPersistingL2CacheSize, prop.persistingL2CacheMaxSize);

cudaStreamAttrValue attr = {};
attr.accessPolicyWindow.base_ptr = hot_ptr;
attr.accessPolicyWindow.num_bytes = 32 * 1024 * 1024;  // 32 MB
attr.accessPolicyWindow.hitRatio = 1.0;   // 窗口大于 set-aside 时应调低，否则持久化行互相驱逐
attr.accessPolicyWindow.hitProp = cudaAccessPropertyPersisting;
cudaStreamSetAttribute(stream,
    cudaStreamAttributeAccessPolicyWindow, &attr);

21.5 陷阱 5：滥用全局 atomic

症状：kernel 性能塌陷，看 SASS 有大量 atomic 指令。

原因：全局内存 atomic 要走 L2 的 atomic 单元，至少是一次 L2 往返（第 4 章：H800 实测 L2 约 263 cycles）；更要命的是同一地址上的高竞争会让它退化成串行——一个 warp 里 32 个线程加同一个地址，就是 32 次串行。编译器帮不上忙：float 的 atomicAdd(out, x) 在 sm_90a 上编译成每线程一条 REDG.E.ADD.F32，只有 atomicAdd(p, 1) 这类常量整数加法会被自动改写成 VOTE + POPC 的 warp 聚合（nvcc 13.4 编译所得）。

反例：第 5 章 v0 的 reduce kernel——N 个线程同时 atomic add 到一个全局变量。

修复：分层归约。

Warp 内 reduce（__shfl）。
Block 内 reduce（SMEM）。
Cluster 内 reduce（DSMEM, Hopper）。
最后只用极少 atomic 写最终结果。

警告：半精度 atomic 硬件早就支持（atomicAdd(__half2*) 从 sm_60、atomicAdd(__half*) 从 sm_70 起），问题从来不在硬件而在精度——被累加的次数越多，fp16 的舍入误差累积越明显。所以 LLM 训练里需要 atomic 累加梯度的地方通常用 fp32 atomic，FlashAttention 的 dQ 也是先 atomic 到一块 fp32 缓冲再转回来（第 16 章）。

21.6 陷阱 6：Block Size 错误

症状：算法正确，性能明显低于预期。

原因：

Block size 太大（比如 1024）：每 block SMEM/寄存器消耗高，SM 上 active block 数减少，occupancy 不足。
Block size 太小（比如 32）：每 SM 最多驻留 32 个 block，32 线程的 block 意味着每 SM 最多 32 个 warp，H100 上 occupancy 被硬性封顶在 50%（每 SM 上限 64 warp）；block 内能协作（共享 SMEM、__syncthreads）的线程也太少。

修复：

常用值：256、512、128。先试 256。
针对算子调整：
- LayerNorm/Softmax：256 或 512（一行一 block）。
- GEMM (Tiled)：128 或 256。
- Reduce：128 或 256。
- Attention (FA2 前向)：128（4 warps）。
用 Occupancy Calculator：Nsight Compute 自带的 Occupancy Calculator，或运行时 API cudaOccupancyMaxActiveBlocksPerMultiprocessor，根据每 block 的线程数、SMEM、寄存器使用估算 occupancy。

21.7 陷阱 7：迷信高 Occupancy

症状：把 occupancy 调到 100% 反而变慢。

原因：高 occupancy 让每个线程的寄存器配额变小，可能导致 spill 或减少 ILP。

反直觉事实：FA、cuBLAS GEMM、CUTLASS 高性能 kernel 的 occupancy 普遍不高。以 FA3 前向 hdim 128 为例：一个 CTA 是 3 个 warpgroup（1 个 producer + 2 个 MMA，共 12 个 warp），用 setmaxnreg 把寄存器分成 MMA 每线程 240 个、producer 24 个（flash-attn/hopper/flash_fwd_kernel_sm90.h:82-83），合计 240 × 256 + 24 × 128 = 64512，几乎占满每 SM 65536 个寄存器，于是每 SM 只能驻留 1 个 CTA，occupancy 只有 12 / 64 ≈ 19%。它们的优势不是"warp 多"，而是"每个 warp 干的活多"。

指导原则：

Bandwidth-bound 算子（reduce、softmax）：较高的 occupancy 有用，因为多 warp 帮助掩盖内存延迟。
Compute-bound 算子（GEMM、FA）：occupancy 不是越高越好，关键是 Tensor Core 利用率。

21.8 陷阱 8：Stream 与 Graph 误用

症状：用了 stream 但 GPU 仍然串行。

原因：Stream 之间的依赖没设置好，或者中间夹了默认 stream 上的操作（legacy 默认 stream 与所有用 cudaStreamCreate 建的阻塞型 stream 隐式同步，见附录 A.1.1），或者 host 端单线程派发太慢。

反例：

// 反例: 默认 stream 阻塞所有
kernel1<<<...>>>(...);
cudaMemcpy(...);  // 默认 stream, 阻塞
kernel2<<<...>>>(...);

修复：

用显式 stream：

cudaStream_t s1, s2;
cudaStreamCreate(&s1);
cudaStreamCreate(&s2);
kernel1<<<..., 0, s1>>>(...);
cudaMemcpyAsync(..., s2);  // 不同 stream, 并行（前提是这次拷贝与 kernel1/kernel2 没有数据依赖）
kernel2<<<..., 0, s1>>>(...);
// 注意 cudaMemcpyAsync 只有在 host 侧内存是 page-locked（cudaHostAlloc /
// cudaHostRegister）时才真的异步；传一块普通 malloc 出来的内存，
// 驱动会退化成同步拷贝，看起来"用了 async 却还是串行"。

用 CUDA Graph：把整个推理 forward 录制成一个 graph，每次只 launch graph：

cudaGraph_t graph;
cudaGraphExec_t graphExec;

cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);
// 录制
forward(input, output, stream);
cudaStreamEndCapture(stream, &graph);
cudaGraphInstantiate(&graphExec, graph, /*flags=*/0);  // CUDA 12 起是 3 参数版

// 后续每次推理:
cudaGraphLaunch(graphExec, stream);  // 一次提交所有 kernel

21.9 陷阱 9：浮点精度问题

症状：FP16 / BF16 训练 loss 不收敛或推理输出异常。

原因：

累加用错精度：FP16 累加溢出/精度损失（FP16 最大只有 65504；逐个加 1.0 的 FP16 累加器到 2048 就不再增长），必须 FP32 累加。Tensor Core 不会自动替你选：mma.sync 有 .f16 累加的变体（sm_90a 上编译成 HMMA.16816.F16），cuBLAS 用 CUBLAS_COMPUTE_16F 也是 FP16 累加，累加器类型要显式指定为 FP32。
Softmax 不 safe：忘了减 max，溢出（FP16 下 exp(x) 在 x > 11.09 就溢出成 inf）。
LayerNorm 用朴素方差公式：第 7 章讲过，E[X²] − (E[X])² 会灾难性消除，要用 Welford（或先求均值的两遍算法）。
量化 scale 选错：per-tensor scale 对异常值敏感。

修复：

累加器永远用 FP32（甚至 FP64）。
所有 reduce 类算子都要数值稳定版本。
量化用 per-channel/per-group。

21.10 陷阱 10：误判带宽 vs 算力 bound

症状：花一周优化算法（减少 FLOPs），性能没变化。

原因：kernel 是带宽 bound，FLOPs 不是瓶颈，HBM 流量才是。

正确诊断：

ncu 看 Roofline：点在带宽屋顶下→带宽 bound；在算力屋顶下→算力 bound。
看 Compute Throughput vs Memory Throughput：哪个高哪个 bound。

修复方向：

带宽 bound：减少 HBM 流量（fusion、tile 复用、量化）。
算力 bound：提升 ALU/Tensor Core 利用（更好的 fragment 调度、消除 stall）。

21.11 陷阱 11：忽略 Kernel Launch 开销

症状：单 kernel 性能不错，整体推理慢。

原因：每次 kernel<<<...>>>() 的 host 端开销是微秒量级（具体值随 driver、参数个数、是否 pinned 而变，自己用 nsys 量一次比记住任何数字都靠谱）。单看很小，但 LLM 推理一次 forward 有几百次 launch，而 decoding 阶段每个 kernel 本身可能只跑几十微秒——开销与有效工作同量级时它就是瓶颈。

诊断：nsys timeline 看 kernel 之间的 gap。

修复：

Kernel Fusion：减少 kernel 数（第 8 章）。
CUDA Graph：把多个 kernel 录成 graph 一次提交。
Persistent Kernel：用一个 kernel 处理多个 tile（第 18 章）。

21.12 陷阱 12：忽略 Host 端瓶颈

症状：GPU 利用率（nvidia-smi 看）只有 50%，但 GPU profiler 看每个 kernel 都很快。

原因：CPU 端阻塞——可能是数据加载、预处理、Python overhead。

诊断：nsys timeline 看 CPU thread 是否在某些点上忙。

修复：

Async data loading：dataloader 用多 worker。
Pin memory + prefetch：减少 H2D 拷贝同步。
Compile heavy logic：用 torch.compile 或写 C++ extension。

21.13 一份 LLM Kernel 优化清单

最后给读者一份 LLM kernel 优化时的快速清单：

□ 1. 用 nsys 确认这是热点 kernel (占总时间 > 5%)
□ 2. 用 ncu 看 Roofline 位置, 判断带宽 vs 算力 bound
□ 3. 检查 spill (ptxas info)
□ 4. 检查 bank conflict (ncu)
□ 5. 检查 occupancy (ncu)
□ 6. 检查 cache hit rate (ncu L2)
□ 7. 检查 warp divergence (ncu)
□ 8. 看 SASS 找次优指令
□ 9. 与 cuBLAS / CUTLASS 同尺寸对比, 看差距来自哪
□ 10. 数值精度验证 (与 reference 实现对比)

每次优化前过一遍这个清单，能避免大部分"白忙活"。

21.14 第五篇收官与下一篇

第五篇我们建立了性能调优的工具链与避坑指南：

第 19 章：用 nsys/ncu 找瓶颈。
第 20 章：读 PTX/SASS 看编译器干了什么。
第 21 章：12 个高频性能陷阱。

到这里本专栏的核心内容（第 1-21 章）全部完成。读者已经掌握了：

基础（第一篇 1-4 章）：GPU 范式、Hopper 架构、编程模型、内存层级。
小算子（第二篇 5-9 章）：Reduction、Online Softmax、LayerNorm、Element-wise Fusion、Quantization。
大算子之 GEMM（第三篇 10-13 章）：朴素到 CUTLASS。
大算子之 Attention（第四篇 14-18 章）：FA1 思想到 FA3 SOTA。
性能工程（第五篇 19-21 章）：诊断与避坑。

附录 A、B、C 会补充三个实用主题：CUDA Graph 与 Stream（异步执行模型）、CUDA C++ 与 Triton 的对比（什么时候选哪个）、与 vLLM·Transformer 那两个专栏的衔接路径（让读者知道下一步该读什么）。

附录写完，本专栏就完成了。

本章动手练习：

写一个故意有 register spill 的 kernel（用 __launch_bounds__ 强制低寄存器），用 ncu 看性能差距。

找你自己写过的一个 CUDA kernel，按 21.13 节清单过一遍，记录每一项的状态。

思考：本专栏介绍的 12 个陷阱中，你之前没意识到、但确实经常踩的是哪几个？