PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2205.14135v2cs.LG

TL;DR

Long-context Transformers are constrained by self-attention’s quadratic time and memory costs. FlashAttention uses IO-aware tiling to reduce GPU memory traffic, requiring up to 9× fewer HBM accesses while supporting higher-quality, longer-context models.

  • Problem

    Long-context Transformers face quadratic self-attention time and memory costs, motivating faster and more memory-efficient attention.

  • Method

    FlashAttention computes exact attention with IO-aware tiling and recomputation to reduce reads and writes between GPU HBM and on-chip SRAM.

  • Results

    Up to 9× fewer HBM accesses than standard attention, with no exact algorithm asymptotically improving on FlashAttention across all SRAM sizes.

  • Takeaways & Limitations

    FlashAttention supports longer-context Transformers that improve model quality and achieve better-than-chance performance on Path-X and Path-256.

  • Takeaways & Limitations

    Implementing IO-aware attention requires substantial CUDA engineering effort and may not transfer across GPU architectures.

Abstract

from arXiv · show

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).

1 Introduction

FlashAttention introduces an IO-aware exact attention algorithm that reduces GPU HBM accesses through tiling, while extending to block-sparse attention for faster long-context modeling. The paper reports faster training, improved quality, and scaling to sequences up to 64K.

  • FlashAttention: FlashAttention computes exact attention while avoiding materialization of the large attention matrix in HBM by using tiling and online softmax reduction.It also avoids storing the large intermediate attention matrix for the backward pass, reducing memory accesses.
  • IO Complexity: FlashAttention requires O(N^2d^2M^-1) HBM accesses versus Ω(Nd + N^2) for standard attention, with up to 9× fewer accesses and an asymptotic lower bound for exact attention.Here d is the head dimension and M is SRAM size; the algorithm is optimal over a range of SRAM sizes.
  • Block-Sparse FlashAttention: Block-sparse FlashAttention applies the IO-aware primitive to approximate attention, achieving 2–4× speedups over FlashAttention while scaling to sequence length 64K.The approach addresses memory-access overhead that limits many sparse and low-rank approximate attention methods.
  • Faster Model Training: FlashAttention trains BERT-large 15% faster, GPT-2 3× faster, and long-range arena models 2.4× faster than baseline implementations.The reported settings use sequence lengths 512, 1K, and 1K–4K, respectively.
  • Higher Quality Models: Longer-context modeling improves GPT-2 perplexity by 0.7, lifts long-document classification by 6.4 points, and enables better-than-chance Path-X performance at sequence length 16K.Block-sparse FlashAttention extends these capabilities to sequence length 64K, where Path-256 reaches 63.1% accuracy.
  • Benchmarking Attention: FlashAttention is up to 3× faster than standard attention for sequence lengths 128–2K, scales to 64K, and block-sparse FlashAttention is faster than existing approximate methods.Through sequence length 512, FlashAttention is faster and more memory-efficient than existing attention methods; beyond 1K, some approximate methods become faster.

2 Background

GPU performance depends on the balance between computation and memory access across a hierarchy where smaller memories are faster but much smaller. Standard attention is slow and memory-intensive because it materializes O(N^2)-sized intermediates in HBM and repeatedly accesses them.

  • GPU Memory Hierarchy: GPU memory is hierarchical: A100 HBM provides 40–80GB at 1.5–2.0TB/s, while on-chip SRAM provides 192KB per multiprocessor at about 19TB/s.SRAM is roughly an order of magnitude faster than HBM but many orders of magnitude smaller.
  • Performance characteristics: Arithmetic intensity distinguishes compute-bound operations from memory-bound operations, whose runtime is dominated by memory accesses rather than computation.Elementwise operations and reductions such as softmax, dropout, and normalization are typical memory-bound examples.
  • Kernel fusion: Kernel fusion reduces HBM traffic by loading shared inputs once, but training still requires intermediate values to be written to HBM for backward passes.This requirement limits the effectiveness of naive fusion despite compiler support for fusing many elementwise operations [53] [65].
  • Standard attention: Standard attention materializes S and P in HBM, requiring O(N^2) memory and many memory accesses that slow wall-clock time, especially when N ≫ d.The implementation writes S, rereads it for softmax, writes P, then rereads P with V to produce O.

