FlashAttention 三代演进：从 2205 到 Hopper 到 Blackwell

前情说明：这是技术深度 D2。上一篇 D1 挖了 DeepSeek V4 的模型侧稀疏 attention，这一篇挖 kernel 侧——FlashAttention 家族三代是如何一步步榨干 GPU 的。从最基本的 tiling + softmax rescaling，到 Hopper 的 warp specialization，再到 Blackwell 时代的重写。素材来自 Tri Dao 团队的 FA3 官方博客 + FA1/FA2/FA3 三篇 arxiv 论文 + Colfax Research 深度解读。

起手：为什么 attention 是 kernel 层最难的战场

深度学习里，GEMM（矩阵乘法）已经被优化到极致——CUTLASS、cuBLAS、CUDNN 十几年的积累，随便一个矩阵乘都能跑到理论峰值的 90%+。

但 attention 不行。Standard attention 有两个 GEMM + 一个 softmax——

S = Q @ K^T           # GEMM 1
P = softmax(S)        # 非 GEMM，逐行 softmax
O = P @ V             # GEMM 2

问题在中间那个 softmax：

需要存中间矩阵 S 和 P（长度 N × N），一个 8K 序列这就是 256 MB——必须写 HBM 再读回来
Softmax 是逐行归一化（需要看到整行才能算），跟 GEMM 的分块并行冲突
Softmax 的 exp 用的是Multi-Function Unit——H100 上算力只有 3.9 TFLOPS，是 FP16 GEMM 的 1/256！

这三个问题让 attention 长期是内存带宽 bound + special function bound + 并行度受限。

FlashAttention 系列的历史，就是每一代想办法——把 attention 从这些约束里一层层解放出来。

第一代 · FlashAttention (2022) · IO-Aware

核心贡献：通过 tiling + softmax rescaling 避免写中间矩阵到 HBM。

天真做法的问题

Standard attention 每一层要做：

1. 写入 S = Q @ K^T 到 HBM  (N² 大小)
2. 从 HBM 读 S,写入 P = softmax(S) 到 HBM
3. 从 HBM 读 P,写入 O = P @ V 到 HBM

HBM 读写次数：O(N²)——序列越长越致命。

FlashAttention 的三个洞察

洞察 1：Tiling 分块

把 Q、K、V 分成 tile（比如 128×128），一次处理一个 tile pair。整个 attention 变成两层 nested loop：

for K_tile in K_tiles:
    for Q_tile in Q_tiles:
        # 一次只算 128×128 大小的 S_tile
        S_tile = Q_tile @ K_tile.T
        P_tile = softmax(S_tile)
        O_tile += P_tile @ V_tile

关键：中间的 S_tile、P_tile 只在 SRAM 里（几十 KB），不写 HBM。

洞察 2：Online Softmax（softmax rescaling）

问题：softmax 是全局归一化，每个 tile 只看到部分数据怎么算？

解法：在线更新分子分母。每处理一个 tile，用当前 max 值 rescale 之前的累加结果：

prev_max = -inf, prev_sum = 0, prev_O = 0

for each K_tile, V_tile:
    S = Q @ K_tile.T
    new_max = max(prev_max, S.max())
    exp_scale = exp(prev_max - new_max)          # rescale 因子
    P = exp(S - new_max)
    new_sum = prev_sum * exp_scale + P.sum()
    new_O = prev_O * exp_scale + P @ V_tile
    # 更新
    prev_max, prev_sum, prev_O = new_max, new_sum, new_O

O = prev_O / prev_sum

关键：每次遇到更大的 max，把之前的累加结果乘以 exp(旧max - 新max) 修正——数学上完全等价于全局 softmax，但只需要流式处理。

洞察 3：Recompute for Backward

Backward 需要中间矩阵。传统做法是训练时把 P 存下来（HBM 巨大开销）。

FA 做法：不存，backward 时重新算一遍 forward（tile 级）。看起来算力翻倍，但backward 本来就 memory-bound——重算反而比读 HBM 快。

收益

HBM 读写 O(N²) → O(N)（跟着 tile 数量走，不是 tile 内元素数）
内存占用 O(N²) → O(N)（不用存 P）
实际 wallclock 加速 2-4x
原生支持 causal mask、sliding window、dropout 等变体

