FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision

要約

ユビキタスな Transformer アーキテクチャのコア層としてのアテンションは、大規模な言語モデルやロングコンテキストのアプリケーションにとってボトルネックとなります。
FlashAttention は、メモリの読み取り/書き込みを最小限に抑えることで GPU でのアテンションを高速化するアプローチを詳しく説明しました。
ただし、最近のハードウェアにある新機能はまだ活用されておらず、H100 GPU では FlashAttendant-2 の使用率はわずか 35% にとどまっています。
私たちは、Hopper GPU での注意を高速化するための 3 つの主な手法を開発しています。1 つは、Tensor コアと TMA の非同期性を利用して、(1) ワープ特化による全体的な計算とデータ移動をオーバーラップする、(2) ブロック単位の matmul およびソフトマックス操作をインターリーブする、および (
3) FP8 低精度のハードウェア サポートを活用したブロック量子化とインコヒーレント処理。
私たちの手法である FlashAttendant-3 は、H100 GPU で 1.5 ～ 2.0$\times$ の高速化を達成し、FP16 では最大 740 TFLOPs/s (使用率 75%) に達し、FP8 では 1.2 PFLOPs/s 近くに達することを実証します。
FP8 FlashAttendant-3 は、ベースラインの FP8 Attendance よりも 2.6$\times$ 低い数値誤差を達成することを検証します。

要約(オリジナル)

Attention, as a core layer of the ubiquitous Transformer architecture, is the bottleneck for large language models and long-context applications. FlashAttention elaborated an approach to speed up attention on GPUs through minimizing memory reads/writes. However, it has yet to take advantage of new capabilities present in recent hardware, with FlashAttention-2 achieving only 35% utilization on the H100 GPU. We develop three main techniques to speed up attention on Hopper GPUs: exploiting asynchrony of the Tensor Cores and TMA to (1) overlap overall computation and data movement via warp-specialization and (2) interleave block-wise matmul and softmax operations, and (3) block quantization and incoherent processing that leverages hardware support for FP8 low-precision. We demonstrate that our method, FlashAttention-3, achieves speedup on H100 GPUs by 1.5-2.0$\times$ with FP16 reaching up to 740 TFLOPs/s (75% utilization), and with FP8 reaching close to 1.2 PFLOPs/s. We validate that FP8 FlashAttention-3 achieves 2.6$\times$ lower numerical error than a baseline FP8 attention.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google