FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, 2023

[AkihikoWatanabe]

https://tridao.me/publications/flash2/flash2.pdf

[AkihikoWatanabe]

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 [5] exploits the asymmetric GPU memory hierarchy to bring significant memory saving (linear instead of quadratic) and runtime speedup (2-4× 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× 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)

Translation (by gpt-3.5-turbo)

• 長いシーケンス長へのTransformerのスケーリングは、過去数年間における主要な問題であり、言語モデリングや高解像度画像理解の性能向上、さらにはコード、音声、ビデオ生成の新しいアプリケーションの開拓を約束しています。
  シーケンスの長さに対して注意層は主なボトルネックであり、ランタイムとメモリがシーケンスの長さに対して二次的に増加します。
  FlashAttention [5]は非対称なGPUメモリ階層を利用して、有意なメモリの節約（二次的ではなく線形的）とランタイムの高速化（最適化されたベースラインと比較して2-4倍）を実現し、近似なしで行います。
  しかし、FlashAttentionは最適化された行列乗算（GEMM）演算に比べてまだ十分に高速ではありませんでした。理論的な最大FLOPs/sの25-40%にしか達しません。
  効率の悪さは、異なるスレッドブロックとワープ間での最適でない作業の分割に起因していることがわかりました。これにより、低い占有率または不必要な共有メモリの読み書きが発生します。
  これらの問題に対処するために、FlashAttention-2を提案します。具体的には、(1) 非matmul FLOPsの数を減らすためにアルゴリズムを調整し、(2) アテンション計算を並列化し、単一のヘッドでも異なるスレッドブロック間で行います。これにより占有率が向上します。(3) 各スレッドブロック内でワープ間の通信を減らすために作業を分散させます。
  これにより、FlashAttentionに比べて約2倍の高速化が実現され、A100では理論的な最大FLOPs/sの50-73%に達し、GEMM演算の効率に近づきます。
  GPTスタイルのモデルのトレーニングにエンドツーエンドで使用する場合、FlashAttention-2はA100 GPUあたり最大225 TFLOPs/sのトレーニング速度に達し、モデルのFLOPs利用率は72%です。

Summary (by gpt-3.5-turbo)

• FlashAttention-2は、長いシーケンス長におけるTransformerのスケーリングの問題に対処するために提案された手法です。
  FlashAttention-2は、非対称なGPUメモリ階層を利用してメモリの節約とランタイムの高速化を実現し、最適化された行列乗算に比べて約2倍の高速化を達成します。
  また、FlashAttention-2はGPTスタイルのモデルのトレーニングにおいても高速化を実現し、最大225 TFLOPs/sのトレーニング速度に達します。

[AkihikoWatanabe]

Flash Attention1よりも2倍高速なFlash Attention 2

[AkihikoWatanabe]

Flash Attention1はこちらを参照
https://arxiv.org/pdf/2205.14135.pdf

QK Matrixの計算をブロックに分けてSRAMに送って処理することで、3倍高速化し、メモリ効率を10-20倍を達成。