FA1 发布在 A100 时代（2022），当时的 attention 从"性能瓶颈"变成"可以忽略的开销"——直接催生了长上下文 LLM 的爆发（GPT-4 128K、Llama 3 1M）。

FA1 是 IO-aware 算法设计的教科书案例——算法的胜利不来自更少的算力，来自更聪明地使用内存层级。

第二代 · FlashAttention-2 (2023) · Work Partitioning

核心贡献：重新设计 work partitioning，让 GPU 各级并行度都吃满。

FA1 的三个遗留问题

FA1 在 A100 上能到理论峰值的 50-70%——但 30% 还没吃到。三个原因：

问题 1：Warp 之间不平衡

FA1 沿 K/V 维度做外层循环，每个 warp 处理不同的 K/V tile。但 softmax 需要全局归一化，每个 warp 算完还要跟其他 warp 通信合并——大量 warp 间同步开销。

问题 2：非 GEMM 计算太多

Softmax 里的 exp、max、sum 都是 non-GEMM 操作，Tensor Core 空转。FA1 里 non-GEMM 占了 30-40% 时间。

问题 3：Backward 效率低

FA1 的 backward 用了跟 forward 相同的分块策略，但 backward 的 work pattern 不同，导致利用率低。

FA2 的三个改进

改进 1：Loop order 反转

FA1: outer loop 沿 K/V，inner loop 沿 Q → warp 之间要合并 O 的中间累加

FA2: outer loop 沿 Q，inner loop 沿 K/V → 每个 warp 独立算完自己的 Q tile，不用跟其他 warp 通信

这个改动看似小，实际是根本性的——一个 warp 完整拥有一个 Q tile 的所有计算，warp 间零同步。

改进 2：减少 non-GEMM 操作

FA1 每处理一个 tile 都做一次 rescale。FA2 观察到——rescale 的 exp(diff) 不需要每次都算，可以延迟到最后才算。这样中间循环只做累加，只在最后一次 rescale 一下。

结果：non-GEMM 时间从 30% 压到 15%。

改进 3：更好的 backward 分块

Backward 沿 K/V outer + Q inner，跟 forward 反过来。避开 backward 的 work pattern 瓶颈。

收益

FA2 vs FA1：加速 2x（A100 上从 30% 峰值到 70% 峰值）
训练速度显著提升——Llama 2 训练里 attention 部分不再是瓶颈
成为 PyTorch 官方 scaled_dot_product_attention 的默认后端

FA2 是**"细节里的胜利"**——没有新算法，就是把 FA1 的 work partition 重新想了一遍，2x 就出来了。

但 FA2 在 H100 上只有 35% 利用率——因为它是为 Ampere 设计的，没用上 Hopper 的新硬件特性。这个坑逼出了 FA3。

第三代 · FlashAttention-3 (2024) · Hopper Async

核心贡献：利用 Hopper 的三大新特性——WGMMA、TMA、FP8——把 H100 峰值利用率从 35% 拉到 75%。

Hopper 新特性

特性 1：WGMMA (Warpgroup Matrix Multiply-Accumulate)

Hopper 引入了新的 tensor core 指令 wgmma，比 Ampere 的 mma.sync 快 50%。关键区别：wgmma 是异步的——发起指令后 CPU（其实是 SM 的其他 warp）可以继续做别的事，等结果时再显式 wait。

特性 2：TMA (Tensor Memory Accelerator)

Hopper 专门加了一个硬件单元来做 HBM ↔ shared memory 的数据搬运：

自动算 index、越界检查
异步——发起 TMA 后 warp 可以继续做别的
释放寄存器压力——不用 warp 手动搬数据

特性 3：FP8

Tensor core 的 FP8 吞吐是 FP16 的 2x（989 TFLOPS → 1978 TFLOPS）。

FA3 的三个核心技术

技术 1：Warp Specialization

传统 kernel 里，一个 warp 什么都干——搬数据、算 GEMM、算 softmax。FA3 把不同职责分给不同 warp group：

Warp Group 1 (Producer):  只做 TMA 数据搬运（HBM → SRAM）
Warp Group 2 (Consumer):  只做 WGMMA + softmax 计算

