FlashAttention：从 IO-Aware Attention 到 GPU Work Partitioning

FlashAttention 并没有把 dense attention 从变成线性复杂度，也没有引入低秩、稀疏或核函数近似；FlashAttention-2 同样没有改变 Attention 的数学定义。两篇论文真正解决的是两个连续的系统问题：

方法	核心问题	主要优化层次
Standard Attention	直接按算子执行，产生大量中间结果
FlashAttention	同样计算 exact attention，如何显著减少 HBM 数据搬运	HBM SRAM
FlashAttention-2	IO 已经优化后，如何进一步提高 GPU 利用率	Thread Block / Warp / Tensor Core

因此，理解 FlashAttention 系列最重要的视角不是“减少了多少 FLOPs”，而是：

graph LR
A(Standard Attention) --> B(IO-aware Attention) --> C(GPU-efficient IO-aware Attention)

Attention 的性能瓶颈

设单个 attention head 的

标准 Attention 可以写成：

其中

从算法复杂度看，两次主要矩阵乘都需要 FLOPs，因此很容易得到“Attention 慢是因为 FLOPS 计算量太大”的结论。但在真实 GPU 上，kernel 的时间不仅取决于计算量，还取决于数据在内存层次之间移动的代价。

标准实现通常经历：

也就是说，完整的和往往会被 materialize 到 HBM，然后再次读回。对于长序列，这两个中间矩阵会迅速成为显存带宽和 activation memory 的主要负担。

以 A100 为例，HBM 容量很大、带宽约为 TB/s 量级，而每个 SM 的片上 SRAM 容量只有数百 KB，却拥有远高于 HBM 的带宽。于是 GPU kernel 大致可分为两类：大型 GEMM 往往具有较高 arithmetic intensity，更接近 compute-bound；softmax、mask、dropout、reduction 等操作 FLOPs 很少，但需要反复读取和写回数据，更容易 memory-bound。

对 memory-bound workload，更合理的性能模型是：

而不是只看 FLOPs。FlashAttention 的核心出发点正是：Attention 的主要瓶颈之一，是中间矩阵造成的 HBM IO。

FlashAttention

FlashAttention 的核心目标是

$仍然计算，但完整不写入$

$仍然精确计算，但完整不写入$

它把切成能够放入片上 SRAM 的 tile，在片上完成局部 score 计算、softmax、mask、dropout 和累积。局部的、使用完立即丢弃，从而避免完整矩阵的 materialization。

真正困难的地方不是矩阵乘，而是 softmax。矩阵乘天然可以分块：

但一行 softmax 的归一化依赖整行所有元素。FlashAttention 能够分块计算 exact softmax，依赖的就是 online softmax。

Online Softmax

考虑一行 score：

为了数值稳定，softmax 通常写成：

如果 SRAM 中一次只能看到一个 score block，我们暂时不知道最终全局最大值，也不知道最终 denominator 。关键观察是：当新的更大最大值出现时，旧的指数和不需要重新计算，只需要重新缩放。

假设已经处理过一部分 scores，保存：

新 block 的局部最大值为：

新的全局最大值：

由于

旧 denominator 可以直接重标定：

因此，完整 softmax 不需要保存完整 score row，只需要持续维护 row-wise 的和。

FlashAttention 前向过程

论文 Algorithm 1 的输入是：

它们存放在 HBM 中，同时假设片上 SRAM 可用容量为。算法首先选择行、列 tile 大小，使正在处理的、局部 score 以及必要的统计量能够放进 SRAM。

随后把：

其中

并有：

算法为每个 query row 维护三个状态：

初始化为：

其中表示目前见过的最大 score，表示对应数值稳定形式下的 softmax denominator，表示目前累计得到的归一化 attention 输出。

第一版 FlashAttention 的循环顺序是：

