FlashAttention — IO-Aware Exact Attention

What it is (intuition first)

Self-attention is the heart of a Transformer: every token looks at every other token, producing an N×N table of "how much should token i pay attention to token j". For a sequence of length N that table has N² entries, which is why long context is expensive.

The naive way to compute attention is to build that entire N×N table in memory, run a softmax over it, then multiply by the values. Two consequences follow:

Memory grows quadratically with sequence length — a 16K-token sequence wants a 256-million-entry score matrix per head, per layer.
You constantly shuttle that giant matrix back and forth between the GPU's big-but-slow memory (HBM) and its tiny-but-fast on-chip memory (SRAM).

The surprising insight behind FlashAttention is that, for attention, the cost is not the math — it's the memory traffic. The GPU's arithmetic units sit idle waiting for the score matrix to be written out and read back in.

FlashAttention is an IO-aware, exact attention algorithm. It computes the bit-for-bit same mathematical result as standard scaled-dot-product attention, but reorganizes the computation so that the full N×N matrix is never written to slow memory. It does this by (a) processing attention in small tiles that fit in fast on-chip SRAM, and (b) using an online softmax trick to combine tiles correctly without ever seeing the whole row at once (Dao et al. 2022).

The word exact matters. Unlike approximate-attention methods (Linformer, Performer, sparse attention) that change the answer to save compute, FlashAttention changes only the order and location of the computation. Same output, far less memory traffic. That is why it became the default attention kernel in PyTorch, vLLM, Hugging Face, and essentially every serious training/inference stack. (See also [[attention-and-transformers]] for the base operation FlashAttention accelerates.)

Why it matters

The bandwidth gap is the whole story. A modern data-center GPU is wildly compute-rich and bandwidth-poor relative to its compute. On an NVIDIA A100:

Compute: ~312 TFLOP/s (FP16 tensor-core math)
On-chip SRAM: ~20 MB total (192 KB per SM × 108 SMs), bandwidth ~19 TB/s
HBM (the big GPU RAM): 40–80 GB, bandwidth only ~1.5–2.0 TB/s (memory-hierarchy figures cross-checked in FA1 §2).

SRAM is roughly 10× faster than HBM but ~4000× smaller. An operation that mostly moves data (rather than doing dense matmuls) is memory-bound: the arithmetic finishes long before the data arrives. Attention's softmax, masking, and dropout are exactly such elementwise, low-arithmetic-intensity steps. Standard attention spends most of its time reading and writing the N×N score matrix to HBM, not computing.

FlashAttention removes that traffic. Concrete wins reported in the papers:

3× faster GPT-2 end-to-end training at 1K context; 15% wall-clock speedup on BERT-large vs. the MLPerf 1.1 record (FA1).
Linear (not quadratic) memory in sequence length — you store O(N) softmax statistics instead of the O(N²) score matrix. This is what unlocked training on much longer contexts; FA1 was the first to push a Transformer through Path-X (16K) and Path-256 (64K).
2× faster again with FlashAttention-2, reaching 50–73% of the A100's theoretical FLOP ceiling (FA1 was only 25–40%) (FA2).
1.5–2× faster again with FlashAttention-3 on H100, ~740 TFLOP/s (~75% util) in FP16 and approaching 1.2 PFLOP/s in FP8 (FA3).

In short: FlashAttention made long-context Transformers both faster and memory-feasible, without sacrificing exactness. Every modern context-window expansion stands on it.

How it works (the real mechanics)

Standard attention, and why it's wasteful

Given queries Q, keys K, values V (each N×d), standard attention computes:

S = Q Kᵀ            # N×N scores            ← write N² to HBM
P = softmax(S)      # row-wise softmax       ← read N², write N² to HBM
O = P V             # N×d output             ← read N² from HBM

The N×N matrices S and P are materialized in HBM. For long N, that traffic dominates. IO complexity of standard attention is Θ(N²) HBM accesses (plus the Nd input/output).

Tiling: never build the whole matrix

FlashAttention splits Q into row-blocks and K, V into column-blocks sized to fit in SRAM. The classic FA1 loop structure is:

for each block of K, V (outer loop):          # columns of the score matrix
    load Kⱼ, Vⱼ  HBM → SRAM
    for each block of Q (inner loop):          # rows
        load Qᵢ, and running output Oᵢ, stats (mᵢ, ℓᵢ)  HBM → SRAM
        Sᵢⱼ = Qᵢ Kⱼᵀ                            # small tile, stays in SRAM
        update Oᵢ and (mᵢ, ℓᵢ) via online softmax (below)
        write back Oᵢ, mᵢ, ℓᵢ  SRAM → HBM