Producer 用 TMA 异步搬数据，Consumer 消费。两组人马各干各的，通过 barrier 同步——就像 pipeline 里的生产者-消费者。

收益：数据搬运和计算完全 overlap，理论峰值从 570 TFLOPS 提到 620 TFLOPS。

技术 2：Pingpong Scheduling — Warp Group 之间 overlap GEMM 和 Softmax

问题：H100 上 GEMM 是 989 TFLOPS，exp 只有 3.9 TFLOPS（差 256 倍）。head dim 128 时，exp 的时间跟 GEMM 的时间差不多——这时候不 overlap 就浪费一半。

Pingpong Scheduling：

时间轴:
Warpgroup 1: [GEMM1][GEMM0'][Softmax0'][GEMM1'']...
Warpgroup 2:     [Softmax1][GEMM1'][Softmax1']...

同色 = 同一次迭代

两个 warpgroup 轮流做 GEMM 和 Softmax——一个在 tensor core 上跑 GEMM，另一个同时在 multi-function unit 上跑 softmax。tensor core 和 MFU 完全并行。

收益：570 → 620 TFLOPS（+9%）。

技术 3：Intra-Warpgroup Overlap — 一个 warp group 内部再 overlap

哪怕在同一个 warp group 里，也能挤时间：

一个 warp group 内部:
  GEMM 输出的一部分先出来 → 立刻开始 softmax
  Softmax 处理时 GEMM 的另一部分还在算

代价：需要更多 register 存中间态。但 register pressure 换来 620 → 660 TFLOPS。

技术 4：Incoherent Processing for FP8

FP8 attention 有个致命问题——LLM activation 有 outlier（少数值极大），量化到 FP8 时误差爆炸。

FA3 借鉴 QuIP 论文的技巧——用 Hadamard transform 打散 outlier：

Q' = Hadamard(Q)     # 打散 Q 的 outlier
K' = Hadamard(K)     # 打散 K 的 outlier
S = Q' @ K'.T        # 数学等价（Hadamard 是正交矩阵）

关键性质：Hadamard 是正交矩阵 → 不改变 attention 的数学结果。

性能技巧：Hadamard transform 是元素级 memory-bound，可以跟 RoPE 融合——for free！

收益：FP8 attention 的量化误差降低 2.6x。

FA3 综合收益

FP16: 740 TFLOPS（H100 峰值 75%）—— 比 FA2 快 1.6-2x
FP8: 接近 1.2 PFLOPS（1200 TFLOPS）—— 误差比 baseline FP8 attention 小 2.6x

FA3 是"硬件深度感知的算法设计"——每一个技术都是为 Hopper 的具体特性量身定做。移植到 Ampere 或 AMD 就跑不出这个数字。

第四代 · FA4 for Blackwell (2025-2026 进行中)

Blackwell (B200/B300) 又带来了新东西：

TMEM (Tensor Memory)：新的片上存储层次，比 SRAM 更大更快
FP4 Tensor Cores：FP8 之外再翻倍
CTA Clusters：多个 CTA 可以直接共享 shared memory，减少 HBM 同步

FA3 的 pingpong 在 Blackwell 上已经不是最优——因为：

CTA cluster 让"数据共享"从 HBM 层下沉到 shared memory 层
TMEM 让中间结果不用回 register 或 shared memory
FP4 让 non-GEMM 的相对开销再次变大（因为 GEMM 又翻倍了）

FA4 目前处于开发中（截止 2026-07-30 尚无正式发布 blog）。已知路线图：

用 CTA cluster 做跨 CTA 的数据共享（替代 warp group 之间的 shared memory）
用 TMEM 存 accumulator，避免 register 溢出
更激进的 FP4 支持（连 KV cache 都 FP4）
针对decode 场景（batch=1，超短 query）的专门优化——因为 decode 时 tiling 策略跟 prefill 完全不同

vLLM 内部已经开始集成 CUTLASS 4.x 提供的 Blackwell attention kernel（FlashInfer for B300 项目在跟进）。

一张表看清三代差异

