Abstract
注意力机制作为无处不在的Transformer架构的核心层，是大语言模型和长上下文应用的瓶颈。FlashAttention通过最小化内存读写提出了一种在GPU上加速注意力计算的方法。然而，它尚未利用最新硬件的新功能，FlashAttention-2在H100 GPU上仅实现了35%的利用率。我们开发了三种主要技术来加速Hopper GPU上的注意力计算：利用Tensor Cores和TMA的异步特性来(1)通过warp特化重叠整体计算和数据移动，(2)交错执行块级矩阵乘法和softmax操作，以及(3)块量化和非相干处理来利用硬件对FP8低精度的支持。我们展示了FlashAttention-3在H100 GPU上实现了1.5-2.0倍的加速，使用FP16达到最高740 TFLOPs/s(75%利用率)，使用FP8接近1.2 PFLOPs/s。
Key Contributions
- 三种利用Hopper GPU功能的新技术：warp特化、操作交错和FP8量化
- Tensor Cores和TMA的异步处理，实现计算与数据移动的重叠
- 矩阵乘法和softmax操作的块级交错，提高效率
- 硬件高效的FP8支持，采用块量化和非相干处理
- 实现75% GPU利用率(740 TFLOPs/s FP16)和接近1.2 PFLOPs/s的FP8性能
- 在H100 GPU上相比FlashAttention-2实现1.5-2.0倍加速
Methodology
FlashAttention-3为Hopper GPU引入了三项关键创新：(1) Warp特化利用Tensor Cores和TMA(Tensor Memory Accelerator)之间的异步性来重叠计算与数据移动，(2)块级操作交错在矩阵乘法和softmax操作之间交替执行以最大化硬件利用率，(3)块级粒度的FP8量化和非相干处理利用原生FP8 Tensor Core支持。这些技术协同工作，从现代GPU架构中提取最大性能。
Experiments
在H100 GPU上的综合性能评估表明：(1) FP16实现了高达740 TFLOPs/s，代表75%的硬件利用率，相比FlashAttention-2的35%有显著提升，(2) FP8精度接近1.2 PFLOPs/s，(3)相比FlashAttention-2实现1.5-2.0倍加速，(4) FP8变体显示出比基线FP8注意力低2.6倍的数值误差，同时展示了速度和准确性的提升。
Evaluation Notes
FlashAttention-3代表了注意力优化的重大进步，有效利用了Hopper GPU架构。这项工作展示了对现代GPU硬件能力的深刻理解以及如何将其应用于注意力计算。三管齐下的方法(异步、交错、FP8)设计精良，取得了令人印象深刻的结果：将GPU利用率从35%提升到75%是一项重大成就。FP8支持对未来扩展特别有价值，而且数值误差实际上低于简单的FP8实现这一事实显示了精心的工程设计。1.5-2倍的加速直接转化为降低LLM的训练成本。这是基础性工作，很可能成为标准的注意力实现。技术深度和实际影响使其成为高效Transformer训练领域的里程碑论文。
Impact Signals
- 业界标准的注意力实现
- GPU利用率大幅提升(35% → 75%)
- 在最新硬件(H100)上实现1.5-2倍加速
- FP8支持用于未来扩展
- 来自原始FlashAttention创造者Tri Dao的工作