FlashAttention: attention was never compute-bound

Dao, Fu, Ermon, Rudra & Ré, 2022

The n2 problem everyone quotes is not the multiplying. It's the moving.

Download flash_attention_from_scratch.py — everything below runs from it. NumPy only, no GPU required.

I'd read the attention page — mine, and everyone else's — and walked away with the same wrong conclusion everybody else does: attention is slow because it does n2 work. So I spent a year assuming the fix had to be doing less work: sparse patterns, low-rank tricks, sliding windows. Then FlashAttention shipped, did exactly the same arithmetic, and was several times faster on the same hardware.

Nobody had told me the n2 in that sentence is a bytes figure, not a flops figure.

You already know this bug

You have written this function. It works on the test file and dies on the real one: 

rows = json.load(open("events.json"))   # 40GB. Your laptop has 16.
rows = [normalise(r) for r in rows]
rows = [enrich(r) for r in rows]
return summarise(rows)

Three passes, three full materialisations, and the machine spends its life in swap. You know the fix without thinking about it. You stream: 

total, count, best = 0.0, 0, None
for chunk in read_in_chunks("events.json"):     # constant working set
    for r in chunk:
        r = enrich(normalise(r))
        total += r.value; count += 1           # carry running statistics
        best = r if best is None else max(best, r)
return total / count, best

Same answer. One pass. Constant memory. That is the entire paper. The 40GB file is the n×n score matrix, the swap is HBM, and the running statistics are a max and a sum — which turn out to be exactly what softmax needs.

TermMeansProgrammer's analogy
HBMThe GPU's main memory. ~40GB, ~1.5TB/s. Where your tensors live.Disk. Big, and you pay every time you touch it
SRAMOn-chip scratch memory, ~100KB per streaming multiprocessor, an order of magnitude fasterL1 cache. Tiny, free, yours only briefly
Memory-boundThe arithmetic units are idle waiting for data to arriveA CPU-bound job that turned out to be IO-bound
TilingProcess the matrix in blocks that fit in SRAMChunked reads. read(8192), not read()

Why this beat the approximations

Between 2019 and 2022 there was a whole genre of "efficient attention" papers: Reformer, Linformer, Performer, Longformer, BigBird. Each cut the asymptotic flop count by computing a cheaper, different function — a sparse pattern, a low-rank projection, a local window — and then argued the quality loss was acceptable. Most were not actually faster in wall-clock time, because they cut the thing that wasn't the bottleneck. FlashAttention cut the thing that was, and computes exactly the same function as the attention you already have. No quality argument to make, no ablation to defend.

This is the interview answer. "FlashAttention is an IO result, not an arithmetic one." The paper's own subtitle is Fast and Memory-Efficient Exact Attention with IO-Awareness. The word doing the work in that title is exact.

The one formula

Not the attention formula — you have that already. This is the online softmax update, the piece of maths that makes streaming possible. You are seeing scores one block at a time and carrying two numbers: m (the biggest score so far) and l (the running sum).

mnew = max(mold, max Sblock)
lnew = e(mold − mnew) lold + Σ e(Sblock − mnew)
Onew = e(mold − mnew) Oold + e(Sblock − mnew) Vblock

Read it as pipeline stages:

StepShapeWhat it does
Sblock(Br, Bc)One score tile. Sized to fit in SRAM. Never written to HBM, never seen again
m(Br, 1)Running maximum per query row. The same overflow-safety max as ordinary softmax — it just learned to work incrementally
e(mold−mnew)scalar per rowThe rescale. The load-bearing line. Everything accumulated against the old max gets converted to the new baseline
l(Br, 1)Running denominator. Applied once, at the very end
O(Br, d)Running unnormalised output. Rescaled alongside l, because it is a weighted sum whose weights moved

Where the time actually goes

I can't benchmark a GPU I don't own and won't pretend to, so the script counts instead: every byte that would cross between fast and slow memory, every multiply-add, tracked by hand in a Counters object. Counting reads and writes is literally what the paper does.

Don't take my word for it. Here's the actual output:

