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

Self-attention’s quadratic time and memory requirements make long-context Transformers difficult to scale. FlashAttention computes exact attention using tiling and recomputation to reduce GPU memory transfers. It accelerates training across BERT-large, GPT-2, and long-range arena, while longer contexts improve model quality.

  • Problem

    Self-attention’s quadratic time and memory complexity makes equipping Transformers with longer context difficult, motivating faster and more memory-efficient attention.

    Source: §1 · p. 1

  • Method

    FlashAttention uses tiling and recomputation to compute exact attention while reducing transfers between GPU high bandwidth memory and on-chip SRAM.

    Source: §1 · p. 2, §3.1 · p. 4

  • Results

    FlashAttention accelerates training across BERT-large, GPT-2, and long-range arena, including 15% faster BERT-large training at sequence length 512 than the MLPerf 1.1 [58] record.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Longer contexts with FlashAttention and its block-sparse variant improve model quality and enable better-than-chance Transformer performance on Path-X and Path-256.

    Source: §1 · p. 3

    Each new attention implementation requires a custom CUDA kernel, demanding significant engineering effort and potentially lacking portability across GPU architectures.

    Source: §5 · p. 10

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 computes exact attention using tiling and recomputation to reduce GPU memory transfers and achieve linear memory use. It accelerates Transformer training and enables longer-context models with improved quality, while its block-sparse extension supports sequences of 64K.

Motivation

  • Quadratic self-attention limits longer contexts, while many approximate methods reduce computation without wall-clock speedups because they neglect memory-access overhead.These approaches include sparse approximations [51], low-rank approximations [12] [50] [84], and combinations [3]. The proposed IO-aware principle explicitly accounts for transfers between fast on-chip SRAM and slower HBM.Source: §1 · p. 1

Exact Attention with Reduced IO

  • FlashAttention computes exact attention with linear memory by tiling inputs, incrementally computing softmax, and recomputing attention during backpropagation instead of storing its matrix in HBM.Saved softmax normalization factors enable on-chip backward recomputation, while a fused CUDA kernel controls memory access despite additional FLOPs. Avoiding attention-matrix transfers delivers up to 7.6× attention-computation speedup over PyTorch on GPT-2.Source: §1 · p. 2, Figure 1 · p. 2

IO Complexity

  • Up to 9× fewer HBM accesses than standard attention support FlashAttention’s IO-efficiency advantage for typical head dimensions and SRAM sizes.FlashAttention requires O(N^2 d^2 M^-1) HBM accesses versus standard attention’s Ω(Nd + N^2), where d is head dimension and M is SRAM size. A lower bound rules out asymptotic improvement by any exact attention algorithm across all SRAM sizes.Source: §1 · p. 2

Block-Sparse Extension

  • 2–4× faster than FlashAttention, block-sparse FlashAttention scales to 64K and outperforms all existing approximate attention methods known to the authors.Its IO complexity improves over FlashAttention by a factor proportional to the sparsity ratio, demonstrating how the primitive can reduce approximate attention’s memory-access overhead.Source: §1 · pp. 2–3

Faster Model Training

  • 3× faster GPT-2 training than HuggingFace and Megatron-LM [77] demonstrates FlashAttention’s end-to-end training benefits.At sequence length 512, BERT-large trains 15% faster than the MLPerf 1.1 training speed record [58]. GPT-2 uses length 1K; long-range arena training at lengths 1K–4K is 2.4× faster than baselines.Source: §1 · p. 3

Higher Quality Models

  • 0.7 better GPT-2 perplexity and 6.4 points of lift on long-document classification [13] result from modeling longer contexts with FlashAttention.FlashAttention enables the first Transformer above chance on Path-X using length 16K. Block-sparse FlashAttention enables the first model above chance on Path-256 using length 64K.Source: §1 · p. 3

Benchmarking Attention

  • Up to 3× faster than standard attention at lengths 128–2K, FlashAttention is faster and more memory-efficient than every existing attention method through length 512.FlashAttention scales to length 64K, but beyond 1K some approximate methods, such as Linformer, begin to become faster.Source: §1 · p. 3

2 Background

