Paper 09~14 min readNeurIPS 2022worked math + runnable code

Paper Breakdown

FlashAttention,
explained.

Everyone knew attention got painfully slow on long inputs, and everyone assumed the problem was too much math. This 2022 paper showed the arithmetic was never the bottleneck — the GPU was spending almost all its time shuffling data between fast and slow memory. Fix the shuffling, and attention gets several times faster while computing the exact same answer.

The paper

FlashAttention

Dao, Fu, Ermon, Rudra, Ré · NeurIPS 2022

Read the original →

The delta The real bottleneck Fast memory, slow memory Tiling and the online softmax The math, worked with real numbers Exact, not approximate Since then Solve it yourself What the paper actually measured Why it still matters Check yourself Where this sits

The delta

What changes against every attempt to make attention cheaper: the diagnosis

The years between the transformer and this rung produced a large literature on making attention cheaper, nearly all of it aimed at the quadratic number of operations — sparsify the pattern, approximate the softmax, low-rank the matrix. This rung's change is to point out that operations were never the bottleneck, and to fix the one that was.

The observation: a standard implementation computes the full score matrix, writes it to the accelerator's main memory, reads it back to apply the softmax, writes it again, reads it again to multiply by the values. On modern hardware, arithmetic is enormously faster than that memory traffic, so the kernel spends most of its time waiting on reads and writes of a matrix it only needs transiently. Attention was slow because it was memory-bound, and every method that traded accuracy for fewer operations was optimising the resource that was not scarce.

The fix is to never build the matrix. Process the sequence in tiles small enough to live in the chip's fast on-die memory, and accumulate the result tile by tile — which requires computing a softmax whose denominator is not known until the end, done by carrying a running maximum and running sum and rescaling the accumulated output as each new tile arrives. Stated in units: the output is bit-for-bit the same exact attention, wall-clock time falls substantially, and memory use drops from quadratic in sequence length to linear, because the quadratic object is never stored.

The real bottleneck

Self-attention has every token compare itself to every other token, producing an N×N scores matrix. For a sequence of length N, that matrix grows with the square of N — at 8,000 tokens it's 64 million numbers. The standard implementation writes that whole matrix out to memory, reads it back to apply softmax, then reads it a third time to multiply by the values.

The paper's diagnosis: the GPU's arithmetic units were mostly idle, waiting on those reads and writes. Attention was memory-bound, not compute-bound. The field had been optimizing the wrong thing.

Fast memory, slow memory

A GPU has two very different kinds of memory. HBM is the big main memory — tens of gigabytes, but relatively slow. SRAM is tiny on-chip memory — only kilobytes to a few megabytes, but roughly an order of magnitude faster.

The GPU memory hierarchy

SRAM~20 MB, ~19 TB/s — tiny but blazing fast. Work here.

HBM~40 GB, ~1.5 TB/s — huge but ~10x slower. Touch as little as possible.

Standard attention parks that giant scores matrix in slow HBM and keeps trudging back to it. FlashAttention's entire strategy is to never put it there in the first place.

Tiling and the online softmax

Instead of computing the full matrix at once, FlashAttention breaks the work into tiles — small blocks of queries and keys that fit inside fast SRAM. It loads a block, computes its piece of attention entirely in SRAM, and moves on, never writing the intermediate scores back to HBM.

The tricky part is softmax, which normally needs to see a whole row at once to normalize it. FlashAttention uses an online softmax: it processes blocks one at a time and keeps a running total and running maximum, rescaling previous partial results as new blocks arrive — so it gets the exact normalized answer without ever holding the full row. All the attention operations are fused into a single pass over fast memory.

One query block walks the key/value blocks — only the amber tile is in SRAM

The math, worked with real numbers

The rescaling trick is the whole paper, so let's do it by hand. Softmax is computed "safely" by subtracting the row's max m before exponentiating (so nothing overflows), then dividing by the sum ℓ:

softmax(s_i) = e^{s_i − m} / ℓ where m = max_i(s_i), ℓ = Σ_i e^{s_i − m}