The N×N matrix is computed one small tile at a time inside SRAM and never written to HBM in full. The whole forward pass is fused into a single GPU kernel — no intermediate round-trip. (FA2 later flips the loop order so the outer loop is over Q-blocks; see below.)

Online softmax: the trick that makes tiling exact

The obstacle: softmax needs the maximum and the sum of exponentials over the entire row for numerical stability and normalization — but with tiling you only see one column block at a time. The online softmax (a.k.a. the "running-max / running-sum" trick) solves this. For each query row i you keep two scalars as you stream through key blocks:

mᵢ — the running maximum score seen so far
ℓᵢ — the running sum of exp(score − mᵢ)

When a new tile arrives with its own local max m_new, you update the global max m = max(mᵢ, m_new) and rescale the partial sum and the partial output by exp(mᵢ − m) so everything is referenced to the new maximum:

ℓ  ← exp(mᵢ − m)·ℓᵢ + Σ exp(Sᵢⱼ − m)
Oᵢ ← exp(mᵢ − m)·Oᵢ  +  (Σ exp(Sᵢⱼ − m)·Vⱼ)
mᵢ ← m

After the last block, divide Oᵢ by ℓᵢ to get the normalized output. Because every term is consistently rebased to the same running maximum, the final result is identical to a softmax computed over the full row — only the order of summation changed. This is the single most important idea: it is what makes FlashAttention exact rather than approximate.

Recomputation in the backward pass

Training needs gradients, which normally require the stored P (the N×N softmax probabilities). Storing it would reintroduce O(N²) memory. FlashAttention instead recomputes the score tile Sᵢⱼ and probability tile Pᵢⱼ on the fly during backprop, using the same tiling. The only thing it needs to persist from the forward pass is one scalar per row, the log-sum-exp statistic:

Lᵢ = mᵢ + log ℓᵢ        →   Pᵢⱼ = exp(Sᵢⱼ − Lᵢ)  recomputed in SRAM

This is a classic memory-for-compute trade: you re-do some matmuls (a modest FLOP increase) to avoid the enormous HBM traffic of storing and reloading P. Because attention was memory-bound, the recomputation is still a net speedup.

IO complexity — the headline theorem

FA1 proves that with SRAM of size M, FlashAttention performs

Θ(N² d² / M)   HBM accesses

versus Θ(N² ) for standard attention. Since d² is typically much smaller than M (e.g. d=64, M≈100K elements), this is a large reduction — and the paper proves it is asymptotically optimal: no exact-attention algorithm can use asymptotically fewer HBM accesses across a realistic range of SRAM sizes (FA1). This is the rare case of a kernel optimization with a matching lower bound.

The version ladder: FA1 → FA2 → FA3

FlashAttention-1 (2022) — Introduces IO-awareness, tiling, online softmax, recomputation, and the optimality proof. Also a block-sparse variant (the one approximate offering: skip whole tiles for sparse masks). Roughly 25–40% of A100 peak — good, but the kernel still had inefficiencies. (arXiv 2205.14135)

FlashAttention-2 (2023) — Same algorithm, better GPU engineering. Three changes: (1) reduce non-matmul FLOPs — tensor cores do matmuls ~16× faster than they do the rescaling/exp work, so minimizing non-matmul ops matters a lot; (2) parallelize over the sequence-length dimension, not just batch×heads, so even a single long sequence keeps all SMs busy (high occupancy); (3) better warp-level work partitioning to cut shared-memory communication inside a thread block. Net: ~2× over FA1, 50–73% of A100 peak, ~225 TFLOP/s and ~72% MFU in real GPT training. (arXiv 2307.08691)

FlashAttention-3 (2024) — Rewritten for the Hopper (H100) architecture, exploiting hardware asynchrony. Three ideas: (1) warp-specialization — producer warps fetch data with the Tensor Memory Accelerator (TMA) while consumer warps run WGMMA matmuls, overlapping data movement with compute; (2) interleave block-wise matmul and softmax so the slow exp/softmax of one block hides under the matmul of the next (ping-pong scheduling); (3) FP8 low precision with block quantization and incoherent processing (a random rotation that spreads outliers) to use the FP8 tensor cores while keeping accuracy. Result: ~740 TFLOP/s (75% util) in FP16 and near 1.2 PFLOP/s in FP8 — and notably FP8 FA3 has 2.6× lower numerical error than a naive FP8 attention baseline. (arXiv 2407.08608)

