FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Tri Dao; Daniel Y. Fu; Stefano Ermon; Atri Rudra; Christopher Ré

论文 · arXiv 2205.14135

FlashAttention：兼顾 IO 感知的高速显存高效精确注意力机制

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré

发表于 2022-05-27cs.LG arXiv PDF

用 XYZNOT 阅读（全文翻译与 AI 讲解）

三句话总结

🎯 标准精确注意力会在较慢的 GPU HBM 中实体化存储 $O (N^{2})$ 的注意力矩阵，导致其受限于内存瓶颈，难以扩展至长上下文。
🔧 FlashAttention 通过将输入分块载入快速片上 SRAM、增量更新 softmax 并反向传播时在片上重计算注意力，在单个融合 CUDA kernel 内计算精确注意力。
📈 FlashAttention 在 GPT-2 上实现了高达 3 倍的端到端训练加速，比 MLPerf 1.1 的 BERT-large 纪录快 15%，并将上下文长度扩展至 64K tokens。

摘要

在处理长序列时，Transformer 的计算速度慢且内存占用高，因为自注意力机制的时间与内存复杂度关于序列长度呈二次方增长。近似注意力方法试图通过牺牲模型质量来降低计算复杂度，但往往无法带来实际的运行时间加速。我们认为，关键在于缺少一个原则，即让注意力算法具备 IO 感知能力——充分考虑不同级别 GPU 内存之间的读写操作。我们提出了 FlashAttention，这是一种 IO 感知的精确注意力算法，它利用分块（tiling）技术减少 GPU 高带宽显存（HBM）与 GPU 片上 SRAM 之间的内存读写次数。我们分析了 FlashAttention 的 IO 复杂度，证明了其所需的 HBM 访问次数少于标准注意力机制，并且在一定 SRAM 容量范围内达到最优。我们还将 FlashAttention 扩展至块稀疏注意力，由此产生了一种比现有任何近似注意力方法都更快的近似注意力算法。FlashAttention 训练 Transformer 的速度显著优于现有基准：与 MLPerf 1.1 的训练速度记录相比，在 BERT-large（序列长度 512）上实现了 15% 的端到端运行时间加速；在 GPT-2（序列长度 1K）上实现 3 倍加速；在 long-range arena（序列长度 1K-4K）上实现 2.4 倍加速。FlashAttention 和块稀疏 FlashAttention 赋予了 Transformer 更长的上下文处理能力，不仅带来了更高质量的模型（GPT-2 的困惑度改善 0.7，长文档分类指标提升 6.4 个百分点），还催生了全新能力：成为首个在 Path-X 挑战（序列长度 16K，准确率 61.4%）和 Path-256（序列长度 64K，准确率 63.1%）上表现优于随机猜测的 Transformer 模型。

英文原文摘要

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 $\times$ speedup on GPT-2 (seq. length 1K), and 2.4 $\times$ 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).

论文解读

问题与动机

Transformer 的时间与内存复杂度关于序列长度 $N$ 呈二次方增长。尽管近似注意力方法减少了理论上的 FLOP 计算量，但它们往往无法获得实际运行时间（wall-clock）的加速，因为现代 GPU 计算的主要瓶颈在于内存传输（IO），而非算术运算。

在标准注意力机制中，计算 $S = Q K^{⊤}$ 、 $P = softmax (S)$ 和 $O = PV$ 时，会将大小为 $N \times N$ 的中间矩阵写入速度较慢的高带宽显存（HBM），并多次读回。这会产生 $Θ (N d + N^{2})$ 次 HBM 访问以及 $O (N^{2})$ 的显存占用，从而限制了实际可用的上下文窗口长度。

核心思想

