FlashAttention-3：通过异步和低精度实现快速准确的注意力机制

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。 Ke...

阅读全文