Source-linked AI summary
FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
Tri Dao
TL;DR
Longer Transformer contexts are limited by attention’s quadratic runtime and memory, while FlashAttention remains less efficient than optimized GEMM because of GPU work-partitioning issues. FlashAttention-2 reduces non-matmul work and improves parallelism across sequence positions and warps, achieving about 2× speedup over FlashAttention and up to 225 TFLOPs/s per A100 in GPT-style training.
Problem
Attention bottlenecks longer-context Transformers because its runtime and memory requirements grow quadratically with sequence length, while FlashAttention remains less efficient than optimized GEMM due to suboptimal GPU work partitioning.
Method
FlashAttention-2 reduces non-matmul FLOPs, parallelizes attention across sequence length, and partitions work between warps to improve occupancy and reduce shared-memory communication.
Results
Around 2× speedup over FlashAttention and up to 225 TFLOPs/s per A100 are achieved, including end-to-end training of GPT-style models.
Takeaways & Limitations
FlashAttention-2 enables training models with 16k context for the same price as previously training an 8k-context model.
Takeaways & Limitations
The work is manually tuned for each head dimension, and the authors leave auto-tuning to future work.
Abstract
from arXiv · showhide
Scaling Transformers to longer sequence lengths has been a major problem in the last several years, promising to improve performance in language modeling and high-resolution image understanding, as well as to unlock new applications in code, audio, and video generation. The attention layer is the main bottleneck in scaling to longer sequences, as its runtime and memory increase quadratically in the sequence length. FlashAttention exploits the asymmetric GPU memory hierarchy to bring significant memory saving (linear instead of quadratic) and runtime speedup (2-4$\times$ compared to optimized baselines), with no approximation. However, FlashAttention is still not nearly as fast as optimized matrix-multiply (GEMM) operations, reaching only 25-40\% of the theoretical maximum FLOPs/s. We observe that the inefficiency is due to suboptimal work partitioning between different thread blocks and warps on the GPU, causing either low-occupancy or unnecessary shared memory reads/writes. We propose FlashAttention-2, with better work partitioning to address these issues. In particular, we (1) tweak the algorithm to reduce the number of non-matmul FLOPs (2) parallelize the attention computation, even for a single head, across different thread blocks to increase occupancy, and (3) within each thread block, distribute the work between warps to reduce communication through shared memory. These yield around 2$\times$ speedup compared to FlashAttention, reaching 50-73\% of the theoretical maximum FLOPs/s on A100 and getting close to the efficiency of GEMM operations. We empirically validate that when used end-to-end to train GPT-style models, FlashAttention-2 reaches training speed of up to 225 TFLOPs/s per A100 GPU (72\% model FLOPs utilization).
1 Introduction
FlashAttention-2 addresses FlashAttention’s remaining GPU-efficiency bottlenecks through reduced non-matmul computation and improved work partitioning. These changes yield about 2× speedup over FlashAttention and up to 225 TFLOPs/s per A100 when training GPT-style models.
- Attention remains difficult to scale because its runtime and memory requirements grow quadratically with sequence length.
- FlashAttention reduces memory usage from quadratic to linear in sequence length and provides 2-4× wall-clock speedup over optimized baselines without approximation.
- FlashAttention reaches only 30-50% forward and 25-35% backward of theoretical A100 throughput because of low occupancy or unnecessary shared-memory reads and writes.
- FlashAttention-2 reduces non-matmul FLOPs while preserving the output, targeting operations that can run up to 16× slower than matrix-multiply FLOPs.
- FlashAttention-2 parallelizes attention across sequence length and partitions work between warps, increasing occupancy and reducing communication and shared-memory accesses.
- Around 2× speedup over FlashAttention and up to 225 TFLOPs/s per A100 are reported for FlashAttention-2, including end-to-end GPT-style model training.
2 Background
Attention computes pairwise interactions across sequence positions, making standard implementations expensive in memory and data movement. FlashAttention uses GPU tiling, online softmax, and recomputation to preserve exact outputs while reducing memory traffic and storage.
- Standard attention: Standard attention forms S = QK^⊤, P = softmax(S), and O = PV, with row-wise softmax and parallelism across heads and batches.The intermediate matrices S and P have quadratic size in sequence length.
- Standard attention: O(N^2) memory is required when standard implementations materialize S and P in HBM, and repeated memory accesses slow execution because many operations are memory-bandwidth bound.Typical settings have N around 1k–8k and d around 64–128.
- FlashAttention: FlashAttention loads input blocks into SRAM, computes blockwise attention, and updates outputs without writing the large intermediate matrices S and P to HBM.Tiling reduces memory reads and writes while maintaining the same output without approximation.
- FlashAttention: Online softmax computes local softmax values for each block and rescales partial outputs so that combining blocks produces the exact attention result.The row-wise maximum and normalization statistics are updated across blocks.
- FlashAttention: FlashAttention's backward pass recomputes S and P after loading Q, K, and V blocks into SRAM, avoiding storage of large intermediates and reducing memory from quadratic to linear in N.It provides 10–20× memory saving and 2–4× wall-clock speedup depending on sequence length.
- Backward pass: The backward implementation is more involved because it must retain more values in SRAM for five matrix multiplications, compared with two in the forward pass.Tiling is applied to the backward equations, which do not require softmax rescaling.
3 FlashAttention-2: Algorithm, Parallelism, and Work Partitioning
FlashAttention-2 refines FlashAttention’s algorithm and GPU work assignment to reduce non-matmul computation, increase occupancy, and reduce shared-memory communication while preserving exact attention outputs.
- Work partitioning between warps: Within each thread block, FlashAttention-2 partitions warp work to reduce communication and shared-memory reads and writes.This addresses the split-K scheme’s need for intermediate shared-memory writes, synchronization, and reduction.
- Algorithmic tweaks: FlashAttention-2 reduces non-matmul FLOPs without changing the output, prioritizing matmul work because A100 matmul throughput can be 16× higher.The stated A100 maxima are 312 TFLOPs/s for FP16/BF16 matmul and 19.5 TFLOPs/s for non-matmul FP32.
- Algorithmic tweaks: The online softmax update keeps an unscaled output and logsumexp statistics, applying final scaling only after processing all blocks.The backward pass stores only row-wise logsumexp rather than both row-wise maxima and exponential sums.
- Causal masking: Causal masking skips blocks that cannot contain valid entries and applies masking only where needed, yielding around 1.7-1.8× speedup over unmasked attention.For each row, masking is needed for only one block when square blocks are used.
- Correctness and complexity: The algorithm returns exact attention outputs with O(N^2d) FLOPs and O(N) additional memory beyond inputs and outputs.The additional memory stores the logsumexp vector.
- Parallelism: FlashAttention-2 parallelizes forward and backward attention over sequence length in addition to batch and heads, improving occupancy for long sequences with small batches or few heads.Forward row-block workers operate independently, while the backward scheme assigns workers to column blocks.
4 Empirical Validation
FlashAttention-2 substantially improves attention runtime across sequence lengths and settings, outperforming FlashAttention and standard implementations while delivering high A100 throughput in end-to-end GPT training.
- Attention benchmarks: FlashAttention-2 is 1.7-3.0× faster than FlashAttention, 1.3-2.5× faster than FlashAttention in Triton, and 3-10× faster than standard attention.Benchmarks vary sequence length and include causal-mask and head-dimension settings.
- Attention benchmarks: FlashAttention-2 reaches up to 230 TFLOPs/s, or 73% of theoretical maximum throughput, on A100 GPUs.The reported attention benchmarks include forward and backward measurements.
- End-to-end training: For GPT-style models with 1.3B or 2.7B parameters and 2k or 8k sequences, FlashAttention-2 yields up to 1.3× speedup over FlashAttention and 2.8× over a baseline without it.The end-to-end result reaches 225 TFLOPs/s per A100 GPU and 72% model FLOPs utilization.
- Attention benchmarks: In the forward pass, FlashAttention-2 is around 1.3-1.5× faster than FlashAttention in Triton and around 2× faster in the backward pass.It is also reported as around 2× faster than FlashAttention and FlashAttention in xformers.
- Additional hardware: On H100 GPUs, the same implementation reaches up to 335 TFLOPs/s without special instructions for newer hardware features.The authors expect further speedup from those instructions but leave that optimization to future work.
5 Discussion and Future Directions
FlashAttention-2 doubles attention speed relative to FlashAttention, enabling longer-context training at the same cost and supporting broader long-sequence applications. The discussion also points to extending the implementation across newer hardware and data types.
- 225 TFLOPs/s per A100 GPU is the reported GPT-style model training speed for FlashAttention-2.The table caption also reports 72% model FLOPs utilization.
- FlashAttention-2 is intended to support understanding long books and reports, high-resolution images, audio, and video.The discussion also identifies training, finetuning, and inference of existing models as applications it will speed up.
- Future work targets broader device and datatype support, including H100 and AMD GPUs and FP8.The authors specifically plan to optimize for H100 hardware features such as TMA, fourth-generation Tensor Cores, and FP8.