GPU attention performance depends on memory traffic as well as computation: fast on-chip SRAM is scarce, while standard attention stores quadratic-sized intermediates in HBM. Kernel fusion reduces repeated reads, but saving intermediates for training limits its effectiveness.

2.1 Hardware Performance

  • An order of magnitude faster than HBM, on-chip SRAM is many orders of magnitude smaller, making efficient use of this scarce memory important.The A100 has 40–80 GB of HBM at 1.5–2.0 TB/s and 192 KB of SRAM per streaming multiprocessor across 108 multiprocessors, with SRAM bandwidth estimated around 19 TB/s [45].Source: §2.1 · p. 3
  • Arithmetic intensity, the number of arithmetic operations per byte accessed, distinguishes compute-bound operations limited by computation from memory-bound operations limited by memory access.Large-inner-dimension matrix multiplication and convolutions with many channels are typically compute-bound; elementwise operations and reductions, including softmax, are typically memory-bound.Source: §2.1 · p. 3
  • Kernel fusion accelerates memory-bound operations by loading shared inputs once, but training still requires writing intermediate values to HBM for the backward pass.These saved intermediates reduce the effectiveness of naive kernel fusion.Source: §2.1 · pp. 3–4

2.2 Standard Attention Implementation

  • O(N^2) memory and quadratic HBM accesses burden standard attention, which materializes S and P in HBM and includes memory-bound operations such as softmax.It computes S = QK^T, P = softmax(S), and O = PV, writing each result to HBM. Masking and dropout exacerbate the memory-access problem.Source: §2.2 · p. 4

3 FlashAttention: Algorithm, Analysis, and Extensions

FlashAttention uses tiling and recomputation to compute exact attention with fewer HBM accesses and without quadratic intermediate storage. Its IO analysis explains these savings, while a block-sparse extension further reduces memory traffic.

3.1 An Efficient Attention Algorithm With Tiling and Recomputation

  • Recomputation avoids storing O(N^2) intermediate values by reconstructing attention matrices during backpropagation from stored outputs and softmax statistics.The matrices are reconstructed from Q, K, V blocks in SRAM; reduced HBM accesses accelerate the backward pass despite additional FLOPs.Source: §3 · p. 4, §3.1 · p. 5
  • O(N) additional memory and O(N^2 d) FLOPs suffice for FlashAttention to compute exact attention, softmax(QK^T)V.Tiling loads Q, K, V blocks into SRAM and combines their outputs using softmax normalization statistics. One fused CUDA kernel performs the attention operations, including optional masking and dropout, avoiding repeated HBM transfers.Source: §3.1 · pp. 4–5

3.3 Extension: Block-Sparse FlashAttention

  • 2.8× speedup on LRA is achieved by block-sparse FlashAttention while performing on par with standard attention.Figure 2 also shows runtime improving proportionally as sparsity increases.Source: Figure 2 · p. 6, §3.3 · p. 7

3.2 Analysis: IO Complexity of FlashAttention

  • Reduced HBM traffic makes FlashAttention faster despite additional FLOPs, but increasing block size beyond 256 exposes other runtime bottlenecks.Larger blocks reduce passes over the inputs until factors such as arithmetic dominate; SRAM capacity also limits block size.Source: Figure 2 · p. 6, §3.2 · p. 6
  • Θ(N^2 d^2/M) HBM accesses replace standard attention’s Θ(Nd + N^2), for sequence length N, head dimension d, and SRAM size d ≤ M ≤ Nd.For typical head dimensions of 64–128 and SRAM around 100KB, d^2 is much smaller than M, yielding many times fewer HBM accesses.Source: §3.2 · p. 6
  • No exact-attention algorithm achieves o(N^2 d^2/M) HBM accesses for every SRAM size M in [d, Nd].The proof uses M = Θ(Nd); parameterized complexity lower bounds in M remain future work.Source: §3.2 · p. 6

3.3 Extension: Block-Sparse FlashAttention

  • Θ(Nd + N^2 d^2 s/M) HBM accesses suffice for block-sparse FlashAttention, where s is the nonzero-block fraction and d ≤ M ≤ Nd.The extension skips zero blocks in a predefined block-sparsity mask, reducing the larger IO term proportionally to s.Source: §3.3 · pp. 6–7
  • N^-1/2 [11] and N^-1 log N [3, 17, 92] are commonly chosen nonzero-block fractions s for long sequences.Downstream experiments use the fixed butterfly sparsity pattern [17], which can approximate arbitrary sparsity [16].Source: §3.3 · p. 7