head dim d = 64, float32

      n     HBM MB    GFLOPs  bytes/flop  n^2 share
    128       0.38     0.004       0.094      66.7%
    512       4.50     0.067       0.070      88.9%
   2048      66.00     1.074       0.064      97.0%
   8192    1032.00    17.180       0.063      99.2%

17 GFLOPs is nothing — a modern GPU eats that in under a millisecond. But it comes with a gigabyte of memory traffic, and 99.2% of that gigabyte is the score matrix being written out and read back three times. The one thing you never wanted to keep is the only thing you're paying for.

The ratio settles at about 0.06 bytes per flop. An A100 does roughly 19.5 TFLOP/s of fp32 against roughly 1.5 TB/s of HBM — it can only feed itself about 0.08 bytes per flop. Attention sits the wrong side of that line: the multiplying waits on the moving.

Online softmax: the whole idea, in two variables

You already met the running-max trick on the attention page, where it existed purely so exp() wouldn't overflow. But "subtract the max" assumes you can see the whole row, and streaming you can't. So you carry the max seen so far, and when a later block contains a bigger one you go back and fix everything already accumulated — by multiplying it by e(mold − mnew). That correction is one line, and deleting it raises nothing. Same row of 32 scores, streamed in four blocks of eight, with and without it: 

ordinary softmax vs online softmax, max abs diff: 0.00e+00  exact
ordinary softmax vs NO-RESCALE version, max abs diff: 6.53e-01  wrong
no-rescale weights sum to 0.190131, not 1.0

first 6 weights, side by side:
  correct:  [0.003957 0.016505 0.003782 0.000028 0.021219 0.005355]
  online:   [0.003957 0.016505 0.003782 0.000028 0.021219 0.005355]  ✓
  broken:   [0.000752 0.003138 0.000719 0.000005 0.004034 0.001018]

Zero difference — not "close enough", zero — against ordinary softmax. And without the rescale, weights that sum to 0.19. You've added numbers that were scaled against three different baselines, like summing a column of prices in three currencies. Nothing throws. The answer is just wrong.

This trick isn't Dao's, and the paper says so: it's Milakov & Gimelshein's online normalizer calculation from 2018. FlashAttention's contribution is running the whole attention pipeline — the @ V included — through it, so the n×n matrix never needs to exist.

Tiling: same answer, flat memory

Now tile it. Q gets cut into row blocks, K and V into column blocks. For each query block you walk every key block, carrying m, l and a running output O per row. When m moves, both l and O get rescaled. 

      n   max abs diff   std peak KB   flash peak KB   std HBM MB  flash HBM MB
    128       2.50e-16          64.0            64.5         0.38          0.19
    256       1.94e-16         256.0            64.5         1.25          0.62
    512       1.39e-16        1024.0            64.5         4.50          2.25
   1024       9.02e-17        4096.0            64.5        17.00          8.50
   2048       8.33e-17       16384.0            64.5        66.00         33.00

The difference column is float noise at 1e−16. The standard peak goes up 16× for every 4× in n; the flash peak does not move, because the working set is two tiles and two running statistics, none of which depend on sequence length. That flat column is why 32k and 128k contexts stopped being a memory question.

Quadratic against linear, in one table. The quadratic term in attention's memory never went away in the maths — it went away in the allocation. You still compute every one of the n2 scores. You just never keep more than a tile of them at a time.

How big should a tile be?

Bigger blocks mean fewer passes over K and V, so less traffic. Why not make the block the whole sequence, then? Because the tile has to fit in SRAM; past that it spills back to HBM and you have reinvented standard attention with extra loop overhead. 

n = 1024, d = 64, SRAM budget assumed 96 KB
standard attention traffic: 17.00 MB

  block     HBM MB   tile working set KB  fits SRAM?
     16      32.50                  13.1  yes
     32      16.50                  28.2  yes
     64       8.50                  64.5  yes
    128       4.50                 161.0  NO -- spills
    256       2.50                 450.0  NO -- spills
    512       1.50                1412.0  NO -- spills
   1024       1.00                4872.0  NO -- spills

Block 16 costs 32.5 MB, worse than standard attention's 17 MB, because K and V get re-read 64 times. Block 128 would be lovely at 4.5 MB and doesn't fit. The optimum is the largest tile that still fits — here, 64. On real hardware that budget is shared memory per streaming multiprocessor: an A100 has 164 KB per SM, of which a kernel can typically use around 96–100 KB. That number and the head dimension pick the block size, not the maths.

