Elamur Research⌕
☰
home / papers / 2023 / FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning

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

Paper · 2023 · Tri Dao

Summary

The time and memory required by the attention layer in a Transformer grow quadratically with sequence length. This makes it difficult to scale to long contexts. The previous version, FlashAttention, already computed attention exactly, without approximation, and with linear memory cost. However, it used only 25–40% of the GPU's peak performance, whereas optimized GEMMs reach 80–90%. The author identified the cause: work is poorly partitioned between thread blocks and warps. As a result, the GPU is underutilized and shared memory is used unnecessarily. FlashAttention-2 addresses this with three changes. First, it reduces the number of operations that are not matrix multiplications. Second, it parallelizes the computation not only over batch and heads but also along the sequence length. Third, it partitions the work between warps within a single block differently. The new version is roughly twice as fast as the previous one and reaches 50–73% of the theoretical maximum on A100. In practice, this makes it possible to train models with twice the context length at the same cost. The code is open source.

Key points

  • A speedup of approximately 2× over FlashAttention across various settings (with and without causal mask, head dimension 64 and 128). The forward pass reaches up to 73%, and the backward pass up to 63%, of the A100's peak throughput.
  • Training GPT-style models (GPT3-1.3B and GPT3-2.7B) on 8×A100 80GB: up to 225 TFLOPs/s per GPU (72% model FLOPs utilization). This is 2.8× faster than the baseline implementation without FlashAttention and approximately 1.3× faster than FlashAttention. For 2.7B with an 8k context: 80 → 175 → 225 TFLOPs/s.
  • Why it is worth reducing non-matmul FLOPs: matrix multiplication on Tensor Cores can run up to 16× faster than other operations.
  • On H100, up to 335 TFLOPs/s was achieved without using new hardware features. Use of TMA, 4th-gen Tensor Cores and FP8 is planned for future work, and the author expects a further 1.5–2× from it.
  • Code: github.com/Dao-AILab/flash-attention. The implementation is built on the CUTLASS 3.x library.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓Stanford UniversityStanford · United States434 works · 314,419 citations🎓Princeton UniversityPrinceton · United States90 works · 39,109 citations

 Fields, topics and concepts

📚 Shelves 6The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 worksComputer Vision and Pattern Recognition · 2,869 worksHardware and Architecture · 81 works

🧭 Topics 1Research subjects and languages the work is about.GPU & parallel computing · 19 works

Cite

Cite as Tri Dao. 2023. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2307.08691 · arXiv 2307.08691

@misc{dao2023flashattention,
  title = {{FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning}},
  author = {Tri Dao},
  year = {2023},
  doi = {10.48550/arxiv.2307.08691},
  eprint = {2307.08691},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2307.08691}
}

References

20 references, 10 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13245, 2023.
  2. Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020. Attention forward + backward speed (H100 80GB SXM5) Pytorch FlashAttention FlashAttention-2 294 Speed (TFLOPs/s) 167 137 160 131 145 128 127 93 OOM 1k 2k 4k 8k 16k Sequenc…
  3. Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Unifying sparse and low-rank attention. In Advances in Neural Information Processing Systems (NeurIPS), 2021. find
  4. Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. In International Conference on Learning Representations (ICLR), 2020.
  5. Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems, 2022. find
  6. Zhe Jia and Peter Van Sandt. Dissecting the Ampere GPU architecture via microbenchmarking. GPU Technology Conference, 2021. find
  7. Zhe Jia, Marco Maggioni, Benjamin Staiger, and Daniele P Scarpazza. Dissecting the nvidia Volta GPU architecture via microbenchmarking. arXiv preprint arXiv:1804.06826, 2018. find
  8. Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pages 5156–5165. PMLR, 2020. find
  9. Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In The International Conference on Machine Learning (ICML), 2020.
  10. Benjamin Lefaudeux, Francisco Massa, Diana Liskovich, Wenhan Xiong, Vittorio Caggiano, Sean Naren, Min Xu, Jieru Hu, Marta Tintore, Susan Zhang, Patrick Labatut, and Daniel Haziza. xformers: A modular and hackable transformer modelling library. https://github.com/facebookresearch/xformers, 2022. ↗
  11. Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. arXiv preprint arXiv:1805.02867, 2018.
  12. OpenAI. Gpt-4 technical report. ArXiv, abs/2303.08774, 2023.
  13. Markus N Rabe and Charles Staats. Self-attention does not need 𝑂 (𝑛 2 ) memory. arXiv preprint arXiv:2112.05682, 2021. find
  14. Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9: 53–68, 2021.
  15. Noam Shazeer. Fast transformer decoding: arXiv:1911.02150, 2019. find
  16. Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-LM: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019. find
  17. Philippe Tillet, Hsiang-Tsung Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, pages 10–19, 2019.
  18. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
  19. Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768, 2020. find
  20. Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in Neural Information Processing Systems, 33, 2020. One write-head is all you need. arX…

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/dao-2023-flashattention-2-faster-attention-better/