3 FlashAttention: Algorithm, Analysis, and Extensions

FlashAttention computes exact attention with sub-quadratic HBM accesses by combining tiling, online softmax normalization, recomputation, and fused GPU kernels. Its IO analysis establishes asymptotic optimality over SRAM sizes, while block-sparse FlashAttention reduces IO proportionally to sparsity.

  • Algorithm: Tiling loads Q, K, and V blocks into SRAM, combines blockwise softmax results using normalization statistics, and avoids repeated HBM transfers through kernel fusion.Recomputation reconstructs attention matrices during the backward pass instead of storing them.
  • Algorithm: FlashAttention returns exact attention with O(N^2d) FLOPs and only O(N) additional memory beyond inputs and outputs.It avoids storing O(N^2) intermediate matrices by retaining the output and softmax statistics for backward-pass recomputation.
  • IO Complexity: FlashAttention requires Θ(N^2d^2M^-1) HBM accesses versus Θ(Nd + N^2) for standard attention, often yielding many fewer accesses for d=64–128 and M around 100KB.The reduction in HBM accesses explains faster execution and lower memory footprint despite higher FLOP counts from backward recomputation.
  • IO Complexity: No exact attention algorithm can achieve o(N^2d^2M^-1) HBM accesses for every SRAM size M in [d, Nd].Thus FlashAttention is IO-optimal across the stated SRAM range, although parameterized lower bounds in M remain open.
  • Block-Sparse Extension: Block-sparse FlashAttention skips zero attention blocks and requires Θ(Nd + N^2d^2M^-1s) HBM accesses, where s is the fraction of nonzero blocks.Its IO complexity improves the larger term by the sparsity factor s.

4 Experiments

FlashAttention improves Transformer training speed and enables longer contexts without changing model quality, while reducing attention runtime and memory use. These gains support higher long-context quality and better-than-chance performance on Path-X and Path-256.

  • Training Speed: 15% faster BERT-large training, up to 3× faster GPT-2 training, and 2.4× faster Long-Range Arena training are achieved with FlashAttention.The GPT-2 comparison is up to 3× versus HuggingFace and up to 1.7× versus Megatron-LM; BERT reaches the same target accuracy from identical initialization.
  • Experimental Validation: FlashAttention preserves baseline perplexity and training curves because it does not change the model definition, while the reproduced Long-Range Arena baselines depend strongly on tuning.The Long-Range Arena experiments use sequence lengths from 1,024 to 4,096 and report accuracy, throughput, and training time.
  • Language Modeling with Long Context: 30% faster GPT-2 training with a 4K context than Megatron-LM with a 1K context also yields 0.7 better perplexity.FlashAttention increases GPT-2 context length by 4× while remaining faster than the optimized Megatron-LM implementation.
  • Long Document Classification: Longer sequences improve document classification by 4.3 points on MIMIC-III at 16K and 8.5 points on ECtHR at 8K, versus length 512.The datasets contain very long documents, with maximum lengths of 14,562 tokens for MIMIC-III and 49,392 for ECtHR.
  • Path-X and Path-256: FlashAttention reaches 61.4% accuracy on Path-X, while block-sparse FlashAttention reaches 63.1% on Path-256 at sequence length 64K.These are reported as the first Transformer results achieving non-random performance on both challenging long-context tasks.
  • Benchmarking Attention: FlashAttention is up to 3× faster than PyTorch exact attention and up to 20× more memory-efficient, with memory growing linearly in sequence length.Block-sparse FlashAttention has the same memory footprint, while its runtime also scales linearly and exceeds existing approximate-attention baselines.

5 Limitations and Future Directions … B.3 FlashAttention: Forward Pass

