FlashAttention 1-4 演进速查

每代优化了什么、目标硬件、关键瓶颈转移一览

总览对比

维度	FlashAttention-1	FlashAttention-2	FlashAttention-3	FlashAttention-4
年份 / 论文	2022 (arXiv:2205.14135)	2023 (arXiv:2307.08691)	2024 (arXiv:2407.08608)	2026 (arXiv:2603.05451)
目标 GPU	Volta / Turing / Ampere	Ampere (A100, A6000)	Hopper (H100, H200)	Blackwell (B200, GB200)
核心主题	IO-aware tiling + online softmax + recomputation	更好并行度与工作划分	异步执行 + 低精度	算法与 kernel pipeline 协同设计
主要瓶颈	HBM 读写量（S/P 矩阵 materialization）	非 matmul 开销、warp 空闲、循环顺序	Tensor Core 利用率低、softmax bubble、同步等待	非 matmul 单元（exp/SFU）跟不上 Tensor Core；共享内存带宽不对称
关键硬件特性	Shared memory, Tensor Core	Shared memory, Tensor Core	TMA, WGMMA, async MMA, FP8	TMEM, 2-CTA MMA, fully async MMA
峰值性能（代表性数字）	—（基准）	~225 TFLOP/s (A100, 50–73%)	740–840 TFLOP/s (H100, 75–85%); FP8 ~1.2 PFLOP/s	~1,613 TFLOP/s (B200, 71%)
相对上一代加速	—	~2×	1.5–2×	1.3× vs cuDNN 9.13; 2.7× vs Triton
精度	FP16 / BF16	FP16 / BF16	FP16/BF16 + FP8（block quant + incoherent processing）	FP16 / BF16（论文聚焦 BF16）

每代关键优化

FlashAttention-1：奠基

• Tiling：把 Q、K、V 分成 tile，每次只在 SRAM 里处理一个小块。
• Online Softmax：用 running max m 和 running sum-exp l 实现增量 softmax，避免全局归约。
• Kernel Fusion：把 QK^T → softmax → PV 融合在一个 kernel 里，避免写回 S、P。
• Recomputation：反向传播只保存 m、l 统计量，重新计算 S、P，显存 O(N)。

工程影响：把 attention 从 memory-bound 变成 compute-bound，开启长上下文训练可能。

FlashAttention-2：并行与工作划分

• Loop order 交换：外层循环遍历 query tile，内层循环遍历 KV tile。这样 K、V 每个 tile 只加载一次，且 query 维度天然可并行。
• Sequence-length parallelism：一个 batch × head 内的不同 query tile 可分配到不同 warps/blocks，提高 occupancy。
• 减少非 matmul 开销：更少的 shared memory read/write，更好的 warp-level work partition。
• 去掉 online softmax 的部分同步：每个 query tile 独立完成，降低 block 间同步。

工程影响：A100 上达到 50–73% 峰值算力；成为工业界默认 attention 实现。

FlashAttention-3：Hopper 异步化

• Warp Specialization：producer warps 用 TMA 异步 load，consumer warps 用 WGMMA 异步 MMA，隐藏内存延迟。
• GEMM–Softmax 交错：一个 warpgroup 算当前 tile 的 MMA 时，另一个 warpgroup 做上一个 tile 的 exp/rescale，消除 softmax bubble。
• FP8 + Block Quantization + Incoherent Processing：在保持数值稳定的前提下利用 Hopper FP8 双倍算力。
• Cluster / dynamic register：利用 warpgroup cluster 和 setmaxnreg 动态分配寄存器。

工程影响：H100 上 Tensor Core 利用率从 FA-2 的 ~35% 提升到 75–85%；FP8 推理接近 1.2 PFLOP/s。

FlashAttention-4：Blackwell 重架构

• Tensor Memory (TMEM)：MMA 结果直接写到 TMEM，减少寄存器压力，允许更大 tile。
• Fully asynchronous MMA pipeline：更深度地重叠 copy、MMA、softmax。
• Software-emulated exp + conditional softmax rescaling：缓解 exp/SFU 等非 matmul 单元瓶颈。
• 2-CTA MMA（backward）：backward 中两个 CTA 协作，减少 shared memory 流量和全局 atomic add。
• LPT Scheduling：对 causal / 变长序列按最长处理时间优先调度，改善负载均衡和 L2 行为。
• CuTe-DSL in Python：用 Python DSL 编写 kernel，编译速度比 C++ 模板快 20–30×，同时保留 PTX 级控制。

工程影响：B200 上 71% 峰值算力，超越 cuDNN 9.13 和 Triton；把非对称硬件缩放作为核心设计约束。

FA-4 前向 roofline（M=N=d=128，每 SM 每时钟周期）

资源	吞吐（ops/clock/SM）	tile 周期数	是否瓶颈
MMA（BF16）	8192（Hopper 4096 的 2×）	1024	—
exp 单元（MUFU）	16（不变）	1024	前向瓶颈
共享内存读	128 B（不变）	768 / 3328（后向）	后向瓶颈

FA-4 三招对应三个子瓶颈：① 软仿真 exp（FMA 多项式与 MUFU 并行）+ 条件 rescale 抬 exp 有效吞吐；② TMEM 累加器 + 2-CTA MMA 减 smem 流量；③ LPT 调度改善负载均衡。

瓶颈转移的直觉