FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré
cs.LG
2022-05-28
FlashAttention在SRAM分块做精确注意力，HBM读写降到约1/9，GPT-2从9.5天降到2.7天。
自注意力的时间和显存都随序列长度 N 平方增长。把注意力近似成线性或近线性的方法不少，墙钟时间却经常快不了。这些方法在减 FLOPs，而 GPU 上很多算子的时间取决于读写次数。A100 的 HBM 带宽约 1.5-2 TB/s，片上 SRAM 约 19 TB/s，但每个 SM 只有大约 192KB。
标准实现先把 S = QK^T 写入 HBM，再读出来做 softmax 得到 P，再把 P 和 V 读回来做乘。S 和 P 都是 N×N。GPT-2 常见配置 N=1024、头维度 d=64，这个矩阵比 Q、K、V 大一个数量级。mask 和 dropout 还要再碰它几遍。
FlashAttention 算的仍是精确注意力，softmax(QK^T)V 的数学结果不变。变的是数据在 HBM 和 SRAM 之间怎么走。Q、K、V 被切成 SRAM 放得下的块。外层每次载入一块 K 和 V，内层再载入一块 Q，在片上完成这一对块的矩阵乘、softmax、可选的 mask 与 dropout、再乘上 V。写回 HBM 的只有输出块，没有 N×N 的 S 或 P。
Softmax 按行把所有 key 耦在一起，块与块不能各算各的再直接相加。每一行额外记两个数：已经看过的最大值 m，以及减完这个最大值之后的指数和 ℓ。下一块到来时，用新旧最大值之差做缩放，把旧的输出和旧的 ℓ 重标定，再加进新块。块全部扫完，结果与一次性 softmax 相同。
反向如果要把整张 P 存下来，显存又回到平方。前向只留下输出 O 和统计量 m、ℓ。反向在 SRAM 里用 Q、K、V 把对应的注意力块重算出来再求梯度，dropout 的掩码也靠保存的随机数状态重放，不必存一张 N×N 的掩码。浮点运算变多，HBM 访问变少。前向和反向都融进一个 CUDA kernel。
标准实现的 HBM 访问是 Θ(Nd + N²)，FlashAttention 是 Θ(N²d²/M)，M 是 SRAM 容量。d 在 64 到 128、M 约 100KB 时，访问量可以少一个数量级。块稀疏版本跳过全零块，主导项再乘非零块比例 s，实验用固定的 butterfly 图案。
快，是因为少搬，不是因为少算。GPT-2 medium、序列长度 1024、头维度 64、16 个头、batch 64，在 A100 上做前向加反向：标准实现 66.6 GFLOPs、HBM 读写 40.3 GB、41.7 ms；FlashAttention 75.2 GFLOPs、4.4 GB、7.3 ms。运算大约多 13%，访存大约少到原来的九分之一，墙钟大约快 5.7 倍。单独看注意力计算相对 PyTorch，图 1 给到 7.6 倍。
BERT 和 GPT-2 的端到端时间都在 8 张 A100 上测。
| 设置 | 基线 | FlashAttention |
| BERT-large，序列 512，MLPerf 1.1 目标 72.0% | Nvidia 20.0±1.5 分钟 | 17.4±1.4 分钟，快 15% |
| GPT-2 small，OpenWebText | HuggingFace 9.5 天，ppl 18.2 | 2.7 天，ppl 18.2，约 3.5 倍 |
| GPT-2 medium | HuggingFace 21.0 天，ppl 14.2 | 6.9 天，ppl 14.3，约 3.0 倍 |
| 相对 Megatron-LM | small 4.7 天，medium 11.5 天 | 2.7 天与 6.9 天，约 1.7 倍 |
| LRA，序列 1K-4K | 标准 Transformer 平均 59.3 | 平均 59.8，快 2.4 倍 |
上下文加长才出现质量差。GPT-2 small 用 FlashAttention 把上下文加到 4K，困惑度 17.5，比 1K 的 18.2 好 0.7，训练 3.6 天，仍比 Megatron 的 1K 配置（4.7 天）少用大约三成时间。长文档分类看 micro-F1。MIMIC-III 从序列 512 的 52.8 到 16K 的 57.1，升 4.3 个点。ECtHR 从 512 的 72.2 到 8K 的 80.7，升 8.5 个点；再加到 16K 回落到 79.2。两点平均 6.4 个点。
Path-X 把 128×128 的黑白图逐像素送进去，序列长度 16K，判断两点之间有没有通路。此前的 Transformer 要么爆显存，要么停在随机猜测。FlashAttention 先在 Path-64 上预训练，再对位置编码做空间插值，准确率 61.4%。块稀疏版本开到 64K，Path-256 准确率 63.1%，同一设置在 Path-X 上是 56.0%。
单卡 A100 40GB 上，序列 128 到 2K，比 PyTorch 注意力最高快 3 倍，显存随 N 线性增长，相对精确注意力最高省约 20 倍。序列过了 512 到 1024，Linformer 等近似方法开始更快。到 64K，除 Linformer 外的基线都 OOM，FlashAttention 显存仍大约只有它的一半。LRA 上块稀疏版本快 2.8 倍，平均分 59.6，与稠密注意力接近。
上下文长度不变时，它就是精确注意力。GPT-2 small 的验证困惑度与 HuggingFace、Megatron 同为 18.2，medium 为 14.3，与 Megatron 一致。省下来的显存和时间用来把上下文拉长，模型质量的变化来自看了更长的文本。
短序列上也有墙钟收益。BERT-large 序列只有 512，仍比 MLPerf 1.1 的单机速度纪录快 15%。块稀疏版本比稠密 FlashAttention 再快 2 到 4 倍，序列可以到 64K。稀疏注意力缺的常常是访存调度。
每种新的注意力变体都要重写 CUDA kernel，比 PyTorch 低一个抽象层，工程量大，换一种 GPU 架构也不保证能直接用。分析停在单卡。多卡之间的传输是另一层 IO，没有进这个下界。
「对一段 SRAM 容量都最优」说得比证明强。证出来的是，不存在一种精确算法能在区间 [d, Nd] 的每一个 M 上把访问量再降一个渐近级别。论证只用到 M 大到 Θ(Nd) 的情形，那时下界退化为输入至少要读一遍。具体某个 M 上的参数化下界没有给出。
精确版的时间仍大致随 N 平方增长。序列超过 1K 后，部分近似方法的墙钟已经更短。块稀疏要求掩码对齐到块，butterfly 是固定图案，不是学出来的稀疏。LRA 对调参敏感，复现基线高于原榜，绝对分不宜跨论文硬比。Path-256 的 63.1% 高于 Path-X 的 61.4%，但 Path-256 的路径相对更短，分高不代表更难。
Rabe 与 Staats 同样分块，为的是峰值显存。附录里他们的速度与标准注意力相当或略慢。FlashAttention 只保留一份输出，反向只重算注意力块，少碰 HBM 的同时把时间也降下来。