FlashAttention-2峰值性能难复现，因其本质是微架构协同设计：需Warp内流水、Shared Memory分块复用、延迟归一化及外循环次序反转，仅按公式逐行写CUDA会导致同步瓶颈与bank conflict。
为什么直接写 CUDA kernel 很难跑出 FlashAttention-2 的峰值性能
因为 FlashAttention-2 不是“把 softmax 搬到 GPU 上”这么简单，它本质是一套微架构协同设计：Warp 内部流水、Shared Memory 分块复用、指数重缩放延迟归一化、以及外循环次序反转。你如果照着公式 softmax(QK^T)V 逐行翻译成 CUDA，大概率会卡在 __syncthreads() 上——每个 Warp 都在等别的 Warp 算完最大值和求和，结果 70% 时间花在同步停顿里。
常见错误现象包括：
- 显存带宽利用率高，但 Tensor Core 利用率始终低于 40%
- 小 batch 下比原生 PyTorch
torch.nn.functional.scaled_dot_product_attention还慢 - 启用
causal=True后结果错位，或梯度反传崩溃
根本原因在于没对齐 FA2 的数学重构逻辑：它把原本分散在内层循环里的 exp(m_old - m_new) 和 O_i = O_i * alpha + V_j * exp(s_j - m_new) 全部推迟到块级归约终点，靠一个 warp-level reduction + shared memory atomic 加法完成最终归一化。
关键三步：tiling 大小、warp 分配、shared memory 布局必须联动调优
FA2 的吞吐不取决于单个 block 能算多大，而取决于每个 warp 是否能连续喂饱 Tensor Core。A100/H100 上最稳的起点配置是：
立即学习“C++免费学习笔记（深入）”；
-
kBlockM = 128（Q 分块高度），kBlockN = 64（K/V 分块宽度）：适配 FP16/BF16 下 64~128 head_dim 的寄存器压力 - 每个 thread block 启动 4 个 warp，每个 warp 负责 1 个 head：避免跨 head 的 shared memory bank conflict
- shared memory 按
float16对齐预分配：__shared__ half s_q[128][64]、s_k[64][64]、s_v[64][64]—— 注意不是二维数组直接声明，而是用线性偏移 + stride 计算地址，否则 bank conflict 会让带宽跌 40%
如果你硬要改 kBlockN=128，那必须同步把 warp 数从 4 改成 8，并重排 s_k 存储为 column-major；否则第二组 64 列会撞进同一 memory bank。
在线 Softmax 的 warp-level reduction 怎么写才不出错
FA2 的核心是把 softmax 拆成两阶段：第一阶段每个 warp 算自己负责的 tile 内最大值和指数和（local max & sum），第二阶段所有 warp 归约出全局 max & sum。这里最容易漏掉的是：
- 必须用
__shfl_sync()而非__syncthreads()做 warp 内规约：32 个线程在一个 warp 里，__shfl_sync(0xffffffff, val, offset)是零开销的寄存器 shuffle - 全局规约不能直接用
atomicAdd写 global memory：要用 shared memory + 一次__syncthreads()后由 warp 0 汇总，否则 atomic contention 会拖垮 latency - exp 计算必须用
hexpf（half 版本），别用expf强转 —— 在 H100 上后者触发隐式类型提升，多占 1 个 cycle
示例片段（简化）：
// Warp 内求 local_max, local_sum
float local_max = -INFINITY;
float local_sum = 0.f;
#pragma unroll
for (int i = 0; i < 64; ++i) {
float s_val = h2float(qk[i]) * scale; // qk 是 half2
local_max = fmaxf(local_max, s_val);
}
#pragma unroll
for (int i = 0; i < 64; ++i) {
float s_val = h2float(qk[i]) * scale;
float exp_val = h2float(hexpf(__float2half(s_val - local_max)));
local_sum += exp_val;
}
// __shfl_sync 规约 local_max/local_sum 到 lane 0
float max_warp = __shfl_sync(0xffffffff, local_max, 0);
float sum_warp = __shfl_sync(0xffffffff, local_sum, 0);
反向传播里重计算（recomputation）不是可选项，是内存安全前提
FA2 的反向 kernel 不存任何中间 QK^T 或 softmax 输出，而是前向时只缓存 q_ptr、k_ptr、v_ptr、o_ptr 和每个 tile 的 lse（log-sum-exp）。反向时重新跑一遍前向流程，用当前梯度 dO 反推 dQ、dK、dV。这意味着：
- 你不能在 forward kernel 里把
qk结果写回 global memory —— 显存带宽全被浪费，还破坏了 zero-copy 前提 -
lse必须以 float32 存（哪怕输入是 half），否则反向里exp(lse)下溢/上溢导致梯度爆炸 - recomputation 的分块逻辑必须和 forward 完全一致，包括
kBlockM/kBlockN、padding 行为、causal mask 的裁剪位置 —— 差一个 index 就梯度错位
真正难的不是算得快，是 forward 和 backward 的 tile 边界、mask 应用点、scale 因子作用时机，三者必须字节级对齐。一旦 mismatch，dQ 的某几行梯度就永远为 0，模型训着训着就崩。
