大模型的上下文从 2K 卷到上百万，注意力的 O(N²) 成了头号公敌。但故事有个反转：让标准注意力慢下来的，其实不是那 N² 次乘加——GPU 干这个飞快——而是把 N×N 的中间矩阵往显存里写、再读出来的那几趟搬运。FlashAttention 系列四代做的都是同一件事：让数据少跑腿。

这篇从「为什么搬数据比算数据贵」讲起，逐代拆解 FA1 到 FA4 各自解决了什么瓶颈、又被什么新瓶颈逼出下一代，最后看一眼它身后的整个生态。

一、先算账：注意力的三次矩阵之旅

标准注意力的公式只有三行：

S = Q · K^T      # N×d 乘 d×N → N×N 打分矩阵
P = softmax(S)   # 逐行归一化  → N×N 概率矩阵
O = P · V        # N×N 乘 N×d → N×d 输出

Q、K、V 都是 N×d（N 是序列长度，d 是每个注意力头的维度，通常 64 或 128）。问题出在前两步：S 和 P 是 N×N 的大矩阵，比输入大得多。以 N = 8192、d = 128、fp16 为例：Q 只有 2 MB，而 S 有 128 MB——放大 64 倍，而且随 N 平方增长。

教科书的实现是三步各干各的：先算 S 整个写进显存；softmax 时再读回来，算完把 P 再写进去；最后读 P 出来乘 V。每一趟都是几百 MB 的显存读写。

二、GPU 的世界观：算得快，搬得慢

为什么几趟读写就能拖垮性能？看 GPU 的内存层级就明白了（数字取自 FA1 论文中的 A100）：

108 × SM 计算核心张量核在这里干活片上 SRAM每 SM 192 KB · 合计 ≈20 MBHBM 显存40–80 GB≈19 TB/s1.5–2 TB/s工作台快 · 小大仓库慢 · 大
A100 的内存层级：SRAM 比 HBM 快一个数量级、小三个数量级，搬运成本决定性能
  • HBM（高带宽显存）：40–80 GB，带宽 1.5–2 TB/s。模型的 Q、K、V 和最终输出都得住在这里。
  • 片上 SRAM：每个流式多处理器（SM）192 KB，A100 有 108 个 SM，合计约 20 MB，聚合带宽估算约 19 TB/s。

打个比方：SM 是干活的工人，SRAM 是手边的工作台，HBM 是园区另一头的大仓库。工人干活再快，若大半时间都耗在往返仓库搬料上，产出也高不了。

于是 GPU 上的操作分两派：

  • 计算受限（compute-bound）：矩阵乘法，FLOPs 多、读写少，能吃满张量核；
  • 内存受限（memory-bound）：softmax、layer norm、逐元素操作，FLOPs 少得可怜，时间全花在读写上。

标准注意力里，softmax 就是典型的内存受限操作，而且它还逼着前后两个矩阵乘把中间结果落到 HBM。kernel 融合（kernel fusion） 是标准解法：把多步操作合成一个 kernel，中间结果只留在 SRAM 里不落地。难就难在 softmax——它需要整行数据才能归一化，而整行 S 根本塞不进小小的 SRAM。

FlashAttention 的全部智慧，就是解决这个「softmax 塞不进 SRAM」的问题。先看一张全景对比：

标准实现：S、P 两次写回 HBM 再读回Q·KᵀsoftmaxP·VOHBM 显存S（N×N）P（N×N）写回 · 读回写回 · 读回写 OFlashAttention：融合成一个 kernel，中间结果不出 SRAMQ 块 × K_j → S_j → 在线 softmax → × V_j → 累加 O逐块循环，全程驻留片上 SRAMHBM 显存读 Q / K / V 块写 O
标准实现让 N×N 中间矩阵两进两出 HBM；FlashAttention 把三步融合进一个 kernel，中间结果不出 SRAM

三、FlashAttention（2022）：IO 感知的精确注意力

FA1（Tri Dao 等人，NeurIPS 2022）的论文标题点明了立场：IO-aware——设计算法时先把「各级内存之间搬多少次数据」算清楚。它做了两件事：分块计算解决「放不下」的问题，重计算解决「存不起」的问题。注意它是精确注意力，输出和标准实现数学等价，不是 Linformer 那类近似方法。

