arXiv。 arXiv:2205.14135（Submitted 27 May 2022，last revised 23 Jun 2022，v2）
摘要
Transformers are slow and memory-hungry on long sequences, since the time and memory complexity of self-attention are quadratic in sequence length. Approximate attention methods have attempted to address this problem by trading off model quality to reduce the compute complexity, but often do not achieve wall-clock speedup. We argue that a missing principle is making attention algorithms IO-aware—accounting for reads and writes between levels of GPU memory. We propose FlashAttention, an IO-aware exact attention algorithm that uses tiling to reduce the number of memory reads/writes between GPU high bandwidth memory (HBM) and GPU on-chip SRAM. We analyze the IO complexity of FlashAttention, showing that it requires fewer HBM accesses than standard attention, and is optimal for a range of SRAM sizes. We also extend FlashAttention to block-sparse attention, yielding an approximate attention algorithm that is faster than any existing approximate attention method. FlashAttention trains Transformers faster than existing baselines: 15% end-to-end wall-clock speedup on BERT-large (seq. length 512) compared to the MLPerf 1.1 training speed record, 3× speedup on GPT-2 (seq. length 1K), and 2.4× speedup on long-range arena (seq. length 1K-4K). FlashAttention and block-sparse FlashAttention enable longer context in Transformers, yielding higher quality models (0.7 better perplexity on GPT-2 and 6.4 points of lift on long-document classification) and entirely new capabilities: the first Transformers to achieve better-than-chance performance on the Path-X challenge (seq. length 16K, 61.4% accuracy) and Path-256 (seq. length 64K, 63.1% accuracy).
摘要指出：自注意力在序列长度上平方，近似方法常不能转化为墙钟加速。缺失原则是 IO-aware，即计入 GPU 存储层次之间的读写。FlashAttention 用分块减少 HBM 与片上 SRAM 之间的搬运，并保持精确注意力。文中给出 IO 复杂度分析、block-sparse 扩展，以及 BERT-large、GPT-2、Long Range Arena 与 Path-X / Path-256 上的测量。
总结
该工作保持注意力在数学上精确，改变的是中间张量存放位置。标准实现把 N × N 分数矩阵写入 HBM。FlashAttention 将 Q、K、V 分块载入 SRAM，维护在线 Softmax，不物化该矩阵。反向传播接受额外算术，因为 HBM 流量更昂贵。2022 年的测量报告 BERT-large、GPT-2 与 Long Range Arena 训练更快，并在 Path-X（16K）与 Path-256（64K）上给出 Transformer 路线首次超过随机水平的结果。算术仍为 O(N²)。后续融合 Kernel 不会自动再现原文墙钟倍数。
详细解读
IO-awareness 与仅减少 FLOPs 的区别
摘要将困难定位于 GPU 存储层次之间的搬运，而非注意力算术本身。标准注意力物化 S = QK^T，逐行 Softmax，再计算 O = PV。长度为 N 时，中间矩阵含 N² 个元素。平方级算术确实存在。然而在 A100 上，HBM 带宽约为 1.5–2.0 TB/s，片上 SRAM 快一个数量级，每 SM 仅约 192 KB。因此，即使 FLOPs 看起来可接受，注意力仍可能受 IO 限制。
Although these methods reduce the compute requirements to linear or near-linear in sequence length, many of them do not display wall-clock speedup against standard attention and have not gained wide adoption.
将理论乘法降为线性或近线性的近似方法，若引入额外 Kernel、不规则读取或物化中间结果，墙钟时间仍可能没有优势。论文将墙钟时间、而非单纯的 FLOP 减少，作为相关判据。
分块与精确 Softmax
(i) We restructure the attention computation to split the input into blocks and make several passes over input blocks, thus incrementally performing the softmax reduction (also known as tiling).
Q、K、V 被切分。Query 块与相继的 Key/Value 块从 HBM 载入 SRAM。Kernel 更新输出统计后丢弃该 Tile。完整的 N × N 注意力矩阵不会写入 HBM。
逐行 Softmax 依赖于全部 Key，各 Tile 独立归一化将得到错误结果。FlashAttention 维护在线 Softmax：运行最大值、指数和，以及部分加权的 Value 累加。后续 Tile 若提高最大值，则将既有累加量缩放到新的参考系。遍历全部 Tile 后，结果与标准 Softmax 一致，仅存在通常的浮点运算顺序差异。该算法是精确的，既非稀疏，亦非低秩近似。
重计算与 IO 复杂度
In contrast, even with more FLOPs, our recomputation speeds up the backward pass due to reduced HBM accesses (Fig. 2).
常规反向传播保存注意力概率矩阵。FlashAttention 保存紧凑的行统计 (m, ℓ)，并在片上重建分数 Tile。额外算术被接受，因为 HBM 流量更为昂贵。论文 Figure 2 显示 FLOPs 更高、HBM 字节更少、运行时间更短。
Standard attention (Algorithm 0) requires Θ(Nd+N²) HBM accesses, while FlashAttention (Algorithm 1) requires Θ(N²d²M^{-1}) HBM accesses.
在 SRAM 容量为 M、Head 维度取常见值时，d² 远小于 M，访问量因此明显下降。论文并指出，在保持精确的前提下，不存在对所有 SRAM 容量都能渐近改进该界的算法。算术复杂度仍为 O(N²)，改变的是局部性。仅统计 FLOPs 的复杂度表会遗漏这一区分。
原文报告的测量
We train BERT-large (seq. length 512) 15% faster than the training speed record in MLPerf 1.1, GPT2 (seq. length 1K) 3× faster than baseline implementations from HuggingFace and Megatron-LM, and long-range arena (seq. length 1K-4K) 2.4× faster than baselines.
在 8 张 A100 上，BERT-large 训练时间报告为 17.4 ± 1.4 分钟，所引 MLPerf 1.1 纪录为 20.0 ± 1.5 分钟。GPT-2 medium 相对 HuggingFace 约 3 倍，相对 Megatron-LM 约 1.7–1.8 倍；困惑度不变，因为模型定义未改。GPT-2 small 在 4K 上下文下仍快于 Megatron 的 1K，困惑度改善 0.7。Path-X（16K）与 Path-256（64K）被报告为 Long Range Arena 上 Transformer 路线首次超过随机水平。
证据边界
上述倍数属于 2022 年的硬件与软件环境。现代框架中的后续融合 Kernel 不会自动再现这些数字。FlashAttention 仍计算稠密配对，算术仍为平方级。NSA 等稀疏设计改变的是哪些配对存在；二者可以组合，但所回答的问题不同。较为持久的设计原则是：融合运算以避免中间结果往返 HBM；按片上容量分块；当搬运成本高于计算时进行重计算；并同时报告墙钟时间与存储层次。