FlashAttention 如何用 IO 感知改写大模型训练速度

事件概述

2022 年斯坦福团队提出 FlashAttention 算法，它不再是单纯减少注意力计算的 FLOP，而是围绕 GPU 内存层级 IO 特性重新设计计算过程。结果表明，该算法将 Transformer 训练速度提升 3 倍、内存消耗降低 20 倍，并让模型首次在 64K 超长序列上取得可用效果。

核心问题：Transformer 的瓶颈不在计算，而在数据搬运

传统注意力机制会产生一个 N×N 的中间矩阵，序列长度翻倍则计算量增至四倍，这常被归咎于算力不足。但 FlashAttention 的作者指出，真正拖慢速度的是 GPU 高带宽显存（HBM）与片上缓存（SRAM）之间的频繁读写。标准实现中，前向与反向传播合计 HBM 访问量高达 O(N²d)，远超输入输出本身的数据量 O(Nd)，大部分耗时都在搬运注意力矩阵而非运算。

技术路线：用分块计算和重计算消除 N×N 矩阵的 HBM 写入

FlashAttention 使用两项经典技术实现 IO 感知：

• 分块计算（Tiling）：将 K、V 按块加载到 SRAM，仅计算当前块与 Q 的局部注意力分数，并通过增量式 softmax 统计量合并结果，避免在 HBM 上生成完整的 N×N 矩阵。
• 重计算（Recomputation）：反向传播时不保存中间注意力矩阵，而是在需要时重新计算，虽然增加少量 FLOP，但换回了约 9 倍的 HBM 访问量下降，实际运行时间反而更短。

理论分析表明，FlashAttention 的 HBM 访问次数为 Θ(N²d/M)（M 为 SRAM 容量），已逼近精确注意力算法的理论下界。

关键表现

• 训练加速：GPT-2 small 训练时长从 HuggingFace 实现的 9.5 天缩短至 2.7 天（快 3.5 倍）；BERT-large 在 MLPerf 上比 NVIDIA 官方记录快 15%。
• 内存效率：内存占用比 PyTorch 标准实现低 20 倍，使单 GPU 即可处理 4K 以上长上下文。
• 长序列能力：GPT-2 上下文从 1K 扩至 4K 后，困惑度降低 0.7；医疗文本分类任务上 16K 序列比 512 序列提升 4.3 分；在法律判决分类上 8K 序列提升 8.5 分。在 Path-256（64K 序列）任务上，块稀疏 FlashAttention 达到 63.1% 准确率，是首个在该长度上超过随机猜测的 Transformer。
• 稀疏扩展：块稀疏版本在长序列推理基准 LRA 上再提速 2.8 倍，且精度无损。
• 硬件普适性：在 A100 上加速 2–4 倍，RTX 3090 上加速 2.5–4.5 倍，受限于 SRAM 大小的 T4 加速较少。

值得关注

FlashAttention 带来的启示不仅是某个算法的改良，而是设计范式的转变：在深度学习加速中，IO-Awareness 与计算优化同等重要，FLOP 数不能直接等同实际效率。分块与重计算这类经典方法在大模型时代仍然能产生突破性影响，并对后续长上下文模型、高效推理方案提供了直接的技术复用价值。