3.1 思想一：分块 + 在线 softmax

先看 softmax 为什么难融合。数值安全的 softmax 需要三遍扫描：

第一遍：m = max(x₁..x_N)              # 求整行最大值，防溢出
第二遍：f_i = e^(x_i - m)，l = Σf_i    # 指数与指数和
第三遍：y_i = f_i / l                  # 归一化

行太长放不进 SRAM，就得反复从 HBM 读。但 2018 年 Milakov 和 Gimelshein 就给出了在线 softmax：维护两个滚动状态——目前的最大值 m、目前的指数和 l，每来一个新数据就更新一次：

m_new = max(m_old, x_i)
l_new = e^(m_old - m_new) · l_old + e^(x_i - m_new)

关键在 rescale（折算）：最大值变大时，把旧的累积和乘上 e^(m_old − m_new) 折算到新基准，最终结果与三遍扫描严格相等。2021 年 Rabe 和 Staats 用它做出了第一个 O(N) 内存的精确注意力（后来进了 xformers），只是 IO 还不够优。

FA1 把这套递推推广到分块矩阵：把 K、V 切成能放进 SRAM 的块，Q 也切块；外循环逐块流入 K、V，在 SRAM 里算出这一块的 S_j，更新滚动状态 m、l，并把部分结果 P_j·V_j 累加进输出 O 的 SRAM 副本——累加前同样用 e^(m_old − m_new) 把旧 O 折算一次。全部块流完后，最后做一次 O ÷ l 归一化。

K 按块流入（外循环）K₁K₂K₃V 按块流入V₁V₂V₃Q 块驻留 SRAM在线更新（SRAM 内）S_j = Q · K_jᵀm ← max(m, rowmax(S_j))l ← l·e^(m_old−m) + Σ e^(S_j−m)O ← O·e^(m_old−m) + P_j·V_j滚动状态m, l（逐行）输出累积 O末步 ÷ l所有中间量只活在 SRAM；全部块流完后 O ÷ l 归一化，写回 HBM
分块 + 在线 softmax：K、V 按块流入，滚动维护每行最大值 m 与指数和 l，输出 O 边算边折算

整个过程中 S、P 从来没离开过 SRAM，HBM 里只有输入和输出。论文给出的 HBM 访问次数对比：

标准实现：      Θ(N·d + N²)
FlashAttention：Θ(N²·d² / M)      # M 是 SRAM 容量

在 d = 64–128、M ≈ 100 KB 的典型配置下，HBM 访问最多减少 9 倍。论文还证明了下界（Proposition 3）：不存在比这更优的精确注意力算法——FA1 在 IO 复杂度上摸到了理论天花板。内存占用则从 O(N²) 降到 O(N)：比精确基线省 20 倍，甚至比近似方法 Linformer 还省 2 倍。

3.2 思想二：反向传播靠重算

训练还要反向传播，标准实现需要从 HBM 读出保存的 S、P 来算梯度——又是 N² 量级的读写。FA1 的做法很反直觉：什么都不存，只保存输出 O 和每行的 (m, l)；反向时在 SRAM 里用 Q、K 把 S、P 重新算一遍。

重算明明增加了 FLOPs，为什么还更快？论文里 GPT-2 medium 的数字（前向 + 反向）是最好的回答：

                FLOPs      HBM 读写      耗时
标准实现        66.6 G     40.3 GB      41.7 ms
FlashAttention  75.2 G      4.4 GB       7.3 ms

FLOPs 多了 13%，HBM 读写少了 9 倍，最终快 5.7 倍——再次印证那条铁律：搬运比计算贵。

3.3 战果

  • BERT-large（seq 512）：端到端训练比 MLPerf 1.1 的速度纪录还快 15%（20.0 → 17.4 分钟）；
  • GPT-2（seq 1K）：相比 HuggingFace 实现最高 3×，相比高度优化的 Megatron-LM 也有 1.7×；small 模型从 9.5 天（HF）压到 2.7 天；
  • 长序列：Long Range Arena 基准 2.4×；注意力算子本身最高 7.6×；省下的内存让模型能吃更长上下文，GPT-2 把上下文拉长后困惑度还降了 0.7。

