FlashAttention

FlashAttention is an IO-aware exact attention algorithm that computes the standard scaled dot-product attention mechanism while dramatically reducing the number of reads and writes to GPU High Bandwidth Memory (HBM), achieving 2–4× wall-clock speedup and sub-quadratic memory footprint relative to sequence length compared to naive implementations. It matters because memory bandwidth—not arithmetic throughput—is the dominant bottleneck in transformer training and inference at scale, making FlashAttention a critical primitive underlying virtually every state-of-the-art large language model. The counterintuitive fact: FlashAttention performs more floating-point operations than standard attention yet runs substantially faster, because FLOP count is not the binding constraint on modern GPU hardware.

The core insight

Standard attention materializes the full N×N attention matrix in HBM, causing memory traffic that scales quadratically with sequence length. FlashAttention never stores the full attention matrix—instead it tiles the computation into SRAM-resident blocks, fuses the softmax normalization with the matrix multiplications using an online softmax trick, and accumulates the output in a single fused kernel. The "aha moment" is realizing that the softmax denominator can be computed incrementally across tiles without requiring a second pass over the data, using the log-sum-exp numerically stable identity—so correctness is preserved while the memory bottleneck is eliminated entirely.

Common trapCommon Misconception

The most pervasive mistake is conflating FLOPs with runtime, leading practitioners to believe FlashAttention must be an approximate method because "doing more work can't be faster." This is wrong for a precise reason: modern A100 GPUs deliver ~312 TFLOPS of FP16 compute but only ~2 TB/s of HBM bandwidth, giving an arithmetic intensity threshold (roofline model) of ~156 FLOPs/byte. Standard attention operates well below this threshold—it is memory-bound, not compute-bound. FlashAttention raises arithmetic intensity by reusing data in SRAM (which has ~19× higher bandwidth than HBM on A100), shifting the operation into the compute-bound regime. The correct mental model is not "fewer operations = faster" but "fewer HBM round-trips = faster," which is an entirely different optimization axis than algorithmic complexity.

How It Works

Scaled Dot-Product Attention (Baseline)

Given query, key, and value matrices $Q, K, V \in \mathbb{R}^{N \times d}$ , standard attention computes:

\quad \text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d}}\right)V

Explanation: Each output row is a convex combination of value vectors, weighted by normalized similarity scores between the corresponding query and all keys. Here $N$ is the sequence length, $d$ is the head dimension, $QK^\top \in \mathbb{R}^{N \times N}$ is the attention logit matrix (requiring $O(N^2)$ memory to store), and the $\sqrt{d}$ scaling prevents vanishing gradients in the softmax.

The Memory Bottleneck

The number of HBM accesses for naive attention is:

\quad \Theta_{\text{HBM}}^{\text{naive}} = O(Nd + N^2)

Explanation: The $N^2$ term dominates at long sequences because the full attention matrix must be written to and read from HBM for the softmax and the subsequent multiplication by $V$ . For $N = 8192$ and $d = 128$ in FP16, the attention matrix alone consumes 1 GB of HBM per head per layer—an enormous bandwidth cost repeated at every forward and backward pass.

Online Softmax Decomposition

The key mathematical primitive enabling FlashAttention is the online softmax identity. For a row of logits $x = (x_1, \ldots, x_N)$ , define:

\quad m^{(j)} = \max_{i \leq j} x_i, \qquad \ell^{(j)} = \sum_{i=1}^{j} e^{x_i - m^{(j)}}

Explanation: The running maximum $m^{(j)}$ and running sum $\ell^{(j)}$ encode all information needed to compute the exact softmax up to block $j$ , and can be updated incrementally when new logits arrive. When a new block of logits $x_{j+1}, \ldots, x_{j+B}$ arrives, the statistics update as:

\quad m^{(j+B)} = \max(m^{(j)},\, m_{\text{new}}), \quad \ell^{(j+B)} = e^{m^{(j)} - m^{(j+B)}} \ell^{(j)} + e^{m_{\text{new}} - m^{(j+B)}} \ell_{\text{new}}

