FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2024
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866909253401313280 |
|---|---|
| author | Shah, Jay Bikshandi, Ganesh Zhang, Ying Thakkar, Vijay Ramani, Pradeep Dao, Tri |
| author_facet | Shah, Jay Bikshandi, Ganesh Zhang, Ying Thakkar, Vijay Ramani, Pradeep Dao, Tri |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2407_08608 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision Shah, Jay Bikshandi, Ganesh Zhang, Ying Thakkar, Vijay Ramani, Pradeep Dao, Tri Machine Learning Artificial Intelligence 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. |
| title | FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision |
| topic | Machine Learning Artificial Intelligence |
| url | https://arxiv.org/abs/2407.08608 |