IO 感知（IO-Awareness）： 明确考虑 GPU 内存层级：快速但较小的片上 SRAM（例如 A100 上每个 SM 拥有 192 KB，带宽约 19 TB/s）对比大容量但较慢的 HBM（40–80 GB，带宽 1.5–2.0 TB/s）。
分块（Tiling，块级 Softmax 计算）： 将输入 $Q, K, V$ 切分为完全能够放入 SRAM 的分块。通过追踪动态更新的 softmax 行最大值与归一化统计量，在各个分块间增量计算 softmax 归约，而无需在 HBM 中实体化整个 $N \times N$ 矩阵。
反向传播中的重计算（Recomputation）： 避免保存用于反向传播的中间 $N \times N$ 注意力矩阵。取而代之的是，仅保存输出 $O$ 和 softmax 统计量 $(ℓ, m)$ ，反向传播时在 SRAM 中逐块重计算 $S$ 和 $P$ 。
算子融合（Kernel Fusion）： 将矩阵乘法、掩码（masking）、softmax 以及 dropout 融合到单个 CUDA kernel 中，消除冗余的 HBM 往返数据传输。

方法细节（及核心公式）

给定 $Q, K, V \in R^{N \times d}$ ，FlashAttention 设置分块大小为 $B_{c} = ⌈ M / (4 d)⌉$ 和 $B_{r} = min (⌈ M / (4 d)⌉, d)$ ，其中 $M$ 为 SRAM 容量。

增量 Softmax 追踪

对于分解为子向量 $x^{(1)}$ 和 $x^{(2)}$ 的向量，其动态统计量通过以下公式更新： $m (x) = max (m (x^{(1)}), m (x^{(2)})),$ $ℓ (x) = e^{m (x^{(1)}) - m (x)} ℓ (x^{(1)}) + e^{m (x^{(2)}) - m (x)} ℓ (x^{(2)}) .$

对于键/值（key/value）的分块 $j$ 和查询（query）的分块 $i$ ：

在 SRAM 上计算未归一化的分数： $S_{ij} = Q_{i} K_{j}^{⊤} \in R^{B_{r} \times B_{c}}$ 。
计算分块局部统计量： $\tilde{m}_{ij} = rowmax (S_{ij})$ 、 $\tilde{P}_{ij} = exp (S_{ij} - \tilde{m}_{ij})$ 以及 $\tilde{ℓ}_{ij} = rowsum (\tilde{P}_{ij})$ 。
更新动态全局归一化因子： $m_{i}^{new} = max (m_{i}, \tilde{m}_{ij}), ℓ_{i}^{new} = e^{m_{i} - m_{i}^{new}} ℓ_{i} + e^{\tilde{m}_{ij} - m_{i}^{new}} \tilde{ℓ}_{ij} .$
更新动态输出分块 $O_{i} \in R^{B_{r} \times d}$ ： $O_{i} \leftarrow diag (ℓ_{i}^{new})^{- 1} (diag (ℓ_{i}) e^{m_{i} - m_{i}^{new}} O_{i} + e^{\tilde{m}_{ij} - m_{i}^{new}} \tilde{P}_{ij} V_{j}) .$

复杂度分析

IO 复杂度： 标准注意力需要 $Θ (N d + N^{2})$ 次 HBM 访问。FlashAttention 需要 $Θ (N^{2} d^{2} M^{- 1})$ 次 HBM 访问。作者证明，对于所有 $M \in [d, N d]$ ，没有任何精确注意力算法在渐近意义上能达到 $o (N^{2} d^{2} M^{- 1})$ 次访问。
FLOPs 与显存占用： FlashAttention 执行 $O (N^{2} d)$ 次 FLOP（反向传播中由于重计算使得 FLOPs 略微增加），且除输入和输出外仅需 $O (N)$ 的辅助显存。

扩展：块稀疏 FlashAttention（Block-Sparse FlashAttention）

根据预定义的稀疏掩码 $M \in {0, 1}^{N / B_{r} \times N / B_{c}}$ 仅对非零分块进行求值，可实现 $Θ (N d + N^{2} d^{2} M^{- 1} s)$ 的 IO 复杂度，其中 $s$ 为非零分块的比例。

实验结果（核心数据）