4 Experiments

FlashAttention accelerates Transformer training while preserving GPT-2 perplexity, and enables longer contexts that improve model quality. Attention benchmarks show faster execution and linear memory growth, with block-sparse FlashAttention supporting sequences up to 64K.

BERT.

  • 15% faster BERT-large [22] training than Nvidia’s MLPerf 1.1 [58] speed record is achieved using FlashAttention.Training starts from the same initialization and targets 72.0% masked-language-modeling accuracy; results average 10 runs on 8×A100 GPUs.Source: Table 1 · p. 7, BERT · p. 7

GPT-2.

  • Up to 3× speedup over HuggingFace and 1.7× over Megatron-LM [77] accelerates GPT-2 training without changing perplexity.Results cover GPT-2 small and medium on 8×A100 GPUs. Matching training and validation curves support numerical stability relative to the baselines.Source: GPT-2 · pp. 7–8, Table 2 · p. 8
  • Up to 2.4× speedup over standard attention is achieved on LRA, while block-sparse FlashAttention outpaces every tested approximate attention method.Tasks span sequence lengths 1024–4096. Accuracy depends strongly on tuning, and the reproduced baselines outperform those in the original comparison.Source: Table 3 · p. 8, GPT-2 · p. 8, §4.2 · p. 8

4.2 Better Models with Longer Sequences

  • 30% faster GPT-2 training with FlashAttention at 4K context than Megatron-LM at 1K context also delivers 0.7 better perplexity.The comparison uses GPT-2 small and reports training time on 8×A100 GPUs.Source: Table 4 · p. 8, §4.2 · p. 8

4.3 Benchmarking Attention

  • Up to 3× faster than PyTorch for forward-plus-backward attention, FlashAttention retains quadratic runtime; block-sparse FlashAttention outpaces all compared exact, sparse, and approximate implementations.Approximate attention runtimes begin crossing below FlashAttention between sequence lengths 512 and 1024. Block-sparse FlashAttention runtime scales linearly with sequence length.Source: §4 · p. 7, Figure 3 · p. 9, §4.3 · p. 10
  • Up to 20× greater memory efficiency than exact attention baselines accompanies linear memory growth for FlashAttention and block-sparse FlashAttention, which share the same footprint.FlashAttention is more memory-efficient than approximate baselines and 2× more efficient than Linformer. All compared baselines except Linformer exhaust A100 memory before sequence length 64K.Source: Figure 3 · p. 9, §4.3 · p. 10

4.2 Better Models with Longer Sequences

  • 4.3-point micro F1 improvement on MIMIC-III at 16K tokens and 8.5-point improvement on ECtHR [6] [7] at 8K compare against 512-token sequences.The experiments extend pretrained RoBERTa [56] by repeating positional embeddings, following Beltagy et al. [3]. The authors suggest differing sensitivity to document-length distribution shifts may explain the unequal gains.Source: §4.2 · pp. 8–9, Table 5 · p. 9
  • 61.4 accuracy on Path-X with FlashAttention and 63.1 accuracy on Path-256 with block-sparse FlashAttention establish the first above-random Transformer results on these tasks.Path-X uses sequence length 16K, while block-sparse FlashAttention enables 64K on Path-256. Transfer to Path-X uses Path-64 pretraining and spatial interpolation of positional embeddings.Source: §4 · p. 7, Table 6 · p. 9, §4.2 · p. 9

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

The appendices derive linear-extra-memory attention passes and describe FlashAttention’s tiled GPU implementation, distinguishing reduced memory usage from reduced HBM traffic. Limitations include handwritten CUDA kernels and the need to extend IO awareness beyond single-GPU attention.

5 Limitations and Future Directions

  • Handwritten CUDA kernels impose engineering and portability costs, motivating high-level compilation, IO-aware modules beyond attention, and multi-GPU methods accounting for inter-GPU transfers.Each new attention implementation requires a new kernel; a proposed high-level compilation approach resembles Halide. Single-GPU attention is optimal within constants, but parallelization across GPUs introduces another IO level.Source: §5 · p. 10

