FlashAttention: The Algorithm That Made Long-Context AI Possible

Standard attention, the Q-K-V formula at the heart of transformers, has a memory cost that grows as the square of sequence length. A 4,096-token sequence requires 16 million entries in the attention matrix. A 128,000-token sequence requires 16.4 billion. For most of 2020 to 2022, that quadratic wall limited context windows to a few thousand tokens. FlashAttention, a 2022 algorithm by Tri Dao and colleagues at Stanford, erased that wall without changing what the computation produces. The answer was hiding in the GPU's memory hierarchy.

The quadratic bottleneck in standard attention

Standard attention computes a score matrix S = QK^T and then an output O = softmax(S) * V. The score matrix is N by N, where N is the sequence length. On a 4,096-token sequence, that matrix holds 16.7 million floating-point numbers. A 128K-token sequence produces 16.4 billion entries, well beyond the 80 GB of high-bandwidth memory (HBM) on an A100 GPU. The naive implementation writes every element of S to HBM, reads it back to compute the softmax, then reads it again to compute the output. On long sequences, most of the wall-clock time is not arithmetic. It is waiting for memory.

What tiling changes

FlashAttention was introduced in the paper 'FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness' by Tri Dao, Daniel Fu, Stefano Ermon, Atri Rudra, and Christopher Re at Stanford (arXiv:2205.14135, accepted at NeurIPS 2022). The paper reorders the attention computation using tiling. Instead of computing the full N by N score matrix at once, it splits Q, K, and V into small blocks that fit inside the GPU's on-chip SRAM. Each block is loaded from HBM, attention is computed entirely in SRAM, and only the result is written back. The large intermediate matrix never exists in HBM at all.

The math requires maintaining a running maximum and running sum of exponential values across blocks so that softmax can be computed correctly without seeing all scores simultaneously. This online softmax trick makes the result numerically identical to standard attention. FlashAttention is exact, not approximate. The output is bit-for-bit the same as the naive implementation, just produced with far fewer HBM reads and writes.

The speedups the original paper measured

The NeurIPS 2022 paper reported a 3x wallclock speedup on GPT-2 training at a 1,024-token sequence length and a 2.4x speedup on the Long Range Arena benchmark at sequence lengths from 1,024 to 4,096 tokens. Memory use dropped from O(N squared) to O(N), because the N by N attention matrix is never materialized in HBM.

Quality improved alongside speed. In the GPT-2 experiments, FlashAttention produced a 0.7-point perplexity improvement over standard attention at the same training budget, because faster training allowed more iterations in the same time. Long-document classification benchmarks improved by 6.4 points. These gains came not from a different model but from a more efficient implementation of exactly the same computation.

Version 2 in 2023 and version 3 in 2024

FlashAttention 2, released in 2023, revised the GPU thread scheduling to reduce wasted compute and better exploit GPU parallelism. It achieved up to 70 percent of the theoretical peak FLOP throughput on an A100 GPU. It also added better support for multi-query attention and grouped-query attention, the variants used in most production LLMs.

FlashAttention 3, published July 11, 2024 (arXiv:2407.08608), targeted the Hopper architecture in NVIDIA's H100. FlashAttention 2 was using only 35 percent of the H100's theoretical maximum, because its code did not exploit the H100's newer instruction set. FlashAttention 3 introduced warp-specialization to overlap data movement and matrix multiplication, a pingpong scheduling scheme to run the softmax exponential in parallel with matrix multiplications, and incoherent processing for FP8 low-precision to reduce quantization error. The result is 1.5 to 2.0 times faster than FlashAttention 2 in FP16, reaching up to 740 TFLOPS, or 75 percent of the H100's theoretical maximum. In FP8 mode it approaches 1.2 petaFLOPS per second.

Why context lengths exploded after 2022

FlashAttention is now the default attention backend in PyTorch (via torch.nn.functional.scaled_dot_product_attention, introduced in PyTorch 2.0) and in HuggingFace Transformers, along with most major training frameworks. Its adoption tracks directly with the jump in context lengths over the past three years. GPT-3 and early OPT models used context windows of 2,048 to 4,096 tokens. GPT-4 expanded to 128,000 tokens. Llama 3 variants have been extended beyond 1,000,000 tokens in research configurations.

Without FlashAttention, a 128K-token context would require a full N by N attention matrix consuming hundreds of gigabytes, far exceeding any single GPU. Memory efficiency was the necessary prerequisite for the long-context push: the positional encoding changes, the dataset curation, the sliding-window architectural patterns. FlashAttention made all of them practical to test on real hardware.

What this means if you are training or fine-tuning today

If you are using PyTorch 2.0 or later with a compatible GPU, you are already benefiting from FlashAttention through PyTorch's scaled dot-product attention dispatcher, with no extra installation required in most frameworks.

• For sequence lengths above a few hundred tokens, verifying that FlashAttention is active will cut memory use substantially compared to a naive attention implementation.
• Training throughput typically improves 2 to 3 times over a naive implementation at long sequence lengths.
• The reference implementation is on GitHub at Dao-AILab/flash-attention, with custom CUDA kernels that expose the full speedup for training loops.
• For inference on H100 hardware, FlashAttention 3 is the version to use for maximum FLOP utilization.

FlashAttention is one of the clearest examples in recent AI history of a systems insight, rather than a new architecture or a larger training run, producing a step-change in what is possible. Tri Dao and his collaborators identified memory bandwidth as the true bottleneck in attention, not arithmetic, and designed the algorithm around that constraint. That insight is now baked into the infrastructure of nearly every serious LLM in production. PromptingIndex covers the hardware and systems layer as well as the model layer, because these engineering decisions increasingly determine what you can actually build.