更重要的是，FA1 几乎立刻成了业界标配：PyTorch 的 SDPA、各大训练框架的注意力后端，都换成了这个思路的实现。

四、FlashAttention-2（2023）：把 A100 榨到 73%

IO 最优不等于硬件利用率最优。FA1 在 A100 上只跑到理论峰值的 25–40%，离矩阵乘 kernel 还差得远。Tri Dao 2023 年单枪匹马写了 FA2，开刀三处。

第一刀：非矩阵乘 FLOPs。 A100 上 FP16 矩阵乘峰值 312 TFLOPS，而非矩阵乘的 FP32 运算只有 19.5 TFLOPS——同样一个 FLOP，前者比后者便宜 16 倍。FA1 的在线 softmax 每轮循环都要做 rescale 和除法，全是贵的那种。FA2 把归一化推迟到循环结束才做一次（循环里只累加未归一的 Õ 和 ℓ）；反向传播也只存一个 logsumexp（L = m + log ℓ），不再分存 m 和 ℓ 两个量。

第二刀：序列维并行。 FA1 只按 batch × 头数分配线程块，长序列、小 batch 时 SM 大量闲置。FA2 前向按 Q 的行块划分——各块完全独立、互不通信；反向按 K/V 的列块划分，列块之间唯一共享的 dQ 更新用原子加（atomic add）解决。

第三刀：warp 分工。 FA1 在一个线程块内把 K、V 切给 4 个 warp（split-K）：每个 warp 算出部分结果后要写共享内存、同步、再归约。FA2 改成把 Q 切给 4 个 warp、K/V 全体共享：每个 warp 独立算完自己那几行输出，warp 之间零通信。

FA1：按 K 切分（split-K）w0w1w2w3行 = Q 块，列 = K 块（每个 warp 负责一列）4 份部分和 → 写 shared memory同步后再归约求和FA2：按 Q 切分（split-Q）w0w1w2w3每个 warp 负责一行块；K / V 全体共享每个 warp 独立算完一整行块直接写出 O，warp 间零通信
warp 分工对比：FA1 按 K 切分要写共享内存再归约；FA2 按 Q 切分，warp 间零通信

顺手的一个优化：因果掩码下 N×N 矩阵近半数的块全零，直接整块跳过，又快 1.7–1.8×。

效果：A100 上前向最高 230 TFLOPs/s，达到理论峰值的 73%（反向 63%），基本追平高度优化的矩阵乘；注意力本身比 FA1 快 1.7–3.0×，比 PyTorch 标准实现快 3–10×。端到端 GPT 训练（8×A100）跑到每卡 225 TFLOPs/s，即 72% 的 MFU——比 FA1 快 1.3×，比没有 FA 的基线快 2.8×。

五、FlashAttention-3（2024）：为 Hopper 重写

2022 年底 H100 发布，FA2 上新硬件一跑：利用率只有 35%（约 335 TFLOPs/s）。不是 FA2 退步了，而是 Hopper 的玩法变了——硬件给了新玩具，不用就等于浪费：

  • TMA：异步数据搬运引擎，显存到共享内存的拷贝不占计算单元；
  • WGMMA：异步的 warpgroup 级矩阵乘指令，张量核更猛（FP16 峰值约 989 TFLOPS）；
  • FP8 张量核：峰值约 1979 TFLOPS，是 FP16 的两倍。

FA3（Jay Shah、Tri Dao 等，2024）对应三招：

第一招：warp 特化（warp specialization）。 线程块里的 warp 分成两班：生产者 warp 专职发 TMA 指令搬 K、V 块，消费者 warpgroup 专职算 GEMM 和 softmax。搬运和计算两条流水线并行跑，SM 不再空等数据。

第二招：把 softmax 藏进 GEMM 的阴影里。 softmax 全是非矩阵乘运算，在 Hopper 上格外拖后腿。FA3 在块与块之间做流水：算第 j+1 块的 GEMM 时，同时做第 j 块的 softmax；再进一步用两个消费者 warpgroup 打乒乓——一个在做 GEMM 时另一个在做 softmax，张量核与非张量核单元都不闲着。