Explanation: The correction factors $e^{m_{\text{old}} - m_{\text{new}}}$ rescale previously accumulated statistics to the new global maximum, guaranteeing numerical stability without a separate pass.

Tiled Output Accumulation

The output accumulator $O$ for a query block is updated at each key-value tile $j$ as:

\quad O^{(j+1)} = \text{diag}\!\left(e^{m^{(j)} - m^{(j+1)}}\right) O^{(j)} + e^{S_{\text{new}} - m^{(j+1)}} V_{\text{new}}

Explanation: The diagonal rescaling matrix corrects the previously accumulated output to be consistent with the updated softmax normalization, allowing the final result to be exact despite never forming the full $N \times N$ matrix. This update has $O(1)$ SRAM overhead per tile regardless of $N$ , reducing total HBM accesses to $O(N^2 d / M)$ where $M$ is the SRAM size—a factor of $M/d$ improvement.

HBM Access Complexity

FlashAttention's total HBM read/write count satisfies:

\quad \Theta_{\text{HBM}}^{\text{Flash}} = O\!\left(\frac{N^2 d^2}{M}\right)

Explanation: Because $M \gg d$ on all modern GPUs (A100 SRAM is 192 KB per SM vs. $d \sim 64$ – $256$ ), this represents a large constant-factor reduction in memory traffic compared to the $O(N^2)$ baseline. Dao et al. (2022) prove this is optimal up to logarithmic factors for algorithms that compute exact attention using an SRAM of size $M$ —FlashAttention is theoretically IO-optimal.

researcher · Tri Dao | b. 1994 | IO-Aware Exact Attention Algorithm

Tri Dao completed his PhD at Stanford under Christopher Ré and Atri Rudra, focusing on the intersection of systems and machine learning theory. His 2022 paper "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness" (with Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré) identified HBM bandwidth as the primary bottleneck and provided both the algorithm and its theoretical IO-complexity lower bound. A surprising detail: Dao developed much of the CUDA kernel implementation working at Together AI, and the paper was initially met with skepticism at NeurIPS 2022 before becoming one of the most practically impactful systems papers of the decade.

Figure

FlashAttention Tiling and Memory Hierarchy

Figure

Online Softmax Update Rule

Intuition

The Mental Model

Imagine you are computing the weighted average of 10,000 documents, but the weights depend on a softmax over all pairwise similarity scores—so you cannot compute any weight until you've seen all scores. A naive librarian would read every document into RAM, compute all 100 million similarity scores, write them back to disk, read them again to normalize, then compute the weighted sum—four disk trips. A smarter librarian processes documents in batches: for each batch, they keep a running notecard with the current maximum score and running sum (just two numbers), update the notecard as each new batch arrives using the rescaling identity, and accumulate the weighted average on a scratchpad—all without ever writing intermediate results to disk. FlashAttention is that smarter librarian: SRAM is the notecard and scratchpad, HBM is the disk, and the online softmax trick is the insight that you only need two running statistics to compute an exact softmax incrementally. The reason this works on a GPU specifically is that SRAM (shared memory on the streaming multiprocessor) is ~19× faster than HBM, so keeping the hot computation in SRAM is worth any amount of extra arithmetic.

Historical Development

Origins · 1 minhow this was discovered

History

2014

Bahdanau et al. introduce neural attention for sequence-to-sequence models, establishing the paradigm of weighted aggregation over hidden states — but without self-attention or softmax bottleneck concerns at practical scales.

2017

Vaswani et al. ("Attention Is All You Need") introduce the Transformer with scaled dot-product self-attention; the

O(N^2)

memory cost is noted but not treated as the primary bottleneck at sequence lengths ≤512.

2018–2020

Sparse attention variants (Child et al., Longformer, BigBird) attempt to reduce the