The throughline: FA1 is a new algorithm; FA2 and FA3 are increasingly hardware-specific re-implementations of the same algorithm, each squeezing closer to the GPU's GEMM ceiling.

Key ideas & tradeoffs

Memory-bandwidth is the real bottleneck, not FLOPs, for memory-bound ops like attention. The entire technique follows from taking the memory hierarchy seriously (IO-awareness).
Exact, not approximate. Reorders/relocates computation; output is unchanged. This is its key advantage over [[state-space-models]] and approximate-attention schemes, which alter the computation to escape O(N²) — FlashAttention keeps O(N²) compute but makes it cheap and memory-linear in practice.
Tiling + online softmax + recomputation are the three legs. Remove any one and you either lose exactness, blow up memory, or reintroduce HBM traffic.
Compute-for-memory trade in the backward pass: deliberately do more arithmetic to do less memory IO. Counterintuitive but correct when the op is bandwidth-bound.
Kernel fusion: the whole forward pass is one kernel — no intermediate tensors in HBM.
Hardware-coupled: each version is tuned to a GPU generation (Ampere → Hopper). The algorithm is portable; the peak performance is not. (FA3 needs H100-class hardware.)

Honest caveats & open questions

Not magic for inference decoding. FlashAttention shines when you have many query rows to tile (training, prefill). Single-token autoregressive decode has one query row, so the classic kernel underutilizes the GPU; this motivated variants like FlashDecoding and FlashDecoding++ that parallelize over the K/V (sequence) dimension instead. Plain FA is not the inference-decode answer by itself.
It does not reduce asymptotic compute. Attention is still Θ(N²) FLOPs. FlashAttention removes the memory quadratic and a large constant factor, but for truly long contexts the compute quadratic still bites — which is why sparse/linear attention and [[state-space-models]] remain active research. FlashAttention and sub-quadratic methods are complementary, not rivals.
Implementation is hard and hardware-specific. The speed lives in hand-tuned CUDA/CUTLASS for specific GPUs. Porting to non-NVIDIA accelerators, new shapes, or unusual masks/biases (ALiBi, custom attention patterns) historically lagged. FlexAttention and Triton ports exist but the bleeding-edge perf is tied to NVIDIA.
FP8 (FA3) trades accuracy for speed. The block-quantization + incoherent-processing tricks shrink the error vs. naive FP8, but FP8 attention is still a precision compromise that must be validated per-model, not assumed safe.
The optimality proof is asymptotic in HBM accesses, under modeling assumptions (range of M, exact attention). It is a strong result but not a claim that any specific implementation is wall-clock optimal on a given GPU — that's an engineering race (FA1→2→3 proves the point).
Numerical order changes. Because summation order differs from naive softmax, bitwise outputs can differ at floating-point granularity from a reference implementation even though the algorithm is mathematically exact. Usually negligible, occasionally relevant for strict reproducibility.

How it connects to OpenAlice

OpenAlice itself doesn't author attention kernels — Alice runs on hosted LLM providers (Codex OAuth gpt-5.x by default; see provider config in lab memory). But FlashAttention is load-bearing infrastructure underneath the whole stack, so a few honest ties:

It is why long context is affordable. Alice's worker-mode tasks (e.g. NAO's diploma-length drafting, large-codebase reasoning) and the long system prompts assembled by the prompt-composer all depend on the provider models having been trained and served with IO-aware attention. The "feed Alice a big context window" pattern is economically possible because of FlashAttention-class kernels on the provider side.
The IO-awareness lesson generalizes beyond GPUs. The Atlas knowledge graph and the inspector face the same "memory traffic, not raw compute, is the bottleneck" pattern at the data-pipeline level — tiling/streaming over large repos rather than materializing everything. FlashAttention is the canonical worked example of designing the algorithm around the memory hierarchy rather than the FLOP count.
If OpenAlice ever runs self-hosted inference (no GPU on blal.de today, per infra notes), the choice of attention kernel — FA2 on Ampere, FA3 on Hopper, FlashDecoding for serving — becomes a direct, concrete tuning knob. Until then the tie is conceptual but real.

(No tie is being overstated here: OpenAlice is a consumer of FlashAttention via its providers, not an implementer.)