The paper develops memory-efficient forward and backward attention computations and FlashAttention’s tiled GPU implementation, while identifying limitations and extensions for IO-aware deep learning and multi-GPU systems. It also situates the approach among efficient Transformer, structured-matrix, sparse-training, and runtime-optimization research.

  • 5 Limitations and Future Directions: The approach’s main limitation is engineering complexity: each new IO-aware attention implementation requires a new low-level CUDA kernel that may not transfer across GPU architectures.The authors propose high-level languages such as PyTorch as a future direction for expressing attention algorithms.
  • 5 Limitations and Future Directions: The authors propose extending IO-aware optimization beyond attention because every deep-learning layer accesses GPU HBM, and they identify multi-GPU data transfer as an additional IO-analysis problem.They note that their attention implementation is optimal within constants on a single GPU, while multi-GPU parallelization requires accounting for inter-GPU transfers.
  • A Related Work: The work connects IO-aware optimization to I/O complexity, working-set models, data locality, Roofline arithmetic intensity, and scalability analyses.These connections frame memory-hierarchy optimization as a longstanding systems principle rather than an attention-specific idea.
  • A Related Work: The paper relates its efficient attention and block-sparse variant to structured matrices, sparse training, and alternative long-context Transformer modules addressing quadratic sequence-length costs.Examples include sparse and low-rank matrices, pruning and lottery-ticket training, and models such as Reformer [51], Performer [12] [54], S4 [31] [36] [37], and FLASH [42].
  • B Algorithm Details: FlashAttention derives memory-efficient forward and backward attention with linear extra memory, reducing HBM accesses to improve runtime and memory footprint.Naive memory-efficient formulations avoid quadratic intermediate storage but still incur quadratic HBM accesses; FlashAttention addresses those accesses directly.
  • B.1 Memory-efficient forward pass: The forward pass avoids storing the quadratic attention matrix by separately maintaining softmax normalization constants and repeatedly accumulating output contributions.This uses O(n) extra memory, with O(n) memory for normalization constants and O(d) extra memory for each output computation.
  • B.2 Memory-efficient backward pass: The backward pass explicitly computes dV, dQ, and dK with linear memory instead of relying solely on gradient checkpointing.The derivation proceeds through dP and dS, while storing only linear-size intermediates such as L and D.
  • B.3 FlashAttention: Forward Pass: FlashAttention’s forward pass tiles Q, K, and V between HBM and on-chip SRAM, computes masked and dropout-adjusted attention on chip, and saves softmax statistics for backpropagation.The algorithm partitions Q into row blocks and K,V into column blocks, then returns O, ℓ, m, and the random-number-generator state R.

B.4 FlashAttention: Backward Pass · B.5 Comparison with Rabe and Staats [66]

