01 · FlashAttention

Never materialize the N×N score matrix.

Attention wants a giant score matrix: 67 million entries per head at 8K tokens. Standard attention writes that matrix to slow GPU memory, reads it back, and writes it again. FlashAttention computes the same result tile by tile on fast on-chip memory, so the full matrix never exists in HBM. Temporary score tiles still exist. They just live and die on the workbench.

dao et al. 2022 · flashattention 1–3

Read the paper

step 1 of 5

find what “it” attends toin progress

The token

The whole of attention is one question: how much should each token listen to every other token? To ask it, every token prepares three vectors. The query is the question it wants answered, the key is the description it advertises, and the value is the actual content it can share. A query scores every key with a dot product, so every pair of tokens gets exactly one number. That number is one cell of the score matrix. Click tokens below to see how strongly “it” listens to each of them, and find the strongest link.

what does “it” attend to?

tap tokens to probe

“it” is the query. Click a token to probe how strongly “it” attends to it. One cell of the score matrix is exactly that number.

step 2 of 5

scores → softmax → weighted sumin progress

Ordinary attention

Ordinary attention, the version from the 2017 transformer paper, does three moves. Move one: compute a score for every pair of tokens, filling a complete N by N matrix with one cell per query-key pair. Move two: run softmax down each row, turning the raw scores into weights that add up to exactly 1, so the row reads like a probability distribution. Move three: mix the value vectors using those weights, one weighted sum per token, and that mix becomes the output. Watch the three moves run below, then zoom out to 8,192 tokens and watch the matrix become 67 million cells.

let it run, or drive it yourself with the buttons

0/3

The

anima

didn’

cross

the

road

The

animal

didn’t

cross

the

road

six tokens sit in a row. attention will let each one consult the others.

rows · tokens asking questions
columns · tokens being considered
cell · one query–key similarity score

step 3 of 5

where the matrix should and shouldn’t livein progress

Store-room & workbench

Attention’s math is cheap. The expensive part is moving numbers around, and a GPU has two very different places to put them. HBM is the store-room: it holds gigabytes, but every trip there costs time, and its bandwidth is the wall attention keeps hitting. SRAM is the workbench: only a few megabytes, but it sits right next to the compute, so touching it is nearly free. Standard attention writes the whole score matrix to the store-room, reads it back for softmax, writes the result, and reads it yet again, about 4N² elements cross the floor per run. FlashAttention does every score on the workbench and only sends the final output back. Step through both paths below and watch each crossing.

floor crossings so far: 0

the floor · every crossing costs time

tokens

QKV

sram · workbench · tiny & fast

on the bench: S · all n² scores

hbm · store-room · huge & slow

stored: Q, K, V

standard · move 1/6 · compute every score at once: the full n×n matrix S is built on the workbench. it does not fit, and that is the whole problem.

crossing the floor: nothing

step through the moves, or let them loop

one full run ≈ 4n² + n·d elements across the floor

step 4 of 5

watch m, ℓ, O survive while tiles diein progress

The tile loop

FlashAttention cuts the job into tiles that fit the workbench, and every tile runs the same four-stage loop. It loads one Q tile, loads one K,V tile, and computes a single small score tile on the bench. It then merges that tile into three running numbers it keeps for each query row: m is the largest score seen so far, ℓ is the running sum that renormalizes, and O is the partial output accumulated so far. Whenever a tile sets a new max, everything seen before is rescaled by the correction factor e^(m_old − m_new). you will see that happen with live numbers at stage 3. Those three numbers carry everything the next tile needs, so the score tile is thrown away and only m, ℓ and O move on. This is called online softmax, and it is why FlashAttention can be exact without ever holding the full matrix. Watch the running-state box on the right.

tile 1/16 · stage 1/4

query block 1/4 · kv block 1/4 · stage 1: load one Q block onto the workbench, one strip of queries.

The grid is the N×N score matrix, cut into tiles. Each square is one query block scored against one key-value block; the bright tile is the one on the workbench right now, and the faded ones were already folded into the running state and discarded. The row and column bands mark exactly which tokens are in play.

running state · row 5

m · biggest score so far…

ℓ · running softmax sum…

O · partial output…

These three numbers, not the score tile, are all that survive between tiles. That is online softmaxThe trick that makes tiling exact: softmax can be recomposed from a running max, sum and partial output, so each row never needs to see all its scores at once., the reason a finished tile can safely disappear.

