Flash Attention Explained

9 min read

Learn

Quick Reading

Estimated 9 mins

Prereq

Intermediate

Basic ML concepts helpful

Interactive

Static Playbook

Static guide & reference tables

The real bottleneck isn't FLOPs, it's memory movement

Standard self-attention computes an $N \times N$ attention matrix (scores between every pair of tokens), applies softmax, then multiplies by $V$. For a sequence of length $N$, this requires $O(N^2)$ memory to materialize the full score matrix. On modern GPUs, the actual performance bottleneck is rarely raw compute (FLOPs) — it's moving data between the GPU's slow, large high-bandwidth memory (HBM) and its fast, tiny on-chip SRAM. Naive attention implementations write the full $N \times N$ matrix to HBM, read it back for softmax, write it again, then read it again for the final matmul — four expensive round trips through slow memory for a computation that's fundamentally about small local operations.

FlashAttention's trick: fused, tiled computation

FlashAttention (Dao et al., 2022) restructures attention to never materialize the full $N \times N$ matrix in HBM at all. It tiles the computation: it loads small blocks of $Q$, $K$, and $V$ into fast SRAM, computes partial attention scores and partial outputs for that block, and incrementally combines them across blocks using an online softmax algorithm — a running, numerically stable way to compute softmax over a stream of values without seeing them all at once. The result is mathematically identical attention output, computed with far fewer HBM reads/writes, because everything stays in fast on-chip memory until the final result is ready.

Online softmax: the algorithmic key

Softmax normally needs the full row of scores before you can divide by their sum. Online softmax avoids this by tracking a running maximum and a running sum as each new block of scores arrives, rescaling previously accumulated partial outputs whenever the running max changes. This lets attention be computed incrementally, block by block, which is exactly what makes tiling in SRAM possible — you never need the whole row in memory at once.

online_softmax_sketch.py

python

import numpy as np

def online_softmax_blocks(score_blocks):
    """Compute softmax normalization incrementally over blocks of scores,
    the core trick behind FlashAttention's tiling."""
    running_max = -np.inf
    running_sum = 0.0
    for block in score_blocks:
        block_max = block.max()
        new_max = max(running_max, block_max)
        # rescale the old running sum to the new max before adding
        running_sum = running_sum * np.exp(running_max - new_max) \
                      + np.exp(block - new_max).sum()
        running_max = new_max
    return running_max, running_sum

blocks = [np.array([1.0, 2.0]), np.array([3.5, 0.2]), np.array([2.9, 4.1])]
m, s = online_softmax_blocks(blocks)
print(f"running max={m:.3f}, running sum={s:.3f}")
# Matches exactly what you'd get from softmax over all 6 values at once.

Standard attention vs FlashAttention

Property	Standard attention	FlashAttention
HBM memory for scores	O(N^2) materialized	Never materialized, O(N) working memory
HBM read/write passes	Multiple (scores, softmax, output)	Single fused pass per block
Output values	Exact softmax attention	Exact, mathematically identical
Speedup source	—	Memory bandwidth, not fewer FLOPs
Typical result	Memory-bound, slower at long sequence length	2-4x faster, enables much longer context windows

Why this matters beyond raw speed

FlashAttention's memory savings are what made long-context training and inference practical at all. Because the $N^2$ score matrix is never stored, peak memory usage scales roughly linearly rather than quadratically with sequence length, which means you can train or serve with 32K, 128K, or longer context windows on the same GPU that previously choked at 2K–4K tokens. FlashAttention-2 and FlashAttention-3 extended the same idea with better GPU occupancy and support for newer hardware (and work asymmetries like causal masking), but the core insight — minimize HBM traffic, fuse the whole computation into tiled SRAM operations — has stayed the same.

It's an implementation detail, not a new attention mechanism

It's easy to confuse FlashAttention with an architectural change like sparse or linear attention. It isn't — it computes the exact same softmax attention output as the naive implementation, bit-for-bit equivalent up to floating point rounding. The difference is purely how the computation is scheduled on hardware. This is why it's a near-universal drop-in optimization rather than a tradeoff you have to reason about for model quality.

What's Next?

FlashAttention optimizes how attention is computed during a single forward pass, but inference has its own separate memory problem — remembering past keys and values across every generation step — which is solved with a different mechanism entirely.

I build these systems professionally.

Whether it's a RAG pipeline, analytics migration, or AI workflow — let's talk.

Work With Me →