FlashAttention-1 的论文标题就点明了方法论：Fast and Memory-Efficient Exact Attention with IO-Awareness。拿到 2–4 倍加速之后，剩下的问题自然是：IO 账本之外还有多少油水？FA-2 和 FA-3 的回答是——还不少，但要动的东西从"算法结构"逐渐变成"硬件流水线"。
FA-2：先算对，再把浪费找回来
FA-2（Tri Dao, 2023）审视 FA-1 kernel 的 profile，发现两类浪费。
循环次序不对。 FA-1 把外层循环放在 key/value 块上、内层循环遍历 query 块：每个 Kⱼ、Vⱼ 块被反复从 HBM 读入，内层结束时还要把 $O$ 的中间状态写回。FA-2 翻过来——外层遍历 query 块 $Q_i$，内层遍历 K/V 块。这样 $Q_i$ 和 $O_i$ 在外层循环期间完全驻留 SRAM，只有 K/V 流式读入。SRAM 内累加器从"每个 K/V 块一份"变成"每个 Q 块一份"，寄存器占用模式对编译器也友好得多。数学上两个循环次序完全等价（都是同一个分块矩阵乘的遍历顺序），纯粹的调度改变，这类"循环交换（loop interchange）换访存"是性能工程里最经典的一手。
非矩阵乘的开销。 在线 softmax 的 max 比较、exp、行和累加、$O$ 的重缩放，这些是非 GEMM 操作，跑在 tensor core 之外的 CUDA core 上。FA-1 里它们能占到相当比例的 wall time。FA-2 做两件优化：把每行的除法（除以 running sum $l$）推迟到内层循环结束，循环体内只做乘加；重缩放 $O \leftarrow O\,e^{m-m^{\text{new}}}$ 也不是每块都做——只有在新块真的刷新了行 max 时才触发，用一次分支省掉大量冗余逐元素乘。softmax 的归一化从"每块一次"摊到"每行一次"。
并行度。 FA-1 在 batch×head×query 块三个维度上切分线程块，长序列、单 head 的场景下（比如 batch=1 的推理 prefill）切不出足够多的 block 喂饱 GPU。FA-2 增加一个维度：把长序列的 K/V 维也切给不同的 warp，各 warp 算出局部 $(m, l, O)$ 后再做一次归并——这本质上是把在线 softmax 从"行内"推广到"行内 + 行外"两级，归并公式和块间重缩放一模一样：
$$ m = \max_a m_a, \qquad l = \sum_a e^{m_a - m}\, l_a, \qquad O = \sum_a e^{m_a - m}\, O_a. $$
结果：FA-2 在 A100 上前向比 FA-1 快约 2 倍，且确认了论文里一个反直觉的结论——在 HBM 带宽打满之前，kernel 的 wall time 竟然是非矩阵乘操作在主导。
FA-3：为 H100 的硬件特性重写
FA-3（2024，配合 H100 / Hopper 架构）面对的是新硬件：第四代 tensor core、TMA（Tensor Memory Accelerator，专门的异步拷贝引擎）、更大的 SRAM（228KB/SM）。改法对应三种硬件能力。
TMA 接管数据搬运。 Hopper 之前，global memory 到 shared memory 的拷贝由计算线程自己发射 load 指令，计算和搬运共用发射端口。TMA 是专门做这个的协处理器：kernel 只写一张描述符（TensorMap，记录形状、步长、块大小），TMA 引擎独立把整块数据搬进 SRAM 并发出完成信号。FA-3 把 Q/K/V 的块加载全部交给 TMA，计算线程零发射负担，这也是 warp specialization 的前提。
Warp specialization：生产者-消费者流水线。 把一个 SM 上的 warp 分组分工：producer warp 只做"等 TMA 完成、发信号"，consumer warp 只做 tensor core 计算和 softmax 更新，两边用异步屏障（barrier）同步。配合 pingpong 调度，两个 consumer warp 组交替工作——一组在等 softmax 的 exp 等非 GEMM 操作时，另一组已经开始下一轮的 GEMM，tensor core 的空窗被填上。FA-2 里那些必须串行的"GEMM → 修正 → GEMM"链条，在 FA-3 里变成两条交错的流水线。
FP8 与低精度通路。 H100 的 tensor core 支持 FP8（E4M3/E5M2），FA-3 给出了 FP8 版本的 attention，需要处理的是量化粒度问题：Q、K 每块量化（tensor-wise 太粗、per-row 在分块下难做），$P$ 和 $V$ 的乘加要在 FP8 输入下保持累加精度，采用"FP8 存、更高精度累加"的混合路径。收益在纯 GEMM 侧接近 2 倍，但端到端被 softmax 修正链和反量化开销吃掉一部分——这正好再次印证 IO 时代结束后的新瓶颈分布。
三代合看：一条收敛的主线
把三代摆在一起，能看清高效 kernel 设计的一般路径：先改算法结构消除数量级级的浪费（FA-1 的 IO），再改调度消除同数量级内的浪费（FA-2 的循环交换与延迟归一化），最后贴着硬件特性消除最后的空窗（FA-3 的 TMA 与异步流水线）。每一步都是先 profile 找到真正的瓶颈，再动刀——和写业务代码找热点是同一套方法论，只是这里的"热点"单位是 warp 周期和字节搬运。
小结
FA-2 的贡献在于把"softmax 修正链"从隐藏的 wall time 大户变成可分析的科目：循环交换减少累加器往返，延迟归一化砍掉冗余除法，warp 归并补上并行度。FA-3 则说明当一个 kernel 已经把 IO 和非 GEMM 都压下去之后，最后的敌人是硬件流水线的空窗，而填空窗的手段是异步化与专精化。对工程实践的启示很朴素：没有一个"终极 kernel"，每一代都是针对当前 profile 的边际改进，先测量再优化这条纪律在纳秒级同样成立。