A Related Work

  • Butterfly matrices [15, 64] motivate the block-sparse pattern through their expressive efficiency [16], but theoretical savings can fail to produce wall-clock speedups.Highly optimized dense multiplication creates this hardware-lottery barrier [41]; butterfly extensions [17, 18] aim to improve hardware compatibility.Source: A Related Work · p. 17
  • Block-sparse FlashAttention fixes a butterfly sparsity pattern throughout training and performs almost as well as dense FlashAttention on Long-range Arena tasks.The authors interpret this fixed pattern as an attention-specific fixed lottery ticket, connecting it to the lottery tickets hypothesis [28] [29] [30].Source: A Related Work · p. 17

B Algorithm Details

  • FlashAttention reduces HBM accesses in both attention passes, improving runtime and memory footprint beyond naive linear-extra-memory formulations that retain quadratic HBM traffic.Reducing intermediate storage alone therefore does not guarantee faster execution.Source: B Algorithm Details · pp. 17–18

B.1 Memory-efficient forward pass

  • O(n) extra memory suffices for the forward pass by computing softmax normalization constants separately, then accumulating outputs by repeated summation.This decouples K/V columns using a technique [60] previously applied to attention [51] [66], although HBM accesses remain quadratic. Normalization constants require O(n) extra memory, while output accumulation requires O(d).Source: §B.1 · p. 18

B.2 Memory-efficient backward pass

  • Linear extra memory suffices for an explicitly derived backward pass, rather than applying gradient checkpointing as suggested by Rabe and Staats [66].Stored softmax normalization constants allow dV to be accumulated by repeated summation without extra memory.Source: §B.2 · p. 18

B.3 FlashAttention: Forward Pass

  • The full forward pass supports masking and elementwise dropout, saving outputs, softmax statistics, and the random-generator state for the backward pass.Masking sets excluded entries to negative infinity; dropout zeros elements with probability p and rescales retained elements by 1/(1-p).Source: §B.3 · p. 19
  • FlashAttention partitions inputs, outputs, and softmax statistics into blocks, loading K/V blocks and then Q/output/statistic blocks into SRAM for on-chip score computation.The outer loop loads each K/V block; the inner loop processes the corresponding Q, output, and statistic blocks.Source: §B.3 · p. 20

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

FlashAttention’s backward pass performs O(N^2) FLOPs using O(N) extra memory. Compared with Rabe and Staats [66], it targets memory traffic and avoids retaining multiple temporary block outputs or recomputing them during backpropagation.

B.4 FlashAttention: Backward Pass

  • O(N) extra memory suffices when saved pseudo-random generator states regenerate the dropout mask during backpropagation, avoiding storage of the O(N^2) forward mask.The generator states are saved during the forward pass.Source: §B.4 · p. 20
  • The softmax-gradient reduction uses D_i = rowsum(dO_i ◦ O_i), avoiding reductions over length-N probability and gradient rows that might not fit SRAM.Algorithm 4 implements the backward derivation blockwise.Source: §B.4 · pp. 20–21
  • O(N^2) FLOPs and O(N) extra memory suffice for FlashAttention’s backward pass, beyond storage for inputs, output, output gradient, and input gradients.These computational and extra-memory complexities match those of the forward pass.Source: §B.4 · p. 21
  • Θ(N^2 d^2 M^-1) HBM accesses are required by FlashAttention’s backward pass, versus Θ(Nd + N^2) for standard attention, when d ≤ M ≤ Nd.N is sequence length, d is head dimension, and M is SRAM size.Source: §B.4 · p. 21

B.5 Comparison with Rabe and Staats [66]

  • Both FlashAttention and Rabe and Staats [66] tile attention blocks, avoid storing the attention matrix during the forward pass, and recompute it during backpropagation.This shared strategy reduces the memory footprint.Source: §B.5 · p. 22
  • 2-4× faster than standard attention, FlashAttention contrasts with Rabe and Staats [66], which is approximately as fast or slightly slower.FlashAttention targets memory accesses, whereas Rabe and Staats [66] targets maximum GPU memory footprint; both substantially reduce memory requirements.Source: §B.5 · p. 22
  • FlashAttention analytically simplifies backpropagation to recompute only the attention matrix, whereas Rabe and Staats [66] also recomputes temporary block outputs using gradient checkpointing.This reduces backward memory requirements and yields speedup. FlashAttention also updates outputs incrementally, retaining one output copy instead of K temporary-output copies for K blocks in Rabe and Staats [66].Source: §B.5 · p. 22

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