维度	FA1 (Ampere)	FA2 (Ampere)	FA3 (Hopper)
主要贡献	Tiling + Online Softmax	Work partition 反转	Async + Pingpong
HBM 访问	O(N²) → O(N)	同 FA1	同 FA1，但 TMA 硬件加速
Warp 分工	每 warp 一 K/V tile	每 warp 一 Q tile	Producer / Consumer warp
GEMM/Softmax overlap	无	部分（延迟 rescale）	完全（Pingpong + intra-WG）
Tensor Core	mma.sync (sync)	mma.sync (sync)	wgmma (async)
量化	FP16	FP16	FP16 + FP8 + Incoherent
A100 峰值利用率	~50%	~70%	N/A
H100 峰值利用率	N/A	~35%	~75%
代码复杂度	中	中	高（多 warp group + async）
发布年	2022	2023	2024

FA 系列的三个通用启示

启示 1：算法的胜利来自 IO 建模

FA1 最深的洞察不是"分块"——分块从 CUDA 诞生就有——而是把内存层次显式建模到算法里：

不要问"这个算法需要多少 FLOPS"，问"这个算法在多少字节 HBM 带宽下能跑多快"。

现代 GPU 已经算力过剩，瓶颈永远是内存（HBM 带宽、SRAM 容量、寄存器数量）。IO-aware 算法设计是 AI Infra 领域最关键的思维方式。

这也是 DeepSeek V4 稀疏 attention 的哲学——它降低的不是 FLOPS，是KV cache 内存。

启示 2：硬件-算法 co-design 有代际红利

FA1 在 A100 上是 SOTA。FA2 依然是 A100 SOTA——但 FA1/FA2 到了 H100 只能到 35%。FA3 用 Hopper 特有的 async + WGMMA + TMA 特性，直接翻倍。

每代新硬件都会重新洗牌——旧代 SOTA 到新代可能连 40% 都吃不到。这是硬件厂商的护城河，也是 kernel 工程师永远有活干的原因。

启示 3：Non-GEMM 是攻不下的天花板

不管怎么优化，H100 的 exp 只有 3.9 TFLOPS——比 GEMM 慢 256 倍。这个数字从设计层面就决定了：MFU 单元是共享的，不可能给它加更多硅面积。

Softmax 是 attention 的天然瓶颈，这就是为什么：

Linear attention（省掉 softmax）一直是研究热点
Mamba/SSM 用状态空间模型完全绕过 softmax
DeepSeek V4 用稀疏 attention 减少 softmax 的输入规模

每一个绕开 softmax 的尝试，本质上都是在攻这个 256x 的差距。

附：素材来源

主线

Tri Dao Blog · FlashAttention-3（2024）
PyTorch Blog · FlashAttention-3
Colfax Research · Deep-dive

论文

FA1 · arxiv 2205.14135
FA2 · arxiv 2307.08691
FA3 · arxiv 2407.08608
QuIP · arxiv 2307.13304
QuIP# · arxiv 2402.04396

代码

Dao-AILab/flash-attention
CUTLASS · NVIDIA/cutlass
ThunderKittens · HazyResearch/ThunderKittens

姊妹篇

整理于 2026-07-31 深夜。原文技术细节归属 Tri Dao 团队与相关论文作者。中文表达可能有解读偏差，请以原文为准。

FlashAttention 三代演进：从 2205 到 Hopper 到 Blackwell

FlashAttention 三代演进：从 2205 到 Hopper 到 Blackwell

起手：为什么 attention 是 kernel 层最难的战场

第一代 · FlashAttention (2022) · IO-Aware

天真做法的问题

FlashAttention 的三个洞察

收益

第二代 · FlashAttention-2 (2023) · Work Partitioning

FA1 的三个遗留问题

FA2 的三个改进

收益

第三代 · FlashAttention-3 (2024) · Hopper Async

Hopper 新特性

FA3 的三个核心技术

技术 4：Incoherent Processing for FP8

FA3 综合收益

第四代 · FA4 for Blackwell (2025-2026 进行中)

一张表看清三代差异

FA 系列的三个通用启示

启示 1：算法的胜利来自 IO 建模

启示 2：硬件-算法 co-design 有代际红利

启示 3：Non-GEMM 是攻不下的天花板

附：素材来源

评论