for each K_j, V_j block:
    load K_j, V_j from HBM to SRAM

    for each Q_i block:
        load Q_i, O_i, m_i, l_i
        compute S_ij = Q_i @ K_j^T
        update online-softmax statistics
        update O_i with V_j
        write O_i, m_i, l_i back

外层固定一个 block，内层让所有依次与它计算。这样一个已经搬入 SRAM 的 K/V tile 能被多个 Q tile 复用，减少重复的 HBM 访问。

对于当前 tile，首先计算：

实际 Transformer 中还包含缩放、causal mask 等操作，这些都可以融合进 tile 内计算，而不需要额外 materialize 一个大矩阵。

然后对每一行计算当前 tile 的局部统计量：

这里的还不是最终归一化后的 softmax probability，因为当前只处理了一个 K/V block。

接下来将旧状态与当前 block 合并。新的最大值：

新的 denominator：

这一步就是 online softmax 的 row-wise recurrence。

仅更新 denominator 还不够，因为 Attention 最终需要的是。对于一行 query，当前输出可以写成：

如果旧结果已经是归一化后的，那么旧 numerator 实际上是：

当最大值从更新为后，旧 numerator 必须按照同一尺度重新缩放；当前 block 的 value contribution 也要转换到新尺度。因此：

论文伪代码中的 diag 只是把每一行各自的写成矩阵形式，本质就是上面的 row-wise 更新。

这一递推始终维持三个不变量。处理到第个 K/V block 时：

当所有 K/V blocks 处理完成后，就与一次性计算完整 softmax 后的结果完全相同。

这也解释了为什么 FlashAttention 可以“边算边丢”。一个局部 tile 的生命周期是：

完整的从未需要出现在 HBM 中。

Tiling、Fusion 与 Backward Recomputation

FlashAttention 的 tiling 与 kernel fusion 是同一件事的两个侧面。标准 Attention 可能跨多个 kernel 依次执行：

FlashAttention 则把：

尽可能融合在同一个 kernel 的片上执行路径中。online softmax 使 tile 可以独立推进，tiling 又使中间矩阵无需落到 HBM，从而为 aggressive fusion 创造条件。

训练时 backward 更能体现“IO 优先”的设计思路。标准 Attention backward 涉及：

以及：

最直接的实现会在 forward 保存完整，导致 activation memory。FlashAttention 选择不保存，而是在 backward 重新计算：

再利用 forward 保存的 normalization statistics 重建：

因此 FlashAttention 的一个关键工程原则是：

也就是：如果重新计算主要发生在 Tensor Core、register 和 SRAM 中，而它替代的是大量 HBM 读写，那么增加 FLOPs 反而可能缩短 wall-clock time。

Softmax backward 还有一个重要化简。定义：

由于：

可以得到：

于是局部 backward 可写成：

其中只需要长度的输出与梯度即可得到，不需要保存长度的完整 probability row。这进一步使 backward 适合 tiled execution。

复杂度

FlashAttention 没有改变 dense attention 的主计算复杂度：

它改变的是中间存储和 HBM IO。完整的不再需要 materialize；额外 row-wise statistics 只需要规模，而必要的输入输出仍保持原本的。

论文给出的 HBM IO complexity 对比为：

其中表示可用 SRAM capacity 参数。在典型 head dimension 和合适 tile 大小时，后者显著减少 HBM traffic。

论文中的一组 GPT-2 medium 实验

指标	Standard Attention	FlashAttention
GFLOPs	66.6	75.2
HBM R/W	40.3 GB	4.4 GB
Forward + Backward	41.7 ms	7.3 ms

FlashAttention 的 FLOPs 甚至更多，但 HBM 读写从 40.3 GB 降到 4.4 GB，运行时间反而从 41.7 ms 降到 7.3 ms。这个结果很好地说明了 FlashAttention 的核心 它主要降低的是 IO complexity，而不是主 FLOP complexity.

FlashAttention-2