The proofs establish FlashAttention’s exactness and analyze its forward and backward HBM accesses. The block-sparse extension skips zero blocks, reducing block-dependent IO while retaining the cost of writing the output.

C Proofs

  • O(N) extra memory stores the normalization statistics ℓ and m.These statistics track normalization sums and row maxima during the blocked computation.Source: C Proofs · p. 22
  • Induction establishes that accumulated row maxima, normalization sums, and partial attention outputs produce the exact final result O = softmax(QK^T)V.After each outer-loop iteration, these quantities correspond to attention over the key and value blocks processed so far.Source: C Proofs · pp. 22–23
  • Θ(NdT_c) HBM accesses suffice for FlashAttention’s forward pass, compared with Θ(Nd + N^2) for standard attention.K and V are each loaded once; Q and O are traversed T_c times.Source: C Proofs · p. 23
  • Tile sizes must allow the B_c×d K/V blocks, B_r×d Q/O blocks, and B_r×B_c score block to fit in on-chip memory.These storage constraints determine the admissible block sizes used in the IO analysis.Source: C Proofs · p. 23
  • Ω(Nd) HBM accesses are unavoidable for exact attention because the input matrices and output have size Nd and reside in HBM.This input/output requirement supplies the contradiction in the lower-bound proof.Source: C Proofs · p. 24
  • Θ(NdT_c) HBM accesses suffice for FlashAttention’s backward pass, compared with Θ(Nd + N^2) for standard attention backward.K and V are loaded once, and dK and dV are written once; Q, O, dO, and dQ require T_c passes.Source: C Proofs · p. 24

D.1 Block-sparse FlashAttention

  • Block-sparse FlashAttention skips zero blocks, scaling block-dependent HBM accesses by the nonzero fraction s while retaining the cost of writing the output.Algorithm 5 otherwise matches Algorithm 2; even when s is small, the N×d output O must still be written.Source: §D.1 · p. 25

D.2 Potential Extensions

Potential extensions of the IO-aware approach include multi-GPU attention, sparse MLP layers, and kernel methods, targeting memory-hierarchy asymmetries and memory-access bottlenecks.

Multi-GPU Attention

  • Multi-GPU attention could exploit memory-hierarchy asymmetries across SRAM, local HBM, and other GPUs’ HBM to cooperatively process very long sequences.Large language models train on hundreds or thousands of GPUs, typically splitting attention across 4-8 GPUs within a node [77].Source: §D.2 · pp. 25–26

Sparse MLP layers

  • IO-aware sparse MLP implementations could alleviate memory bottlenecks that prevent speedups from scaling with sparsity.Typical dense MLP layers are compute-bound; sparse weight matrices can improve efficiency [17], but many sparse MLP layers become memory-bound.Source: §D.2 · p. 26

Kernel machine learning

  • Kernel methods could prioritize reducing IO rather than only FLOPs, exploiting the low-dimensional inputs underlying pairwise kernel-matrix entries.FlashAttention reduces HBM access by reloading Q and K and recomputing attention blocks from QK^T, whose rank d is much smaller than N. The KeOps library [8] demonstrates that reducing memory reads and writes can accelerate kernel operations.Source: §D.2 · p. 26

E Full Experimental Results · E.1 BERT

BERT-large training follows the MLPerf 1.1 reference procedure, measuring wall-clock time to a 72.0% masked-language-modeling validation accuracy target. The evaluation compares against Nvidia’s reported MLPerf 1.1 training speed using the same data split and validation examples.