FlashAttention’s backward pass preserves exact gradients with O(N^2) FLOPs, O(N) extra memory, and fewer HBM accesses than standard attention. Compared with Rabe and Staats [66], it targets memory-access reduction and achieves greater speed while retaining substantial memory savings.

  • B.4 FlashAttention: Backward Pass: FlashAttention computes backward gradients blockwise on SRAM, reconstructing attention and dropout information while producing dQ, dK, and dV without materializing the full attention matrix.Its softmax-gradient computation avoids reductions over N-sized rows that may not fit in SRAM, using an equivalent reformulation.
  • B.4 FlashAttention: Backward Pass: O(N^2) FLOPs and O(N) extra memory characterize FlashAttention’s backward pass beyond its inputs, outputs, output gradient, and input gradients.The backward pass avoids storing an O(N^2) dropout mask by saving and replaying pseudo-random generator states.
  • B.4 FlashAttention: Backward Pass: Θ(N^2d^2M^-1) HBM accesses for FlashAttention backward versus Θ(Nd + N^2) for standard attention, when d ≤ M ≤ Nd.This is the stated IO-complexity comparison for sequence length N, head dimension d, and SRAM size M.
  • B.5 Comparison with Rabe and Staats [66]: FlashAttention and Rabe and Staats [66] both tile attention blocks, avoid storing the large forward attention matrix, and recompute it during the backward pass.Both methods use tiling or softmax scaling [51] [60].
  • B.5 Comparison with Rabe and Staats [66]: FlashAttention reduces memory accesses, whereas Rabe and Staats [66] primarily reduces maximum GPU-memory footprint; memory accesses are identified as the primary runtime determinant.Reducing memory accesses also necessarily reduces total memory requirements.
  • B.5 Comparison with Rabe and Staats [66]: 2-4× faster than standard attention, FlashAttention contrasts with Rabe and Staats [66] at around standard-attention speed or slightly slower, while both save substantial memory.The comparison attributes FlashAttention’s speed advantage to reducing memory accesses rather than only total memory footprint.
  • B.5 Comparison with Rabe and Staats [66]: FlashAttention incrementally updates block outputs, requiring less total memory than Rabe and Staats [66], which stores temporary outputs and later combines them using normalization statistics.This difference concerns how information is summarized and propagated across attention blocks.
  • B.5 Comparison with Rabe and Staats [66]: FlashAttention analytically simplifies backward computation, recomputing only the attention matrix rather than each block’s temporary output, which reduces memory requirements and yields speedup over Rabe and Staats [66].Rabe and Staats [66] instead uses gradient checkpointing to recompute both quantities.

C Proofs · D Extension Details · D.1 Block-sparse FlashAttention

The proofs establish that FlashAttention computes exact attention while reducing HBM accesses for forward and backward passes, and its block-sparse extension scales accesses with the fraction of nonzero blocks.

  • C Proofs: The forward algorithm is exact: after the final outer-loop iteration, its output equals softmax(QK^T)V.Correctness follows by induction over key-value blocks while maintaining row maxima, exponential row sums, and the accumulated output.
  • C Proofs: FlashAttention’s forward pass requires Θ(NdT_c) HBM accesses because each K and V element loads once while Q and O are revisited across T_c passes.Standard attention instead requires Θ(Nd + N^2) global-memory accesses.
  • C Proofs: The IO bounds depend on choosing K/V, Q/O, and score blocks that fit within on-chip SRAM.The proof separately derives constraints for blocks of sizes B_c×d, B_r×d, and B_r×B_c.
  • C Proofs: The exact-attention IO lower bound follows because Q, K, V, and O collectively require at least Ω(Nd) HBM accesses.This contradicts any proposed exact algorithm whose access count would fall below that input-output cost in the M = Θ(Nd) regime.
  • C Proofs: FlashAttention’s backward pass also requires Θ(NdT_c) HBM accesses, loading K and V once and writing dK and dV once.Standard attention backward requires Θ(Nd + N^2) HBM accesses.
  • D Extension Details: The block-sparse extension preserves FlashAttention’s tiled computation while loading only blocks selected by the nonzero sparsity mask.Algorithm 5 tests each mask block and performs on-chip computation only when M_ij ≠ 0.
  • D.1 Block-sparse FlashAttention: Block-sparse FlashAttention skips zero blocks, reducing HBM accesses proportionally to the nonzero-block fraction s while still writing the N×d output.The algorithm is otherwise identical to FlashAttention, and the IO reduction is limited by the cost of writing O when s is small.

D.2 Potential Extensions

The IO-aware approach suggests extensions beyond single-GPU attention, including multi-GPU cooperation, sparse MLP optimization, and kernel machine learning. These directions exploit memory-hierarchy asymmetries or repeated computation from low-dimensional inputs to reduce memory or computational costs.

  • Multi-GPU Attention: Multi-GPU attention could coordinate GPUs on the same node for very long sequences while accounting for multiple memory-hierarchy levels [77].The hierarchy includes GPU SRAM, local GPU HBM, and the HBM of other GPUs; attention is typically split across 4–8 GPUs.
  • Sparse MLP layers: IO-aware implementations could make sparse MLP layers more efficient when memory traffic prevents speedups proportional to sparsity [17].The proposed direction aims to reduce the computational requirements of large models.
  • Kernel machine learning: Kernel machine learning presents a similar opportunity because each N×N kernel entry depends on two vectors of dimension d ≪ N, enabling repeated input loads and recomputation.FlashAttention uses the analogous low-rank structure of QK^⊤ to reduce HBM access by recomputing needed attention blocks.