The backward pass: pay flops, save bytes

Standard attention's backward pass wants the n×n probability matrix that the forward pass computed, so it stashes it. FlashAttention refuses, keeps only m and l — two floats per row — and recomputes the score tiles on the way back. 

      n   std stored MB   flash stored MB   std GFLOPs   flash GFLOPs  extra flops
    512            1.38              0.38         0.17           0.20          20%
   2048           17.50              1.52         2.68           3.22          20%
   8192          262.00              6.06        42.95          51.54          20%

At n=8192 that's 256 MB saved per attention head for 20% more arithmetic, and it still comes out faster in wall-clock, because — say it with me — the arithmetic was never the bottleneck. You have made this trade before: it is what gradient_checkpointing=True does in every training config you've ever set. FlashAttention applies it to one operator instead of a whole layer, and the accounting is so favourable it stops being a trade — you get the memory and the speed.

Exact means exact

This is the claim people get wrong, so here it is against a genuine approximation. The comparison method is a 128-wide local window — a stand-in for the sparse family FlashAttention displaced. 

flash vs standard, max abs diff:  1.665e-16
flash vs standard, allclose:      True
local-window vs standard, max abs diff: 3.224e-01
local-window vs standard, allclose:     False

1.665e−16 is accumulation order, not approximation. 0.32 is a different function wearing attention's name. That gap is the whole argument for why FlashAttention was everywhere within a year and Linformer was not.

And no, it is not doing less arithmetic

The belief I want to kill. Every multiply the standard version does, the tiled version also does: 

      n   std GFLOPs   flash GFLOPs   std HBM MB   flash HBM MB  traffic cut
    512        0.067          0.067         4.50           2.25         2.0x
   1024        0.268          0.268        17.00           8.50         2.0x
   2048        1.074          1.074        66.00          33.00         2.0x
   4096        4.295          4.295       260.00         130.00         2.0x

Identical flop columns. Only traffic moves. (My 2× is modest because I hardcoded a 64-wide tile for a 64-dim head — the ratio here scales with block size over head dimension, which is what demo 4 is about. The paper reports substantial wall-clock speedups from real kernels; I can't reproduce those on a laptop, so I'll leave the figures to the paper.)

FlashAttention-2, briefly

Dao's 2023 follow-up doesn't change the idea. Same tiling, same online softmax, same exactness. What changes is the CUDA work partitioning: fewer non-matmul operations (the rescaling gets restructured so it happens less often), parallelism over the sequence dimension as well as batch and heads, and a better split of work between warps inside a thread block so they share less through shared memory.

It matters enormously in practice and it is not conceptually new. If you understood the tile and the running statistics, you have understood FlashAttention-2 too. FlashAttention-3 goes further again, targeting Hopper-specific hardware.

The code

Pure NumPy. No torch, no GPU, no fake benchmark. The memory hierarchy is a counter object — every read and write is added up by hand, because reads and writes are the paper's actual subject. Every function under 30 lines.

Download the full script — about 180 lines of mechanism plus nine runnable demos. softmax and standard_attention are in there too: the ones from the attention page, with counter calls added. Then: 

python3 flash_attention_from_scratch.py
def online_softmax(row, block=8, rescale=True):
    m = -np.inf        # running max
    l = 0.0            # running sum of exp(x - m)
    for j in range(0, len(row), block):
        chunk = row[j:j + block]
        m_new = max(m, chunk.max())
        if rescale:
            l = l * np.exp(m - m_new) if m > -np.inf else l
        l = l + np.exp(chunk - m_new).sum()
        m = m_new
    # Second pass only because we want the full vector back for comparison;
    # attention does not need it, because it accumulates O as it goes.
    return np.exp(row - m) / l