FlashAttention 已经解决了最严重的 HBM IO 问题，但第一版 kernel 的吞吐仍明显低于高质量 GEMM。FlashAttention-2 的 profiling 发现，新的瓶颈集中在 GPU work partitioning：thread block 数量不足、occupancy 不够、warp 分工不理想、shared-memory communication 偏多，以及 non-matmul FLOPs 的实际代价较高。

因此 FA2 的重点从：

继续下沉到：

减少 Non-Matmul FLOPs

在 A100 上，FP16/BF16 Tensor Core matmul 的理论吞吐远高于普通 FP32 non-matmul operation。也就是说，从硬件执行成本看，一个 elementwise/reduction FLOP 与一个 Tensor Core matmul FLOP 并不“等价”。

FA1 在每个 K/V tile 后维护已经归一化的，因此 inner loop 中反复执行缩放和除法。FA2 改为维护 unnormalized numerator accumulator，不再每个 block 都归一化

当新 block 到来时：

只有所有 column blocks 处理完成后，才做一次：

数学结果不变，但 inner loop 中减少了昂贵的 non-matmul work。

FA2 还把 forward 保存的 softmax statistics 压缩成 LogSumExp。由于：

定义：

即可写成：

因此 backward 不必同时依赖，保存即可完成 probability reconstruction。

Sequence Parallelism

第一版 FlashAttention 的主要并行维度近似是：

如果 batch 很小、head 数也有限，就可能没有足够的 thread blocks 填满 GPU。长序列尤其容易出现这个问题，因为显存压力往往迫使 batch size 进一步减小。

FA2 将 Q 的 row block 也纳入并行维度。FA1 的逻辑顺序近似是：

1
2
3

for K/V column block:
    for Q row block:
        compute tile

FA2 forward 改为：

1
2
3

parallel for Q row block:
    for K/V column block:
        compute tile

也就是把：

作为 thread-block 级的独立工作单元。并行任务数量由近似：

提升到：

其中：

例如：

则：

并行 work units 从约提升到：

这使大量 SM 更容易保持忙碌。

Forward 适合 row-oriented parallelism，因为每个可以独立遍历所有 K/V blocks 并最终生成自己的。Backward 的依赖更复杂：为了高效计算，FA2 可以采用 column-oriented work partition，而多个 worker 对同一的贡献需要累积，因此会引入 atomic add。这里的关键不是保持 forward/backward 形式对称，而是让 schedule 服从数据依赖。

Warp Work Partitioning

增加 thread-block 数量解决的是 block-level parallelism，FA2 还进一步优化一个 block 内多个 warp 的分工方式。

FA1 forward 的典型方案可以理解为 Split-K：多个 warp 共享同一组 Q rows，而把 K/V 方向上的工作拆开。这样在阶段，不同 warp 会得到同一组 output rows 的 partial contribution：

于是需要把 partial results 写入 shared memory、同步、重新读取并 reduction。真正昂贵的不是几次加法，而是：

FA2 改为 Split-Q。把 Q rows 在 warp 之间切分：

而 K/V 对这些 warp 可复用。每个 warp 负责独立的一部分 output rows：

最终输出是拼接：

而不是多个 partial output 的求和，因此不需要 warp 间 reduction。

这体现了一个很通用的 GPU kernel 原则：

$尽量让一个拥有最终，而不是共同拥有$

后者通常意味着更多 synchronization、shared-memory traffic 和 reduction。

Tile Size 与 Causal Attention

FlashAttention 并不是 tile 越大越好。增大可以提高数据复用、减少某些访问，但同时会提高 register pressure 和 shared-memory usage，可能降低 occupancy；严重时还会发生 register spilling，甚至超过 shared-memory capacity。

因此实际调优必须同时考虑：

Causal Attention 还可以利用 tile 结构进一步减少无效计算。对于：

如果一个 tile 完全位于 causal mask 的上三角区域，那么整个 tile 都不会贡献输出，可以直接跳过，不必计算。只有跨越主对角线的 tile 才需要进行 elementwise masking。这里 tiling 的价值已经不只是改善内存访问，它还把结构化条件从 element level 提升到了 block level。