训练加速：
- BERT-large（序列长度 512）： 较 MLPerf 1.1 的训练速度记录提升 15%（在 $8 \times$ A100 上耗时 17.4 分钟对比 20.0 分钟）。
- GPT-2（序列长度 1K）： 比 HuggingFace 加速高达 3.5 倍，比 Megatron-LM 加速高达 1.8 倍。
- Long-Range Arena (LRA)： 较标准注意力实现 2.4 倍加速；块稀疏 FlashAttention 实现 2.8 倍加速。
长上下文能力：
- 4K 上下文的 GPT-2： 训练速度快于运行在 1K 上下文下的 Megatron-LM，同时验证困惑度改善了 0.7（17.5 对比 18.2）。
- 长文档分类： 将上下文从 512 扩展到 16K，在 MIMIC-III 上的 micro- $F_{1}$ 提升了 4.3（52.8 提升至 57.1），在 ECtHR 上提升高达 8.5（在 8K 长度下从 72.2 提升至 80.7）。
- Path-X 与 Path-256： FlashAttention 是首个解决 Path-X（序列长度 16K）的 Transformer，达到 61.4% 的准确率。块稀疏 FlashAttention 在 Path-256（序列长度 64K）上达到了 63.1% 的准确率。

局限性与开放问题

CUDA 工程开销： 实现 IO 感知 kernel 需要定制的底层 CUDA 编程，而非标准的 PyTorch 或编译器抽象，使得跨硬件架构的移植具有挑战性。
多 GPU 注意力机制： 单节点最优的 IO 调度策略无法直接平滑推广到基于互连结构（NVLink、InfiniBand）的分布式多 GPU 注意力计算中。
精确注意力的理论下界： 尽管在一般的 $M$ 条件下达到最优，但证明以 $M$ 为参数的更精细参数化复杂度下界仍是一个开放问题。

关键术语

IO-awareness（IO 感知）: 在设计算子算法时，明确考虑各级硬件内存之间数据读写（如快速片上 SRAM 与较慢 HBM）的设计原则。
High Bandwidth Memory (HBM)（高带宽显存）: 容量较大但带宽显著低于片上缓存的 GPU 主存。
SRAM（静态随机存取内存）: 现代 GPU 上位于每个流式多处理器（SM）内部的高速片上暂存内存（scratchpad memory）。
Tiling（分块技术）: 通过将输入矩阵拆分为能够容纳在 SRAM 中的较小分块，并在多个计算阶段逐步累加结果来重构计算的过程。
Selective gradient checkpointing（选择性梯度检查点）: 在前向传播过程中避免保存庞大的中间注意力矩阵，而在反向传播中利用缓存的 softmax 统计量在片上重新计算该矩阵的技术。
Kernel fusion（算子融合）: 将多个独立的运算（矩阵乘法、掩码、softmax、dropout）合并到单个 GPU kernel 中执行，以消除对 HBM 的重复读写。
Block-sparse FlashAttention（块稀疏 FlashAttention）: FlashAttention 的一种扩展，仅计算由块稀疏掩码指定的非零分块，从而按稀疏比例降低 IO 复杂度。
Path-X / Path-256: 来自 Long-Range Arena 的长序列视觉推理基准任务（序列长度分别为 16K 和 64K），先前的 Transformer 模型在其上的准确率均未能超越随机猜测水平。

常见问题

FlashAttention 是对注意力矩阵进行了近似计算吗？

不是。FlashAttention 计算的是精确自注意力，其数学输出结果与标准注意力完全一致（仅受浮点数值精度误差影响）。

既然 FlashAttention 在反向传播中增加了 FLOP 计算量，为什么速度反而更快？

因为标准注意力是内存密集型的（受限于显存带宽）；在 HBM 和 SRAM 之间搬运 $N \times N$ 矩阵所耗费的时间，远多于在高速 SRAM 上执行额外算术运算的时间。

FlashAttention 在反向传播时需要存储哪些辅助显存？

FlashAttention 仅存储前向输出 $O$ 、softmax 归一化因子 $ℓ$ 和 $m$ （仅需 $O (N)$ 的额外内存），以及用于在片上复现 dropout 掩码的伪随机数生成器状态。

FlashAttention 与标准注意力的 IO 复杂度相比如何？

标准注意力需要 $Θ (N d + N^{2})$ 次 HBM 访问，而 FlashAttention 只需要 $Θ (N^{2} d^{2} M^{- 1})$ 次访问，其中 $M$ 为 SRAM 大小， $d$ 为注意力头维度。

本页解读由 AI 根据论文生成，可能有错误；细节请以原文为准。

用 XYZNOT 阅读（全文翻译与 AI 讲解）