FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
打开互动全文版（逐段中英对照 + 图/公式 + 论文问答）→将 Transformer 扩展到更长的序列长度是近年来的一大难题，有望提升语言建模和高分辨率图像理解的性能，并解锁代码、音频和视频生成的新应用。注意力层是扩展到更长序列的主要瓶颈，因为其运行时间和内存随序列长度二次增长。FlashAttention 利用非对称 GPU 内存层次结构，在不进行近似的情况下，显著节省内存（线性而非二次）并加速运行（相比优化基线快 2-4 倍）。然而，FlashAttention 仍远不如优化的矩阵乘法（GEMM）操作快，仅达到理论最大 FLOPs/s 的 25-40%。我们观察到，效率低下是由于 GPU 上不同线程块和线程束之间的工作分区次优，导致占用率低或不必要的共享内存读写。我们提出 FlashAttention-2，通过更好的工作分区来解决这些问题。具体来说，我们（1）调整算法以减少非矩阵乘法 FLOPs 的数量；（2）将注意力计算并行化，即使对于单个头，也在不同线程块之间进行以提高占用率；（3）在每个线程块内，在线程束之间分配工作以减少通过共享内存的通信。这些改进相比 FlashAttention 实现了约 2 倍的加速，在 A100 上达到理论最大 FLOPs/s 的 50-73%，接近 GEMM 操作的效率。我们通过端到端训练 GPT 风格模型的经验验证，FlashAttention-2 在每个 A100 GPU 上达到高达 225 TFLOPs/s 的训练速度（72%的模型 FLOPs 利用率）。
Scaling Transformers to longer sequence lengths has been a major problem in the last several years, promising to improve performance in language modeling and high-resolution image understanding, as well as to unlock new applications in code, audio, and video generation. The attention layer is the main bottleneck in scaling to longer sequences, as its runtime and memory increase quadratically in the sequence length. FlashAttention exploits the asymmetric GPU memory hierarchy to bring significant memory saving (linear instead of quadratic) and runtime speedup (2-4$\times$ compared to optimized baselines), with no approximation. However, FlashAttention is still not nearly as fast as optimized matrix-multiply (GEMM) operations, reaching only 25-40\% of the theoretical maximum FLOPs/s. We observe that the inefficiency is due to suboptimal work partitioning between different thread blocks and warps on the GPU, causing either low-occupancy or unnecessary shared memory reads/writes. We propose FlashAttention-2, with better work partitioning to address these issues. In particular, we (1) tweak the algorithm to reduce the number of non-matmul FLOPs (2) parallelize the attention computation, even for a single head, across different thread blocks to increase occupancy, and (3) within each thread block, distribute the work between warps to reduce communication through shared memory. These yield around 2$\times$ speedup compared to FlashAttention, reaching 50-73\% of the theoretical maximum FLOPs/s on A100 and getting close to the efficiency of GEMM operations. We empirically validate that when used end-to-end to train GPT-style models, FlashAttention-2 reaches training speed of up to 225 TFLOPs/s per A100 GPU (72\% model FLOPs utilization).