E.1 BERT

  • 72.0% masked-language-modeling validation accuracy is the stopping target for wall-clock training-time comparisons against Nvidia’s reported MLPerf 1.1 training speed.The supplied passages describe the comparison protocol but provide no comparative speed result.Source: §E.1 · p. 26
  • 3.75e-3 is the LAMB optimizer learning rate, following the MLPerf 1.1 reference training procedure and hyperparameters for BERT-large.Source: §E.1 · p. 26
  • 448 is the training batch size, with training limited to at most 7100 steps.Source: §E.1 · p. 26
  • FP16 training uses Apex AMP with the O2 optimization level.Source: §E.1 · p. 26
  • 10000 validation examples and the MLPerf 1.1 reference train/validation split match those used by the Nvidia baseline.Source: §E.1 · p. 26
  • 10 runs are averaged when measuring wall-clock training time to the validation accuracy target.Training uses 8×A100-80GB GPUs, and each run takes between 16 and 19 minutes.Source: §E.1 · p. 26

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

FlashAttention matches HuggingFace GPT-2 validation perplexity curves, while almost all tuned attention methods achieve similar accuracy across five LRA tasks. These sections document the training and comparison protocols, including pretraining and fine-tuning for Path-X and Path-256.

E.2 GPT-2

  • GPT-2 experiments use standard Huggingface transformers and Nvidia Megatron-LM implementations, following the Megatron-LM training recipe.For long-document classification on MIMIC-III and ECtHR, the authors follow the hyperparameters of Dai et al. [13].Source: §E.2 · p. 26
  • 512 is the effective batch size for GPT-2 training, using gradient accumulation, AdamW, and PyTorch AMP mixed precision for 400K steps.AdamW learning rates are 6e-4 for GPT-2 small and 1.5e-4 for GPT-2 medium, with weight decay 0.1. All implementations use the same hyperparameters.Source: §E.2 · p. 26
  • 0.5% of Openwebtext is randomly selected once for shared validation across all GPT-2 models, with the remainder used for training.The experiments use the GPT-2 BPE tokenizer.Source: §E.2 · p. 26
  • 2.7-9.5 days for GPT-2 small and 6.9-21.0 days for GPT-2 medium are reported for training on eight A100-40GB GPUs.The reported measurement is wall-clock training time.Source: §E.2 · p. 26
  • FlashAttention produces nearly identical validation perplexity curves to the HuggingFace baseline for GPT-2 small and medium.The two implementations’ validation perplexity curves almost overlap throughout training.Source: §E.2 · p. 26, Figure 4 · p. 27

E.3 LRA details

  • Almost all attention methods achieve similar accuracy across all five LRA tasks after hyperparameter tuning.Hyperparameters follow Long-range arena, its repository, and the Nyströmformer reproduction; unreproduced baselines receive the better reported performance from Tay et al. or Xiong et al.. Mixed precision excludes Performer because of instability and Local Attention because FP16 is unsupported; overall wall-clock speedup is the geometric mean across the five tasks.Source: §E.3 · p. 27
  • Roughly 4 points of Path-X accuracy are added to FlashAttention by an additional 200-epoch fine-tuning stage, with overfitting afterwards.Following PathFinder-32 hyperparameters from the long-range arena paper [80], Path-X and Path-256 pretrain on Path-64 for 200 epochs, duplicate positional embeddings gridwise, and fine-tune for 200 epochs with one warmup epoch and cosine learning-rate decay. The additional Path-X stage starts from the best validation-accuracy checkpoint and uses the same warmup and learning rate.Source: §E.3 · p. 27

E.4 Comparison with Apex FMHA

FlashAttention builds on Apex FMHA with tiling and recomputation to save memory and support sequences up to length 64K. For short sequences, FlashAttention has a slightly faster forward pass but a slightly slower backward pass than FMHA.