Both m and ℓ look like they need the whole row. They don't. Take a row of scores s = [2, 1, 3, 0] and process it in two blocks — [2, 1] first, then [3, 0]:

Step	Running max m	Running sum ℓ	What happened
Block [2, 1]	2	e⁰ + e⁻¹ = 1.3679	Everything measured against max-so-far = 2
Block [3, 0] arrives	3	—	New max! The old sum used the wrong max
Rescale old ℓ	3	1.3679 × e²⁻³ = 0.5033	One multiply fixes every earlier term at once
Absorb new block	3	0.5033 + e⁰ + e⁻³ = 1.5530	Exactly what the full row gives: e⁻¹+e⁻²+e⁰+e⁻³ = 1.5530 ✓

Why does one multiply fix everything? Every old term was e^s−2 but should have been e^s−3 — and e^s−3 = e^s−2·e²⁻³. The correction factor is the same for every term, so it can be applied to their sum. That's the update rule the kernel runs per block:

m_new = max(m_old, m_block)
ℓ_new = ℓ_old · e^{m_old − m_new} + Σ_i∈block e^{s_i − m_new}
O_new = O_old · e^{m_old − m_new} + Σ_i∈block e^{s_i − m_new} · V_i

The third line is the punchline: the output accumulator O (the weighted sum of values built so far) is corrected by the same factor. Final answer = O / ℓ. Finishing the example: with m = 3 and ℓ = 1.5530, the weights come out [0.2369, 0.0871, 0.6439, 0.0321] — identical to computing softmax on the whole row at once. No approximation anywhere; the N×N matrix simply never existed.

Exact, not approximate

This is what set FlashAttention apart. Earlier attempts to speed up attention approximated it — sparse patterns, low-rank tricks — trading accuracy for speed. FlashAttention changes only the memory schedule. The result is bit-for-bit the same attention; you give up nothing.

Worth knowing

Because it never stores the N×N matrix, FlashAttention's memory grows only linearly with sequence length instead of quadratically. That's the door to long context windows — the reason models could suddenly handle tens of thousands of tokens.

Since then

What the authors admit, and what happened next

The cost is that this is a hardware-specific kernel, not an algorithm you can express in a framework and have run everywhere. Its advantage comes from the ratio between on-die and main memory bandwidth on a particular accelerator, so it has to be rewritten per architecture, and successive versions exist largely to track new hardware. The paper is explicit that the analysis is about the memory hierarchy.

There is nothing to contest in the result — exactness makes it unusually unarguable — so what is worth noting instead is the opportunity cost the field paid for the wrong diagnosis. Several years of approximate-attention research was aimed at a resource that was not binding, and that line largely stopped being competitive the moment exact attention got fast.

How it held up: it is in everything, and it is the single most-transferable lesson on this site for anyone who optimises anything. Measure which resource is saturated before redesigning the mathematics — the first-order win was available without giving up a single bit of accuracy, and it went unnoticed because the community was reasoning about complexity rather than about the memory hierarchy. Long context became practical as a side effect, since the quadratic memory term was what made it impossible.

SOLVE IT YOURSELF

Online softmax & tiled attention, for real

This is the paper's core algorithm in ~40 lines of plain Python, running in your browser through WebAssembly. naive_softmax needs the whole row; online_softmax streams it in blocks with the rescale trick; flash_attention extends the same trick to the full attention output. Run it and watch the tiled version match the standard one to 16 decimal places — then change the row, the block size, or Q/K/V and check it still does.

YOUR TASK

Write the body of flash_attention. Everything else in the file is already written; this is the function the paper is about. Press ▶ Run to check it against the 5 assertions below — they start red.

WHAT IT HAS TO SATISFY — 5 ASSERTIONS

Online softmax == naive softmax
Block size never changes the answer
Flash attention == standard attention
The online softmax is still a valid distribution (sums to 1)
Tiled flash attention equals standard attention on the worked example

CPython · WebAssembly

import math

