《FlashAttention：具备 IO 感知的快速精确注意力》

注意力为什么是内存问题

注意力（attention）让序列里的每个位置都能看到其他所有位置，这正是 Transformer 有效的原因，也是序列变长后代价急剧上升的原因。多数加速尝试着眼于减少运算量。FlashAttention 的作者——Tri Dao、Daniel Y. Fu、Stefano Ermon、Atri Rudra 和 Christopher Ré——认为此前缺少一条原则：让注意力算法具备 IO 感知（IO-aware），也就是把 GPU 各级内存之间的读写算进去，而不只是数浮点运算次数。论文于 2022 年 5 月 27 日提交到 arXiv ^[1]。

促使他们这样想的硬件差距很直观。一块 A100 GPU 有 40-80GB 的高带宽内存（high bandwidth memory，HBM），但它的 108 个流式多处理器每个只有 192KB 的片上静态随机存储器（SRAM）^[1]。数据在这两级之间搬运，比在数据已经就位时做计算要慢得多。常规的注意力实现会把规模很大的注意力分数矩阵写到 HBM，再在后续步骤里读回来，于是相当一部分时间花在搬运数据而不是计算上。

分块，但不牺牲精确性

FlashAttention 是一种具备 IO 感知的精确注意力算法。它采用分块（tiling）：把输入切成小到能放进片上 SRAM 的块，逐块完成注意力计算，中间结果留在片上并逐步合并，而不是写回 HBM 再读出来。论文指出，这样做所需的 HBM 访问次数少于标准注意力，并且在一系列 SRAM 容量下都是最优的 ^[1]。

“精确”这个词是关键。此前很多降低注意力成本的办法是近似：跳过或压缩分数矩阵的一部分，用模型质量换速度。FlashAttention 算出的是与标准写法相同的注意力结果，省下来的是数据搬运，而不是换算了别的东西。

实测的加速幅度

论文报告了在若干工作负载上快于已有基线的训练速度。与 MLPerf 1.1 的训练速度纪录相比，它在 BERT-large 上取得 15% 的端到端实际耗时（wall-clock）加速；在 GPT-2 上达到 3 倍，在 Long Range Arena 基准上达到 2.4 倍 ^[1]。这些数字比较的是真实训练运行的耗时而非运算次数，因此直接体现了关于内存搬运的论证：同样的数学工作更早做完。

更长的上下文与一个难题

速度并不是唯一的后果。FlashAttention 及其块稀疏（block-sparse）扩展让更长的上下文变得可行，作者把这一点与更高质量的模型联系起来。他们报告了首批在 Path-X 挑战上表现优于随机猜测的 Transformer，该任务的序列长度为 16K，准确率 61.4% ^[1]。Path-X 值得注意，是因为如此长的序列此前超出了标准实现的能力范围，模型至多只能靠猜。

这项工作改变了问题的提法

论文所描述的持久贡献，一半是方法，一半是视角：把内存层级当作算法设计的一等要素，注意力就会明显变便宜，而且不需要任何近似。这个视角把抽象的运算——缩放点积注意力——与运行它的硬件的具体约束连在一起，也解释了为什么能腾出这么大的空间。它还移动了实际边界：一个 Transformer 能训练多长的序列，时间与内存都受同一件事支配，即 HBM 与 SRAM 之间的搬运 ^[1]。