E.4 Comparison with Apex FMHA

  • 4% slower at sequence length 128, 8% faster at 256, and 5% faster at 512: FlashAttention’s overall runtime compared with Apex FMHA.Table 7 measures attention with masking and dropout on an A100-SXM4-40GB GPU, using batch size 64, 16 heads, and head dimension 64.Source: Table 7 · p. 28, §E.4 · p. 28
  • Apex FMHA targets BERT, supporting sequence lengths at most 512, head dimension 64, and A100 GPUs only.It was the fastest attention implementation known to the authors when the project began; almost all Nvidia-GPU MLPerf BERT training submissions used FMHA as of MLPerf 1.1 [58].Source: §E.4 · pp. 27–28
  • 64K-length sequences are supported by applying tiling and recomputation to FMHA code, alongside broader head-dimension and GPU support.Supported head dimensions are 16, 32, 64, and 128; supported GPUs include all Turing and Ampere GPUs at the time of writing.Source: §E.4 · p. 28
  • FlashAttention avoids storing the attention matrix and recomputes it during backpropagation, explaining its slightly faster forward pass and slightly slower backward pass than FMHA.FMHA fuses attention into one CUDA kernel but stores the attention matrix in HBM for gradient computation, so it does not offer substantial memory savings.Source: §E.4 · p. 28

E.5 Speedup On Different Hardware and Configurations

FlashAttention’s speedup over standard PyTorch attention varies across GPUs and attention configurations. Memory bandwidth, SRAM capacity, head dimension, and masking affect the observed gains.

A100

  • 2-4× speedup over standard PyTorch attention is generally observed on A100, with larger gains from dropout and masking due to kernel fusion.Measurements span different sequence lengths with batch size 8, head dimension 64, and 12 attention heads.Source: Figure 5 · p. 28, §E.5 · p. 28

A100, Head Dimension 128

  • Up to 3× speedup over standard PyTorch attention remains possible on A100 with head dimension 128 and a causal mask, despite lower overall gains.Larger heads require smaller blocks to fit SRAM, while causal masking skips half the blocks. Measurements use batch size 16 and 12 heads.Source: Figure 6 · p. 29, §E.5 · p. 29

RTX 3090

  • 2.5-4.5× speedups over standard PyTorch attention on RTX 3090 slightly exceed A100 speedups because RTX 3090 has lower memory bandwidth.RTX 3090 bandwidth is roughly 900 GB/s versus A100’s 1.5 TB/s; measurements use batch size 12 and 12 attention heads.Source: Figure 7 · p. 29, §E.5 · p. 29

T4

  • T4 achieves less speedup because its smaller SRAM requires smaller FlashAttention blocks, consistent with the IO complexity analysis.Results cover combined forward and backward passes and forward-only execution, reflecting T4’s common use for inference.Source: §E.5 · p. 29, Figure 8 · p. 30

E.6 Full Benchmarking Results

This section benchmarks attention runtime and memory across sequence lengths on a single A100, comparing exact, approximate, and sparse implementations under multiple dropout and masking configurations.

E.6 Full Benchmarking Results

  • Benchmarks compare PyTorch/HuggingFace and Megatron exact attention against five approximate and three sparse attention implementations.Approximate baselines are Reformer [51], Local Attention, Linformer Attention [84], Smyrf, and LongShortFormer; sparse baselines are OpenAI Block-Sparse Attention [11], Longformer [3], and BigBird Attention. Approximate and sparse methods use a compression ratio of 1/8 or a compressed sequence length of 256, whichever is smaller.Source: §E.6 · p. 30
  • 8 heads of dimension 64 and batch size 16 define the attention workload on one A100 GPU with 40 GB HBM.Random Q, K, and V inputs exclude hidden-layer projections; runtime averages 100 measurements, while memory is measured once. Dropout is 0.1, and padding-mask lengths are uniformly random between the sequence length minus 20 and the full sequence length.Source: §E.6 · p. 30
  • Runtime benchmarks cover forward, backward, and combined passes under dropout and masking configurations; memory benchmarks measure combined passes without dropout or masking.Runtime tables report milliseconds, while the memory table reports MB.Source: §E.6 · p. 31, Table 11 · p. 32, Table 21 · p. 34
  • Block-Sparse, Longformer, and BigBird are measured without masking because external-library bugs prevented their masked backward passes.This exception was intended to give these baselines a generous comparison.Source: §E.6 · p. 31
  • Measurements use FP16 except for Local Attention, whose implementation supports only FP32.Source: §E.6 · p. 31
  • Megatron, OpenAI Block-Sparse, and Longformer/BigBird support sequence lengths only up to 2048, 4096, and 8092, respectively.Other baselines are tested at increasing sequence lengths until GPU memory is exhausted.Source: §E.6 · p. 31

Related papers

Loading 2205.14135v2…