step 5 of 5

the same math, a fraction of the trafficin progress

The payoff

Here is the whole point, side by side. Standard and flash run the exact same attention, so the answers are identical. the compute bars below are literally the same length. The only difference is how many elements cross between memory and chip, and that is what the coral waiting time is made of. Drag the sequence length and watch what happens: standard traffic scales like the square of the length, flash scales almost linearly, so the gap, and the speedup, widens fast. The score matrix that standard materializes is the cost. Flash simply never builds it.

real-scale accounting · d = 128

sequence length8k

drag the slider and watch the traffic gap widen

full score matrix / head

67.1Mentries

flash keeps on-chip

16Kat a time

hbm traffic reduction

3.7×

standard · hbm traffic138.4M elems

flash · hbm traffic37.7M elems

where the time goes · one full run of the layer

Each bar is wall-clock time, split in two. Coral is the GPU sitting idle, waiting for numbers to arrive from the store-room at 2 TB/sAssumed HBM bandwidth: 2 TB/s, roughly an A100-class GPU.. Indigo is real computing at 312 TFLOP/sAssumed math throughput: 312 TFLOP/s of fp16 matmul, roughly an A100-class GPU.. The compute part is identical for both: same algorithm, same FLOPs. Only the waiting shrinks.

waiting on memory computing

standard193 µs

waiting 138 µs

computing 55 µs

72% of the run is pure waiting · moves 138.4M elements

flash2.1× faster93 µs

waiting 38 µs

computing 55 µs

41% of the run is pure waiting · moves 37.7M elements

The same attention algorithm, without approximation, computed with a different memory schedule. Longer sequences become affordable; training and inference get faster.

show the accounting +

standard: 2N² + 4Ndd: the dimension of each head’s query/key/value vectors (128 here). → 138.4M elements moved

flash: 4Nd + 2N²d²/MM: the number of elements that fit in on-chip SRAM (64K here). → 37.7M elements moved

floating-point addition order differs between the two schedules, so outputs agree only within numerical tolerance. the official tests verify equality up to that tolerance, not bit-for-bit.

canonical paper

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness Dao et al. · 2022

Loading the interactive visualization…

01 · Reference notes

← All visual guides

FlashAttention Explained

FlashAttention is an exact, IO-aware attention algorithm that reduces GPU memory traffic by computing attention in tiles in fast on-chip memory instead of materializing the full N×N attention matrix in HBM.

01 · Definition

What is FlashAttention?

It changes the schedule for the usual attention calculation, not the calculation itself. The output is the same as standard softmax attention; the win comes from moving far less intermediate data to and from slow GPU memory.

02 · Mechanism

Why standard attention is slow

For a sequence of N tokens, standard attention forms an N×N score matrix. At long context lengths that matrix is enormous. Writing scores to high-bandwidth memory (HBM), reading them back for softmax, and writing weights again can cost more time than the matrix multiplications themselves.

03 · Mechanism

HBM vs SRAM

HBM is large GPU memory: it stores a lot, but moving data there is expensive. SRAM is the much smaller on-chip workspace next to the compute units. FlashAttention organizes work so temporary score blocks fit in SRAM and are discarded before they ever need to live in HBM.

04 · Mechanism

How FlashAttention tiling works

The algorithm loads a block of queries and a block of keys and values, computes their small score tile, and immediately folds that tile into the output. It repeats this for the next key/value block. Only the final output and a small running state survive between tiles.

05 · Mechanism

How online softmax stays exact

Softmax usually appears to require every score in a row at once. FlashAttention instead tracks each row's running maximum m, normalizing sum ℓ, and partial output O. When a later tile contains a larger score, the earlier contribution is rescaled. Those three values preserve the same result as a full-row softmax.

06 · Mechanism

Is FlashAttention approximate?

No. FlashAttention is exact attention with a different memory-access pattern. It does not select fewer tokens or replace softmax; it avoids storing the full intermediate matrix.

07 · Mechanism

FlashAttention vs standard attention

Both compute QKᵀ, apply softmax, and combine V. Standard attention materializes large intermediates in HBM. FlashAttention processes the intermediates in SRAM-sized tiles, which makes the gap more valuable as sequences get longer.

Source · Original paper

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness ↗Dao et al. · 2022