Skip to content
The methodOpen data
Back to the collection

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré

2022 · arXiv

Language modelsAlgorithms
Read the paper
HN score
0
Mentions
8
Readers
8
Months
7
What’s in a score?

From the abstract

Transformers are slow and memory-hungry on long sequences, since the time and memory complexity of self-attention are quadratic in sequence length. Approximate attention methods have attempted to address this problem by trading off model quality to reduce the compute complexity, but often do not achieve wall-clock speedup. We argue that a missing principle is making attention algorithms IO-aware -- accounting for reads and writes between levels of GPU memory. We propose FlashAttention, an IO-aware exact attention algorithm that uses tiling to reduce the number of memory reads/writes between GPU high bandwidth memory (HBM) and GPU on-chip SRAM. We analyze the IO complexity of FlashAttention, showing that it requires fewer HBM accesses than standard attention, and is optimal for a range of SRAM sizes. We also extend FlashAttention to block-sparse attention, yielding an approximate attention algorithm that is faster than any existing approximate attention method. FlashAttention trains Transformers faster than existing baselines: 15% end-to-end wall-clock speedup on BERT-large (seq. length 512) compared to the MLPerf 1.1 training speed record, 3$\times$ speedup on GPT-2 (seq. length 1K), and 2.4$\times$ speedup on long-range arena (seq. length 1K-4K). FlashAttention and block-sparse FlashAttention enable longer context in Transformers, yielding higher quality models (0.7 better perplexity on GPT-2 and 6.4 points of lift on long-document classification) and entirely new capabilities: the first Transformers to achieve better-than-chance performance on the Path-X challenge (seq. length 16K, 61.4% accuracy) and Path-256 (seq. length 64K, 63.1% accuracy).

Metadata: arXiv. Abstracts retain their original authors’ rights.

In the conversation

Original HN comments, with labels from simple text rules. A label describes the whole comment and may miss context or irony. Read the discussion to make up your own mind. About the labels

8 selected excerpts · 8 total mentions