O(N^2)

cost by computing attention over subsets of token pairs, trading exactness for efficiency — the dominant approach before IO-awareness is recognized.

2021

Rabe & Staats ("Self-attention Does Not Need

O(N^2)

Memory") show that attention can be computed in

O(N)

memory using a sequential scan, foreshadowing the online softmax approach but without GPU systems analysis or practical CUDA implementation.

2022

Dao, Fu, Ermon, Rudra & Ré publish "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness" (NeurIPS 2022); the paper provides the first GPU-optimized tiled attention kernel with IO complexity analysis, achieving 2–4× speedup and enabling 4–8× longer context at the same memory budget.

2022

FlashAttention is adopted in the GPT-NeoX codebase and later incorporated into HuggingFace Transformers and Megatron-LM, becoming the de facto standard attention implementation for LLM training.

2023

Dao publishes "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning" (ICLR 2024 track); achieves ~2× further speedup over v1 by improving thread-block partitioning, reducing non-matmul FLOPs, and exploiting tensor core utilization more aggressively.

2023

FlashAttention-2 is integrated into xFormers (Meta), PyTorch 2.0's `torch.nn.functional.scaled_dot_product_attention` (SDPA), and becomes the default backend for training GPT-4, Llama 2, Mistral, and virtually every frontier model.

2024

Dao & colleagues publish "FlashAttention-3" targeting H100 GPUs specifically, exploiting asynchronous memory copy (TMA), WGMMA instructions, and FP8 support to reach ~750 TFLOPS on H100—close to theoretical peak—while also extending to non-causal and multi-query attention variants.

Why It Matters

Real-World Impact

GPT-4 and OpenAI Training InfrastructureFlashAttention-2 is used throughout GPT-4's training pipeline, enabling the 8k and 32k context variants by making the previously prohibitive memory cost of long-context attention tractable on A100 clusters without tensor parallelism overhead.

Llama 2 and Llama 3 (Meta AI)Meta's Llama 2 (70B) and Llama 3 (405B) training codebases use FlashAttention-2 as the default attention kernel; the 128k context window in Llama 3.1 is directly enabled by the sub-quadratic memory profile of FlashAttention combined with grouped-query attention.

Mistral 7B and Sliding Window AttentionMistral AI's sliding window attention mechanism is implemented as a masked variant of FlashAttention, achieving efficient

O(N \cdot w)

attention over windows of size

w

while retaining the IO-optimal tiling structure.

PyTorch 2.0 `torch.compile` and SDPAPyTorch's `scaled_dot_product_attention` operator (introduced in 2.0) automatically dispatches to FlashAttention-2 when inputs are on CUDA, making IO-optimal attention the default for any model written in PyTorch without requiring manual kernel integration.

State Space Models and Hybrid ArchitecturesMamba-2 and other SSM-attention hybrid architectures (e.g., Jamba, Zamba) use FlashAttention selectively on attention layers, with the same tiling philosophy informing the design of the selective state space scan kernels.

AlphaFold 3 and Protein Structure PredictionDeepMind's AlphaFold 3 uses attention over amino acid tokens with sequence lengths up to thousands; FlashAttention's memory efficiency makes the full pairwise attention in the Evoformer and Diffusion modules computationally feasible at scale.

Diffusion Transformers (DiT, Stable Diffusion 3, FLUX)Text-to-image models using the DiT architecture (e.g., Stable Diffusion 3, FLUX.1) apply FlashAttention across spatial patch tokens; at 1024×1024 resolution with 16×16 patches,

N=4096

tokens per image make IO-optimal attention critical for training throughput.

FlashAttention

How It Works

Scaled Dot-Product Attention (Baseline)

The Memory Bottleneck

Online Softmax Decomposition

Tiled Output Accumulation

HBM Access Complexity

Intuition

Historical Development

Why It Matters

Prerequisites

What This Leads To

See Also

Before this

Where this leads