FlashAttention 与 FlashAttention-2 对比

维度	FlashAttention	FlashAttention-2
Attention 数学定义	Exact dense attention	Exact dense attention
主计算复杂度
主要瓶颈	HBM IO	GPU utilization
核心优化层次	HBM SRAM	Thread Block / Warp
Softmax	Online softmax	Online softmax + 算术精简
matrix	不写入 HBM	不写入 HBM
Forward accumulator	每 block 更新 normalized	维护 unnormalized ，最后统一归一化
Forward statistics		LogSumExp
Backward	Recomputation	Recomputation + 更好的 work partition
主要并行维度	Batch × Heads	Batch × Heads × Sequence
Forward loop orientation	K/V → Q	Q → K/V
Warp partition	Split-K	Split-Q
Shared-memory communication	较多	更少

从 Roofline Model 看，两代算法的关系尤其清楚。一个 kernel 的实际性能受：

限制，其中：

是 arithmetic intensity。

标准 Attention 中，大量 materialization 让 pipeline 存在严重 memory-bound 部分。FA1 通过减少 bytes transferred 提高有效 arithmetic intensity，使 workload 更接近 compute-bound。到了这个阶段，新的问题才暴露出来，计算单元是否真的被充分利用？FA2 因此继续优化 occupancy、thread-block parallelism、warp-level work partition、shared-memory communication 和 non-matmul FLOPs。

FA1 的成功并没有结束优化，反而暴露了 FA2 要解决的问题。

两篇论文最值得迁移到其他高性能 kernel 的经验也可以归纳为同一条逻辑：算法复杂度不能替代硬件性能模型；数据移动本身就是成本；recomputation 不一定比 memory access 贵；并行任务数量和 work partitioning 同样重要；tile size 必须在数据复用、寄存器、shared memory 与 occupancy 之间综合权衡。

总结

FlashAttention 的突破，不是发明了新的 Attention 数学形式，而是重新安排 exact attention 的执行顺序。通过 tiling、online softmax、kernel fusion 和 backward recomputation，它避免把完整的 materialize 到 HBM，从而显著降低 HBM IO 和 activation memory。

FlashAttention-2 则进一步认识到，当 IO 瓶颈被缓解后，性能会转而受 GPU 并行映射限制。于是它减少 inner loop 中的 non-matmul FLOPs，引入 sequence-level thread-block parallelism，并将 warp work partition 从 Split-K 改成 Split-Q，减少 shared-memory communication，提高 occupancy 和 Tensor Core utilization。

因此，FlashAttention 系列真正展示的是一条完整的 algorithm-hardware co-design 路径：

高性能 Attention 的关键不只是“算多少”，而是同时设计：数据放在哪里、什么时候搬、以什么粒度计算、由谁计算，以及中间结果是否值得保存。

参考资料

arXiv | Computer Science > Machine Learning | FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

arXiv | Computer Science > Machine Learning | FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning

FlashAttention：从 IO-Aware Attention 到 GPU Work Partitioning

Attention 的性能瓶颈

FlashAttention

Online Softmax

FlashAttention 前向过程

Tiling、Fusion 与 Backward Recomputation

复杂度

FlashAttention-2

减少 Non-Matmul FLOPs

Sequence Parallelism

Warp Work Partitioning

Tile Size 与 Causal Attention

总结

参考资料

论文精选

评论

搜索文章

Attention 的性能瓶颈

FlashAttention

Online Softmax

FlashAttention 前向过程

Tiling、Fusion 与 Backward Recomputation

复杂度

FlashAttention-2

减少 Non-Matmul FLOPs

Sequence Parallelism

Warp Work Partitioning

Tile Size 与 Causal Attention

总结

参考资料

论文精选

相关文章

从 Switch Transformer 到 GLaM

大模型分布式训练并行策略

语言模型缩放定律与计算最优训练

评论