E Full Experimental Results · E.1 BERT

The BERT-large experiment follows MLPerf 1.1 procedures and compares training speed against Nvidia’s reported submission under matched evaluation conditions. Training uses FP16 on 8×A100-80GB GPUs, with runtime averaged across 10 runs.

  • E.1 BERT: BERT-large training follows the reference MLPerf 1.1 procedure and hyperparameters, including LAMB optimization, batch size 448, and at most 7100 steps.The learning rate is 3.75e-3.
  • E.1 BERT: Training stops when masked-language-modeling validation accuracy reaches 72.0%, after which wall-clock runtime is measured.Runs use FP16 precision with Apex AMP at O2 optimization level.
  • E.1 BERT: The experiment compares its training speed with Nvidia’s reported result submitted to MLPerf 1.1.The comparison is presented in Table 1.
  • E.1 BERT: Evaluation uses the same train/validation split as the MLPerf 1.1 reference implementation, including the same 10,000 validation examples as Nvidia’s baseline.
  • E.1 BERT: The model is trained on 8×A100-80GB GPUs.
  • E.1 BERT: Each training run takes 16–19 minutes, and reported results average 10 runs.

E.2 GPT-2 · E.3 LRA details

GPT-2 experiments used standard Hugging Face and Megatron-LM implementations under matched training conditions, while LRA comparisons followed published settings and found similar tuned accuracy across tasks. FlashAttention matched baseline GPT-2 validation perplexity curves, and evaluation procedures accounted for reproducibility and implementation constraints.

  • E.2 GPT-2: GPT-2 used standard Hugging Face and Megatron-LM implementations, following the Megatron-LM training recipe.
  • E.2 GPT-2: GPT-2 models used a 512 effective batch size, AdamW, model-specific learning rates, weight decay 0.1, 400K steps, and mixed-precision training.The learning rates were 6e-4 for GPT-2 small and 1.5e-4 for GPT-2 medium.
  • E.2 GPT-2: GPT-2 training used OpenWebText with the GPT-2 BPE tokenizer and a fixed random 0.5% validation split shared across models.
  • E.2 GPT-2: GPT-2 wall-clock training was measured on 8×A100-40GB GPUs, taking 2.7–9.5 days for small and 6.9–21.0 days for medium models.The reported durations correspond to Table 2.
  • E.2 GPT-2: FlashAttention produced validation perplexity curves nearly identical to Hugging Face baselines for GPT-2 small and medium.The comparison appears in Figure 4 and indicates equivalent validation behavior between implementations.
  • E.3 LRA details: LRA experiments followed settings from the Long-range arena paper and repository and the Nyströmformer reproduction [80].When reproductions underperformed, the better reported baseline result from Tay et al. [80] or Xiong et al. was used.
  • E.3 LRA details: Almost all attention methods achieved similar accuracy across all five LRA tasks after hyperparameter tuning.
  • E.3 LRA details: LRA overall wall-clock speedup was computed as the geometric mean across the five task-specific speedups.Mixed precision was used except for Performer, which was unstable, and Local Attention, whose implementation lacked FP16 support.

E.4 Comparison with Apex FMHA