生产者 warp消费者 WG0消费者 WG1TMA K₁V₁TMA K₂V₂TMA K₃V₃TMA K₄V₄TMA K₅V₅GEMMsoftmaxGEMMsoftmaxGEMMsoftmaxGEMMsoftmaxGEMMsoftmax一方的 softmax 与另一方的 GEMM 同时进行■ GEMM（张量核）　□ softmax（非矩阵乘）　┄ TMA 异步搬运时间
FA3 的乒乓调度：生产者 warp 用 TMA 异步搬数据，两个消费者 warpgroup 的 GEMM 与 softmax 互相掩护

第三招：FP8 低精度。 直接上 FP8 误差太大（打分矩阵 S 的离群值多）。FA3 用了两手：块量化（每块一个 scale，贴合局部范围）加非相干处理（incoherent processing，给数据乘一个随机 ±1 翻转式的正交变换，把离群值「打散」摊平）。最终数值误差比朴素 FP8 注意力低 2.6×。

成绩：H100 上 FP16 最高 740 TFLOPs/s，75% 利用率，比 FA2 快 1.5–2×；FP8 冲到接近 1.2 PFLOPs/s。

六、FlashAttention-4（2026）：不对称扩展时代

故事到 Blackwell 又反转了。B200 的 FP16/BF16 张量核冲到 2.25 PFLOPS（H100 约 1 PFLOPS，翻倍还多），但共享内存带宽和指数单元（MUFU，每 SM 每时钟只有 16 次）几乎原地踏步。FA4 论文给这个现象起了名字：不对称硬件扩展（asymmetric hardware scaling）。后果是瓶颈反转：以一个 128×128×128 的 tile 为例，MMA 要 1024 个时钟周期，光指数运算也要 1024 个周期；共享内存流量加指数运算的总开销超出 MMA 25–60%。非矩阵乘从「可以藏在阴影里」变成了主角。

FA4（Ted Zadouri、Markus Hoehnerbach、Jay Shah 共同一作，Tri Dao 等，Princeton / Meta / Colfax / NVIDIA / Georgia Tech / Together AI，2026 年 3 月挂出）的应对是算法与流水线的协同再设计：

其一，条件重缩放。 回看在线 softmax 的 rescale：只有新块的最大值超过旧基准才真正需要折算。FA4 再放宽一步——m 的增量不超过阈值 τ = log₂(256) = 8 就跳过。跳过的代价是中间累积值最多偏离正确缩放 256 倍，FP32 的巨大动态范围完全兜得住；而 m、l 始终在正确追踪，最后用真最大值统一归一，最终结果严格不变。一旦 running max 已经够大，后续块很少再超出它 8 个单位，绝大多数迭代的 rescale 就此省掉。实现上按 warp 粒度统一判断，避免分支分歧。

其二，软件仿真指数函数。 MUFU 只有 16 ops/clk/SM，那就让一部分 exp 改到 FMA 单元上用多项式软算：Cody-Waite 约简把 2^x 拆成整数幂（直接操作 IEEE 754 的指数域，一次移位加）乘小数幂（Sollya 软件求出的最优多项式，Horner 法求值）。FMA 单元和 MUFU 本来就能并行，两条硬件路径分摊指数运算。但全部仿真会撑爆寄存器反而掉速，所以只对每行 10–25% 的元素仿真，比例按 MMA 与指数的吞吐比调优。精度上，三次多项式的最大相对误差 8.77×10⁻⁵，看似比硬件差几百倍，但 BF16 本身的量化误差（约 3.9×10⁻³）远大于它——softmax 输出以 BF16 消费，99% 的输入上与硬件结果相差不超过 1 个 ULP。

其三，新流水线与反向重构。 前向用 128×128 大 tile 和全异步 MMA，乒乓调度中把输出重缩放拆给独立的 correction warpgroup；反向利用 Blackwell 的 TMEM 和 2-CTA MMA，把共享内存流量从 3328 周期压到 2688 周期（从超出 MMA 30% 降到 5%），dQ 的全局原子加次数减半。还提供了确定性反向模式（信号量锁排序），性能约为非确定性版本的 75%。调度器按「最长处理时间优先」（LPT）分配任务块，在 H200 上也白捡 4–14%。

其四，工程上的大解放：CuTe-DSL。 FA3 是 C++/CUDA 手写，编译一次前向 55 秒、反向 45 秒；FA4 改用 Python 的 CuTe-DSL 实现，编译只要 2.5 秒 / 1.4 秒——快 20–30 倍，改 kernel 的迭代成本骤降。