def flash_attention(Q, K, V, c, Br=64, Bc=64, rescale=True, keep_full_S=False):
    n, d = Q.shape
    O = np.zeros((n, d))

    for i in range(0, n, Br):
        Qi = Q[i:i + Br]
        c.move(Qi.size)                       # load the query tile
        Oi = np.zeros_like(Qi)
        mi = np.full((Qi.shape[0], 1), -np.inf)
        li = np.zeros((Qi.shape[0], 1))
        # working set: one Q tile, one K tile, one V tile, one score tile,
        # plus the two running statistics. Independent of n.
        c.scratch(Qi.size + 2 * Bc * d + Qi.shape[0] * Bc + 2 * Qi.shape[0])

        for j in range(0, n, Bc):
            Kj, Vj = K[j:j + Bc], V[j:j + Bc]
            c.move(Kj.size + Vj.size)         # load the key/value tile

            Sij = Qi @ Kj.T                   # (Br, Bc) -- lives in SRAM only
            c.matmul(Qi.shape[0], d, Kj.shape[0])

            m_new = np.maximum(mi, Sij.max(axis=1, keepdims=True))
            P = np.exp(Sij - m_new)           # unnormalised weights
            if rescale:
                correction = np.exp(mi - m_new)
                correction[np.isnan(correction)] = 0.0   # -inf minus -inf
                li = correction * li + P.sum(axis=1, keepdims=True)
                Oi = correction * Oi + P @ Vj            # rescale the output too
            else:
                li = li + P.sum(axis=1, keepdims=True)   # THE BUG
                Oi = Oi + P @ Vj
            c.matmul(Qi.shape[0], Kj.shape[0], d)
            mi = m_new

        O[i:i + Br] = Oi / li                 # normalise once, at the end
        c.move(Oi.size)

    return O

Now break it

This is the part that actually makes it stick. Reading won't do it. You debug for a living — reverse-engineering from broken behaviour is your native mode, and every row here is a real edit to the script you just downloaded.

Break thisWatch for
Pass rescale=False to flash_attention — drop the correction * on l and ODemo 2's row: weights sum to 0.190131 instead of 1.0, max abs error 6.53e−01. No exception, no warning.
Call flash_attention(..., Br=512, Bc=512) on demo 4's SRAM budgetWorking set jumps to 1412.0 KB against a 96 KB budget. On real hardware the tile spills to HBM and the traffic saving evaporates.
Set Br=Bc=16Traffic climbs to 32.50 MB — worse than standard attention's 17.00 MB, because K and V get re-read 64 times.
Pass keep_full_S=True to keep the score matrix "just for debugging"Peak extra memory goes 64.5 KB → 4096.0 KB. 64× worse. The tiling saved you nothing; the saving was in never allocating.
In online_softmax, replace np.exp(chunk - m_new) with np.exp(chunk)Demo 9's failure: on scores near 800, exp() returns inf and the weights come back [nan nan nan nan].
Believe the speedup is fewer flops — then read demo 7's flop columns0.067 / 0.067, 0.268 / 0.268, 1.074 / 1.074. Identical. Only the traffic column moves.
Then read the paper — the introduction, the algorithm section with Algorithm 1, and the IO-complexity analysis that follows it. Skip the CUDA-kernel details and the benchmark tables on a first pass. The proof you actually want is the one showing standard attention moves Θ(nd + n2) bytes while the tiled version moves Θ(n2d2/M), where M is the SRAM size — that's demo 4, written as a theorem.

On the maths anxiety

There is almost none here, and what there is you already have. The online softmax update is one algebraic identity: e(a−c) = e(a−b)·e(b−c). That's it. That's the paper's maths. If you can see why rescaling a running total by a constant factor keeps it correct, you can derive the update yourself.

What the paper actually demands is a skill you already have: knowing a memory hierarchy exists and that crossing it costs. FlashAttention took five years to appear after the transformer not because the maths was hard, but because ML people were counting flops.

Where to go next

This page gave you one thing: the traffic counter, and why the number in it is the whole story. It deliberately skips the CUDA, the proof, and anything about inference. Four people cover those better than I would.

Order I'd take them in: Horace He first for the mental model, then the FlashAttention paper, then Milakov if the softmax update still feels like magic, then FlashAttention-2 last. And if you want the same memory-bound argument at inference time rather than training time, that's the KV cache page — same lesson, different bottleneck. The maths underneath both is on the attention page.