FlashAttention-3：基于异步与低精度的快速精准注意力
机器学习
2024-07-16 v2 人工智能
摘要
注意力机制作为 Transformer 架构的核心层，是大型语言模型和长上下文应用的瓶颈。FlashAttention 通过最小化内存读写次数来加速 GPU 上的注意力计算，但尚未充分利用近期硬件新能力，FlashAttention-2 在 H100 GPU 上仅实现 35% 的利用率。我们开发了三项主要技术，以在 Hopper GPU 上加速注意力计算：(1) 通过warp专门化利用 Tensor Core 和 TMA 的异步特性实现计算与数据传输重叠；(2)交叉进行块级矩阵乘法和 Softmax 操作；(3)块量化和非一致处理，利用硬件对 FP8 低精度的支持。我们展示，FlashAttention-3 在 H100 GPU 上可实现 1.5-2.0 倍的加速，FP16 最高可达 740 TFLOPs/s (75% 利用率)，FP8 可接近 1.2 PFLOPs/s。我们验证，FP8 FlashAttention-3 的数值误差比基准 FP8 注意力低 2.6 倍。
引用
@article{arxiv.2407.08608,
title = {FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision},
author = {Jay Shah and Ganesh Bikshandi and Ying Zhang and Vijay Thakkar and Pradeep Ramani and Tri Dao},
journal= {arXiv preprint arXiv:2407.08608},
year = {2024}
}