成绩：B200 上最高 1613 TFLOPs/s，71% 利用率；比 cuDNN 9.13 的注意力快 1.1–1.3×，比 Triton 实现快 2.1–2.7×。（FA3 不支持 B200，所以没有同门前代对比。）

七、四代一张图

2022202320242026FA1 · IO 感知FA2 · 并行重构FA3 · 异步流水FA4 · 协同设计分块 + 在线 softmax反向重计算内存 O(N²)→O(N)减少非矩阵乘序列维并行warp 按 Q 分工warp 特化 + TMAGEMM/softmax 乒乓FP8 低精度条件重缩放软件仿真 exp2-CTA 反向 · TMEMA100　25–40% 峰值A100　73%230 TFLOPs/sH100　75%740 TFLOPs/sB200　71%1613 TFLOPs/s
FlashAttention 四代演进：瓶颈从 IO 到调度、到异步开销、再到不对称扩展，利用率一路推上 70%
代际年份主战场核心招式峰值利用率
FA12022A100分块 + 在线 softmax、反向重计算25–40%
FA22023A100减非矩阵乘、序列维并行、warp 按 Q 分工73%（230 TF/s）
FA32024H100warp 特化、GEMM/softmax 乒乓、FP875%（740 TF/s）；FP8 ≈1.2 PF/s
FA42026B200条件重缩放、软件仿真 exp、2-CTA 反向71%（1613 TF/s）

围绕主干还有一片生态值得知道：

  • Flash-Decoding：推理 decode 阶段 batch 小、序列长，FA2 的并行度又不够用了；解法是沿 KV 长度再切一刀并行算，最后用 logsumexp 把各段结果归并——还是在线 softmax 那一套；
  • 推理框架：vLLM、SGLang 里的 PagedAttention 管的是 KV cache 的显存布局（另一个维度的问题），注意力计算本身仍跑 FlashAttention 系或 FlashInfer 等 kernel；
  • 训练框架：PyTorch 的 F.scaled_dot_product_attention 的 flash 后端、Megatron、DeepSpeed、HF 的 attn_implementation="flash_attention_2"——今天训练 LLM 几乎默认就在用 FA；
  • 同门师兄弟：xformers 的 memory-efficient attention 源自 Rabe & Staats 的 O(N) 内存注意力；cuDNN 9.x 的图优化注意力内核也是 FA 思路的官方实现，FA4 论文里它就是 Blackwell 上的主要对标。

结语：算法追着硬件跑

四代 FlashAttention 讲的是同一个故事的四章：瓶颈在哪，优化就跟到哪。FA1 发现瓶颈在 IO，把 HBM 访问压到理论下界；FA2 发现瓶颈在并行与调度，把 A100 利用率从 25–40% 推上 73%；FA3 赶上 Hopper 的异步硬件，把 softmax 藏进 GEMM 的阴影；FA4 赶上 Blackwell 的不对称扩展，干脆连指数函数都自己软件仿真，把闲置的 FMA 单元也拉进来干活。

「让算法感知硬件」（IO-aware / hardware-aware）这个朴素的原则，如今早已超出注意力本身，成了大模型系统优化的通用方法论。下一次硬件换代时，FlashAttention 大概还会有第五代——瓶颈搬到哪里，它就追到哪里。

参考资料

  • Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness（NeurIPS 2022）.arXiv:2205.14135
  • Tri Dao. FlashAttention-2: Better Parallelism and Work Partitioning（2023）.arXiv:2307.08691
  • Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision（2024）.arXiv:2407.08608
  • Ted Zadouri, Markus Hoehnerbach, Jay Shah, Timmy Liu, Vijay Thakkar, Tri Dao. FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling（2026）.arXiv:2603.05451
  • Maxim Milakov, Natalia Gimelshein. Online normalizer calculation for softmax（2018）.arXiv:1805.02867
  • Markus N. Rabe, Charles Staats. Self-attention Does Not Need O(n²) Memory（2021）.arXiv:2112.05682
  • Tri Dao 等. Flash-Decoding for long-context inference（2023）.PyTorch Blog
  • Dao-AILab/flash-attention（GitHub 仓库，含 FA2/FA3/FA4 实现）