elegantly
Neutral mention · rule
Next Monday, 28 Apr 2025, we are launching the public permanently free unlimited, no payment, no catch of Purem – a revolutionary CPU-accelerated AI kernel that matches the performance of leading in-house infrastructures like OpenAI, Tesla, and Meta. A major leap for macOS users – explore Purem here: https://worktif.com/ After strict recalculation accounting for GPU-to-CPU architecture and memory parallelism, our CPU implementation of Softmax (6500 ops/sec) outperforms the GPU-optimized FlashAttention recalculated for standard CPUs (estimated 800–5000 ops/sec) Benchmark Highlights: • Purem achieves ~6500 ops/sec for large-scale softmax operations on Apple M2 CPU core.…
Read the full comment
Comparable to production-grade benchmarks: • OpenAI's FlashAttention: Achieves 800–5000 ops/sec on CPU. • Meta’s Xformers (PyTorch 2.0): Achieves 1300–1500 ops/sec on CPU. Key Advantages: • Local Execution: Runs entirely on your MacBook without the need for specialized hardware or cloud services. • Cost-Efficient: Delivers performance comparable to billion-dollar infrastructures at a fraction of the cost. • Open and Transparent: Built with transparency in mind, allowing for easy inspection and integration. Official Comparing References: • OpenAI FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness: Tri Dao et al., 2022, https://arxiv.org/abs/2205.14135 • FlexAttention: The Flexibility of PyTorch with the Performance of FlashAttention: PyTorch Blog, 2024, https://pytorch.org/blog/flexattention/ With Purem, experience the power of industrial-grade AI performance directly on your personal device.
Original comment 400x AI Acceleration
araes
Neutral mention · rule
For those who have no idea what's being discussed, quick background. Discussing: Transformer [1] memory issues and approximate attention [2] in machine learning training. Specifically: FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. [3] As a side comment, this entire industry is sorely in need of at least intros. The entire space has moved so fast in the last year I need an entire new dictionary and thesaurus for all the terms they've created. Notably, because of this, found out Google has a glossary of machine learning terms. Actually somewhat handy. [1] Google Machine Learning Glossary (Transformer): https://developers.google.com/machine-learning/glossary/#tra... [2] Same (Attention): https://developers.google.com/machine-learning/glossary/#att... [2] arXiv: https://arxiv.org/abs/2205.14135
Original comment Show HN: Flash Attention in ~100 lines of CUDA
kristjansson
Neutral mention · rule
FlashAttention(2)[0] reduces context-length space complexity to linear. Compute is still O(n^2) in length though, AFAIK, so we'd expect these long sequence lengths to take some time to compute. I'm a bit out of my depth, but I think ultra-long exact-attention work like this also probably has to answer some questions about where to put the KV-cache before it can be used in practice? [0]: https://arxiv.org/abs/2205.14135
Original comment LongRoPE: Extending LLM Context Window Beyond 2M Tokens
EnricoShippole
Neutral mention · rule
Introducing three new open-source PaLM models trained at a context length of 8k on C4. Open-sourcing LLMs is a necessity for the fair and equitable democratization of AI. The models of sizes 150m, 410m, and 1b are available to download and use here: https://github.com/conceptofmind/PaLM If you have any questions about the models or training be sure to reach out and ask! I will try to respond promptly. https://twitter.com/EnricoShippole/status/165559930145459404... The models are also compatible with many of Lucidrain's popular repositories such as Toolformer-pytorch, PaLM-rlhf-pytorch, and PaLM-pytorch. Please be sure to sponsor and help support Phil's great work:…
Read the full comment
https://github.com/lucidrains/PaLM-rlhf-pytorch You can find the weights on Hugging Face if you prefer to download the PyTorch .pt files from there instead: https://huggingface.co/conceptofmind/palm-1b All of the C4 data has been pre-tokenized with the GPTNEOX tokenizer and blocked at sequence lengths of 8192. This will help to save you the large cost of preprocessing data. The datasets are available on Hugging Face. An example chunk can be found here: https://huggingface.co/datasets/conceptofmind/c4_0-to-20_neo... The original unprocessed C4 dataset used can be found on the Hugging Face hub here: https://huggingface.co/datasets/c4 If you would like to preprocess your own dataset for training there is a dataset builder script provided. This uses Hugging Face datasets to efficiently map, tokenize, and block the data: https://github.com/conceptofmind/PaLM/blob/main/build_datase... A distributed training script is provided so that you may train or fine-tune your own PaLM models using Hugging Face accelerate. More information and experiments about the training will be detailed in the repository: https://github.com/conceptofmind/PaLM/blob/main/train_distri... The models were trained with Flash Attention, Xpos Rotary Embeddings for better length extrapolation, and multi-query single-key-value attention for more efficient decoding: https://arxiv.org/abs/2205.14135 Further instruction-tuning will be done on the new FLAN datasets we have released. https://huggingface.co/datasets/conceptofmind/cot_submix_ori... You can find his twitter here: https://twitter.com/ShayneRedford A basic inference script was provided in the repository which you can play around with. You may want to experiment with hyperparameters in order to get generations of varying quality. Changing a variable such as temperature matters a lot: https://github.com/conceptofmind/PaLM/blob/main/inference.py Different inference optimizations such as Flash Attention, Hidet, and Torch compile are used. You can read more about the Hidet compiler and project here: https://pytorch.org/blog/introducing-hidet/ Our work on Toolformer, PaLM, and related projects is all thanks to the generous sponsorship by CarperAI and StabilityAI: https://github.com/CarperAI This is not an official Google or StabilityAI product.
Original comment HN discussion #35874554
dpflan
Neutral mention · rule
Sharing a comment from a similar line of discussion: """ Complexity is quadratic in sequence length. For 512 tokens it is 262K, but for 4000 tokens it becomes 16M and goes OOM on a single GPU. We need about 100K-1M tokens to load whole books at once. Since 2017 there have been hundreds of attempts to bring O(N^2) to O(N), but none of them replaced the vanilla attention yet in large models. They lose on accuracy. Maybe Flash attention has a shot (https://arxiv.org/abs/2205.14135). """ - source: https://news.ycombinator.com/item?id=34171503
Original comment The Coming of Local LLMs
visarga
Neutral mention · rule
Complexity is quadratic in sequence length. For 512 tokens it is 262K, but for 4000 tokens it becomes 16M and goes OOM on a single GPU. We need about 100K-1M tokens to load whole books at once. Since 2017 there have been hundreds of attempts to bring O(N^2) to O(N), but none of them replaced the vanilla attention yet in large models. They lose on accuracy. Maybe Flash attention has a shot (https://arxiv.org/abs/2205.14135).
Original comment GPT3/DALL-E2 in Discord, chat like ChatGPT, generate images, and more

Keep following the thread.

More in Language models

More from the same topics.