FlashAttention builds on Apex FMHA with tiling and recomputation to reduce memory use, support longer sequences and broader hardware, while achieving comparable short-sequence runtime. It is slightly faster in the forward pass but slightly slower in the backward pass because it recomputes rather than stores the attention matrix.

  • Comparison with Apex FMHA: Apex FMHA targets BERT models with head dimension 64 and supports only A100 GPUs and sequence lengths up to 512.It fuses dropout(softmax(mask(QK^T)))V into one CUDA kernel but stores the attention matrix in HBM for gradient computation, limiting memory savings.
  • Comparison with Apex FMHA: FlashAttention supports sequences up to 64K, head dimensions 16, 32, 64, and 128, and all Turing and Ampere GPUs described at writing.These extensions use tiling and recomputation to handle long sequences and save memory.
  • Comparison with Apex FMHA: FlashAttention is generally slightly faster than FMHA in the forward pass and slightly slower in the backward pass.The backward-pass slowdown results from not storing the attention matrix during the forward pass and recomputing it during backpropagation.
  • Comparison with Apex FMHA: FlashAttention is about 4% slower than FMHA at sequence length 128.The comparison uses masking and dropout on an A100-SXM4-40GB GPU with batch size 64, 16 heads, and head dimension 64.

E.5 Speedup On Different Hardware and Configurations

FlashAttention speedups vary across GPU hardware and configurations because HBM bandwidth and SRAM size affect IO efficiency. Across tested settings, it achieves substantial gains, with masking, dropout, and lower memory bandwidth often increasing speedup while larger head dimensions and smaller SRAM reduce it.

  • A100: 2-4× speedup over standard PyTorch attention appears on A100 across sequence lengths, increasing with dropout and masking through kernel fusion.The configuration uses batch size 8, head dimension 64, and 12 attention heads.
  • RTX 3090: 2.5-4.5× speedup occurs on RTX 3090, slightly exceeding A100 because its memory bandwidth is lower, roughly 900 GB/s versus 1.5 TB/s.The RTX 3090 configuration uses batch size 12 with 12 attention heads.
  • T4: T4 shows less speedup than A100 because its smaller SRAM requires smaller FlashAttention blocks, consistent with the IO complexity analysis.Because T4 GPUs are commonly used for inference, both combined forward-plus-backward and forward-only speedups are reported.

E.6 Full Benchmarking Results

The section benchmarks exact, approximate, and sparse attention implementations on an A100 across varied sequence lengths, reporting forward, backward, combined runtime, and memory usage under multiple conditions. Comparisons include FlashAttention-related baselines such as Reformer [51], Local Attention, Linformer [84], Smyrf, LongShortFormer (LSFormer), Block-Sparse Attention [11], Longformer [3], and BigBird.

  • E.6 Full Benchmarking Results: The benchmark compares exact, approximate, and sparse attention against reference implementations from PyTorch/HuggingFace, Megatron, Reformer [51], Local Attention, Linformer Attention [84], Smyrf, LongShortFormer (LSFormer), Block-Sparse Attention [11], Longformer [3], and BigBird.These baselines cover the main exact, approximate, and sparse attention categories.
  • E.6 Full Benchmarking Results: Experiments use 8 heads of dimension 64, batch size 16, random Q, K, and V vectors, and sequence lengths varied on one 40 GB A100 GPU.Attention projection from the hidden layer is excluded; dropout is 0.1, and masking uses uniformly random mask lengths.
  • E.6 Full Benchmarking Results: The results tables cover runtime by sequence length for unconditioned, masked, dropout, and combined dropout-and-masking settings, plus memory usage.Tables identify best and second-best methods with bold and underlined formatting, respectively.
  • E.6 Full Benchmarking Results: Runtime is reported for forward, backward, and combined passes with and without dropout, masking, or both, while memory is measured for combined forward and backward passes without either.Measurements use FP16 except Local Attention, whose implementation supports only FP32.
  • E.6 Full Benchmarking Results: Sequence lengths increase until GPU memory exhaustion, subject to implementation limits: Megatron 2048, Block-Sparse 4096, and Longformer and BigBird 8092.Block-Sparse, Longformer, and BigBird were measured without masked backward passes because external-library bugs prevented those runs.
Loading 2205.14135v2…