def naive_softmax(row):
    """The textbook way - needs the WHOLE row before it can start."""
    m = max(row)                              # pass 1: global max
    exps = [math.exp(s - m) for s in row]
    total = sum(exps)                         # pass 2: global sum
    return [x / total for x in exps]

def online_softmax(row, block=2):
    """FlashAttention's way - one block at a time, never the whole row."""
    m, l = float("-inf"), 0.0                 # running max, running sum
    for b in range(0, len(row), block):
        chunk = row[b:b+block]
        m_new = max(m, max(chunk))
        # THE trick: one multiply re-bases every earlier term at once
        l = l * math.exp(m - m_new) + sum(math.exp(s - m_new) for s in chunk)
        m = m_new
    return [math.exp(s - m) / l for s in row]

row = [2.0, 1.0, 3.0, 0.0]                    # the row from the worked example
print("naive :", [round(x, 4) for x in naive_softmax(row)])
print("online:", [round(x, 4) for x in online_softmax(row)])

# ---- the same trick, extended to full attention ----
def attention(Q, K, V):
    """Standard attention - materializes every score."""
    d = len(Q[0]); out = []
    for q in Q:
        scores = [sum(qi*ki for qi, ki in zip(q, k)) / math.sqrt(d) for k in K]
        w = naive_softmax(scores)
        out.append([sum(wi*v[j] for wi, v in zip(w, V)) for j in range(len(V[0]))])
    return out

def flash_attention(Q, K, V, block=2):
    """Tiled - walks K/V in blocks, rescaling the accumulator O too."""
    # TODO: write this. Everything else on the page is already here —
    # this is the one function the paper is actually about.
    raise NotImplementedError

Q = [[1, 0], [0, 1], [1, 1], [0.5, 0.5]]
K = [[1, 0], [0, 1], [1, 1], [1, 0.5]]
V = [[1, 2], [3, 4], [5, 6], [7, 8]]
std, fl = attention(Q, K, V), flash_attention(Q, K, V)
print("standard:", [[round(x, 4) for x in r] for r in std])
print("flash   :", [[round(x, 4) for x in r] for r in fl])
diff = max(abs(x-y) for ra, rf in zip(std, fl) for x, y in zip(ra, rf))
print("max difference:", diff)                # ~1e-16. Exact, not approximate.

# Try it: change block=2 to block=1 or 3 - the answer must not move.

What the paper actually measured

The headline numbers from the NeurIPS 2022 paper, not paraphrases:

Benchmark	Against	Result
BERT-large training (seq 512)	MLPerf 1.1 record	15% faster end-to-end wall clock
GPT-2 training (seq 1K)	HuggingFace / Megatron-LM	up to 3× / ~1.7× faster end-to-end
Long Range Arena	standard attention	2.4× speedup
Path-X (16,384 tokens)	all prior Transformers	first to beat chance (61.4%); block-sparse variant reached Path-256 at 64K tokens
Attention memory	quadratic baseline	linear in N — up to ~20× less at long lengths

And the structural comparison that explains why:

	Writes full N×N matrix?	Attention memory	Accuracy
Standard attention	Yes — to slow HBM	Quadratic in N	Exact
Sparse / low-rank	No	Lower	Approximate
FlashAttention	No	Linear in N	Exact

What it did not solve

FlashAttention makes attention IO-efficient, not sub-quadratic — the arithmetic is still O(N²); only the memory traffic became linear. The backward pass recomputes attention on the fly (trading FLOPs for memory), and each hardware generation needs its kernel rewritten — which is exactly what FlashAttention-2 and -3 turned out to be.

Why it still matters

FlashAttention became a default building block baked into every major training and inference stack — PyTorch's scaled_dot_product_attention, vLLM, TensorRT-LLM, xFormers. If you've used a model with a long context window, you've almost certainly used it. And it grew into a lineage:

Version	Year	What changed
FlashAttention	2022	Tiling + online softmax + recomputation; IO-aware exact attention
FlashAttention-2	2023	Better parallelism & work partitioning between GPU thread blocks and warps — roughly 2× again, reaching 50–73% of A100 peak FLOPs
FlashAttention-3	2024	Rebuilt for Hopper (H100): warp specialization, asynchronous TMA copies, FP8 support — another ~1.5–2×

Its deeper lesson is one every systems engineer learns eventually: on modern hardware, data movement is the cost, not arithmetic. The fastest code isn't the one that does the least math — it's the one that moves the least memory. FlashAttention is that principle, applied to the single most important operation in modern AI.

Check yourself

5 questions on this paper

Answer each one out loud before you open it. Recalling an answer is what makes it stick; recognising one you have just read is not. Your score is kept on this device.

What is FlashAttention?

Show answer

FlashAttention is a 2022 algorithm that computes exact attention faster and with far less memory by being IO-aware — minimizing data movement between a GPU's slow high-bandwidth memory (HBM) and its fast on-chip memory (SRAM). It never writes the full attention matrix to slow memory; instead it processes attention in tiles that fit in SRAM, using an online softmax to combine them.

Why was standard attention slow?

Show answer

Standard attention builds the full N-by-N scores matrix and writes it to the GPU's slow HBM, then reads it back for the softmax and again for the value multiply. For long sequences that matrix is huge, and all that reading and writing of slow memory — not the arithmetic — is the bottleneck. Attention was memory-bound, but everyone had been optimizing the math.

What is the difference between SRAM and HBM?

Show answer

HBM is the GPU's large main memory — many gigabytes, but relatively slow. SRAM is tiny on-chip memory — only kilobytes to megabytes, but roughly an order of magnitude faster. FlashAttention's whole strategy is to do as much work as possible inside fast SRAM and touch slow HBM as little as possible.

Is FlashAttention exact or approximate?

Show answer

Exact. Unlike earlier efficient-attention methods that approximated the attention matrix (sparse or low-rank), FlashAttention computes the mathematically identical result as standard attention — just with a smarter memory schedule. You lose no accuracy; you only save time and memory.

How does FlashAttention enable long context windows?

Show answer

Standard attention's memory grows with the square of the sequence length because it stores the full scores matrix. FlashAttention never materializes that matrix, so its memory grows only linearly with sequence length. That makes much longer context windows practical, which is why it became a default building block for modern long-context models.

Where this sits

Serving it cheaply

Five papers about the bill. Each one buys throughput back from a different part of the stack: the weights, the memory hierarchy, the KV cache, the decode loop.

All 110 papers, in 19 ladders →

Finished this one? 0 / 110 Paper Breakdowns done

Explore the topic

See this alongside everything else on the same subject — handbooks, system designs, challenges and tools, in one place.

LLM Engineering

FlashAttention,
explained.

What changes against every attempt to make attention cheaper: the diagnosis

The real bottleneck

Fast memory, slow memory

Tiling and the online softmax

The math, worked with real numbers

Exact, not approximate

What the authors admit, and what happened next

Online softmax & tiled attention, for real

What the paper actually measured

Why it still matters

5 questions on this paper

What is FlashAttention?

Why was standard attention slow?

What is the difference between SRAM and HBM?

Is FlashAttention exact or approximate?

How does FlashAttention enable long context windows?

Serving it cheaply

Explore the topic

More Paper Breakdowns

Explore more from Vibe Engines

FlashAttention,explained.

What changes against every attempt to make attention cheaper: the diagnosis

The real bottleneck

Fast memory, slow memory

Tiling and the online softmax

The math, worked with real numbers

Exact, not approximate

What the authors admit, and what happened next

Online softmax & tiled attention, for real

What the paper actually measured

Why it still matters

5 questions on this paper

What is FlashAttention?

Why was standard attention slow?

What is the difference between SRAM and HBM?

Is FlashAttention exact or approximate?

How does FlashAttention enable long context windows?

Serving it cheaply

Explore the topic

More Paper Breakdowns

Explore more from Vibe Engines

FlashAttention,
explained.