Flash Attention, From First Principles

In this article we are going to see, in detail, flash attention. Attention Is All You Need Vaswani et al., 2017 was a watershed moment in the world of generative modeling, especially in NLP. Hard to believe it was nearly a decade ago — and here we are, in a world where LLMs are chipping away at long-unsolved Erdős problems , designing drugs (Nobel Chemistry 2024 — AlphaFold) , and (for better and worse) even orchestrating cyber operations (Anthropic, 2025) . Agents are an integral part of our lives now, and in this article we are going to see one of the atomic concepts that powered this revolution. Sit tight.. You will enjoy this ride, I can assure you that.

What if I told you — you can do MORE computation, in LESS time, and still deliver the exact same output? Not an approximation. The same attention, computed in a different order. We are generous on computation but frugal on time; we would rather compute — or recompute — than wait. In this article we will see exactly how and why flash attention is faster while being exactly attention itself.

Don't take my word for it. Here is the receipt, straight from the paper's own measurements (FlashAttention paper, Fig. 2 — GPT-2 medium, N=1024, A100):

Read that again. Flash attention does MORE arithmetic — and finishes almost six times faster, moving nine times less data. How on earth?

Here is my promise: by the end of this article, you will not have read the answer — you will have derived the machine that produces the first two of those numbers, with a pencil, starting from a 3×2 matrix: why the work goes up, why the data movement collapses, and how far each one moves. The work and the data movement are both yours to count. The third number, the clock, you will not derive — and the reason why is one of the more interesting things in here. You will earn that one by measuring it. Let's go.

What is attention?

$$O = \mathrm{softmax}(QK^\top)\,V$$

notebox — the missing $1/\sqrt{d}$. you may have met this formula as $\mathrm{softmax}(QK^\top/\sqrt{d})\,V$ — scores divided by $\sqrt{d}$ before the softmax, so that dot products of long vectors don't grow large enough to push softmax into its flat corners. I drop that factor everywhere in this article, deliberately. It is a fixed rescaling applied before the softmax, so it changes none of the counting, none of the two walls, and none of the repair algebra we are about to derive — it would only add a symbol to every line. It does matter in practice, and the companion code applies it: you will find 1/sqrt(d) in both implementations.

I know some of you may be thinking: we have been seeing this equation forever, what do you offer new? My answer — nothing new, technically speaking. But what I promise is that by the end of this article you would have internalised it in an intuitive fashion, derived every piece of it with your own hands. That's the dimension I really care about. It's a promise, and I believe I can deliver.

Let's go to the smallest case we can comfortably hold in our heads. (I was tempted to go square, but am deliberately staying rectangular — later sections live in rectangular block sizes $B_r$, $B_c$, and I don't want the two side-lengths masquerading as each other. Square cases we reserve for simplifying upper-bound calculations, and we will say so when we do.)

Consider three matrices, each of shape 3×2:

$$Q = \begin{bmatrix} 1 & 0 \\ 0 & 2 \\ 1 & 1 \end{bmatrix} \qquad K = \begin{bmatrix} 1 & 1 \\ 0 & 1 \\ 2 & 0 \end{bmatrix} \qquad V = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{bmatrix}$$

terminology freeze — pin this table before anything else, we will use it for the entire article:

notebox — the two levers. $d$-wide things are representations (a token's vector). $N$-wide things are score-lists (a token judged against everyone). Keep them separate in your head at all times; half the confusion in this subject is these two masquerading as each other.

Now, I ask you to compute $O$. The journey has three stops.

Stop 1: the raw scores

$$S = QK^\top = \begin{bmatrix} 1 & 0 & 2 \\ 2 & 2 & 0 \\ 2 & 1 & 2 \end{bmatrix}$$

Note $S$ is square, unlike its inputs — and think about what it means. Its rows are the $N$ tokens... And its columns? You are right — also the $N$ tokens. Entry $S_{ij}$ is the dot product of token $i$'s query row with token $j$'s key row: token $i$ asking, token $j$ being asked. $S$ is the full questionnaire — every token interrogating every token.

And pay close attention: $S_{12} = 0$ but $S_{21} = 2$. They need not be the same. The importance of token 2 to token 1 is a different question from the importance of token 1 to token 2 — relevance in language is directional ("sat" urgently needs "cat"; "cat" barely needs "sat").

notebox — why two matrices? if each token had just ONE vector $x_i$, the score would be $x_i \cdot x_j$ — and dot products don't care about order, so the matrix would be forced symmetric. Giving each token separate query and key vectors is precisely what buys the freedom to be asymmetric. (can $S_{12} = S_{21}$ still happen? Sure — specific vectors can align that way, and if you set $Q = K$ the whole matrix turns symmetric. The split doesn't forbid symmetry; it stops forcing it.)

Stop 2: from scores to proportions — why softmax?

Here is the plan for stop 3, stated upfront so stop 2 has a purpose: each token's output will be a blend of the value rows — some amount of $v_1$, some of $v_2$, some of $v_3$, mixed according to relevance. And there is the problem: raw scores cannot be blend amounts. A score of 2 is not "2 of something"; scores can even be negative, and a recipe cannot call for −0.3 cups of milk. We need to convert each row of raw scores into legal proportions. Normalization is not a math ritual — it is what turns opinions into recipe amounts.

So let's design the converter. Take a row vector and ask: how much share does each element deserve?

candidate 1 — divide by the sum. try $s_{\text{row}} = [1, 2, 3]$: sum is 6, shares $[1/6,\, 2/6,\, 3/6]$. Works! Now try $s_{\text{row}} = [1, -1, 0]$:

$$\text{sum} = 0 \qquad s_{\text{norm}} = [\,1/0,\; -1/0,\; 0/0\,]$$

Oops. The machine just beeps. Division by zero — and even when the sum isn't zero, negative "shares" sneak through. Dead.

candidate 2 — take modulus, then divide by the sum. for $[1, -1, 0]$: modulus sum is 2, shares $[1/2,\, 1/2,\, 0]$. Positive, sums to one... But do you see the problem? Try $[1, 1, 0]$:

$$[1, -1, 0] \;\longrightarrow\; [1/2,\, 1/2,\, 0] \qquad [1, 1, 0] \;\longrightarrow\; [1/2,\, 1/2,\, 0]$$

Two vectors living in different parts of the space — one dislikes its second element, the other likes it — collapse to the identical shares. The modulus erased the sign, and the sign carried the opinion. Semantic confusion.

So let's write down what we actually demand. three properties, and notice each dead candidate died against a different one:

1. every share ≥ 0 — a recipe holds no negative amounts (candidate 1 died here, plus the fire)
2. shares sum to 1 — proportions of a whole (think percentages)
3. faithfulness — the converter must preserve what the scores said: higher score, bigger share; dislike stays small; different opinions land on different shares (candidate 2 died here — it met properties 1 and 2 perfectly and still lobotomised the meaning)

notebox — the weights ambush. the word "weights" is doing double duty in ML. model weights — learned parameters — are signed floats, negative all the time. attention weights — these shares we are building — are computed fresh from each input and must be legal proportions. Same english word, different mathematical citizens. When this article says weights, it means the shares.

The survivor: the exponential — followed by the very same divide-by-the-sum we started with, which works now because $e^x$ makes everything positive first. $e^x > 0$ always, and $e^x$ is strictly increasing, so order is preserved and different opinions stay different. That pair — exponentiate, then normalize — is softmax:

$$\mathrm{softmax}(s)_j = \frac{e^{s_j}}{\sum_k e^{s_k}}$$

Check it on the killer case $[1, -1, 0]$:

$$e^1 + e^{-1} + e^0 \approx 2.718 + 0.368 + 1 = 4.086 \qquad \Rightarrow \qquad [\,0.665,\; 0.090,\; 0.245\,]$$

All shares positive, sum to 1, and the disliked element gets a small but honest share. And on $[1, 1, 0]$: $[\,0.422,\; 0.422,\; 0.155\,]$ — different vector, different shares, no collapse. Faithfulness lives.

notebox — a name for later. numbers that are all ≥ 0 and sum to 1, used to mix vectors, make what mathematicians call a convex combination — the blend can never escape the region spanned by its ingredients. Remember the phrase; it pays rent in stop 3.

Stop 2.5: don't set the machine on fire

Now consider $s_{\text{row}} = [101, 100, 102]$. Apply softmax directly: $e^{101} \approx 10^{43.9}$. The machine beeps again — float16 overflows at ~65,504, and even float32 dies near $e^{89}$. These are perfectly reasonable scores and the formula is uncomputable.

The rescue: subtract the row's max from every element first. Max is 102, so

$$[101, 100, 102] \;\longrightarrow\; [-1, -2, 0]$$

Every shifted score ≤ 0, every exponential ≤ 1. Fireproof forever. But hang on — am I allowed to just mutilate the input because its numbers are inconvenient? Here is the beautiful part:

s₁ = [101, 100, 102] and s₂ = [−1, −2, 0] have the exact same softmax.

DON'T BELIEVE ME? Three layers of conviction, escalating:

SOFTMAX IS SHIFT INVARIANT. the overflow was the motive for subtracting the max; shift invariance is the license that makes it legal. Motive plus license — remember this pair, it returns with a vengeance later in this article.

Stop 3: the pour

Apply the (max-shifted) softmax to every row of our $S$:

$$P = \begin{bmatrix} 0.245 & 0.090 & 0.665 \\ 0.468 & 0.468 & 0.063 \\ 0.422 & 0.155 & 0.422 \end{bmatrix}$$

(every row $\ge 0$, every row sums to 1 — verify!)

And pour: row $i$ of $O$ is token $i$'s blend of the value rows, mixed by its own proportions:

$$o_i = P_{i1}\, v_1 + P_{i2}\, v_2 + P_{i3}\, v_3 \qquad\qquad O = PV = \begin{bmatrix} 0.910 & 0.755 \\ 0.532 & 0.532 \\ 0.845 & 0.578 \end{bmatrix}$$

Look at what happened geometrically: each output row sits inside the region spanned by $v_1, v_2, v_3$ — the convex combination promise cashing in. Token 1's blend leans hard toward $v_3$ (share 0.665, its top score); token 2 sits nearly on the fence between $v_1$ and $v_2$ (0.468 each). The questionnaire decided where each token stands.

notebox — what softmax does and does NOT constrain. the mandate (≥ 0, sums to 1) lives on $P$ and only on $P$. Rows of $O$ inherit no such law — their entries are whatever the ingredients dictate, negative and unnormalised as real learned values are. softmax disciplines the shares, never the goods.

So what has a token actually DONE by replacing itself with such a blend? Before attention, token 2 was only itself. After, its representation is built from material that lived in other tokens' rows — pulled over in proportion to measured relevance. One token asked, everyone answered, and the relevant answers flowed in.

There is an everyday word for what these tokens just did.

ATTENTION IS HOW TOKENS COMMUNICATE.

Each token asks every token how relevant it is ($S$, the questionnaires), converts the answers into faithful shares ($P$, the laundering), and rebuilds itself as a blend of what the relevant ones were carrying ($O$, the pour). That one sentence is this entire article's anchor — because communication, it turns out, is exactly the thing that gets expensive. Every token talking to every token is $N^2$ conversations. Beautiful for expressiveness. Brutal, as the next section counts, for the courier.

Anatomy of standard attention — the crime scene

By now you are familiar: the journey to $O$ passes through $S$ and $P$. Now let's watch how that journey actually runs on hardware — and count what it costs.

Your inputs $Q, K, V$ all reside in HBM (High Bandwidth Memory) — the GPU's large main memory. Large, but slow relative to where the work happens. The work happens on the SRAM workbench — tiny (a few hundred KB per compute unit) but blazing fast. And between the two runs the courier — the memory machinery hauling numbers back and forth. One law governs this whole article: compute on the bench is nearly free; courier trips are the expensive currency. (paper §2.1: A100 — HBM 40 GB @ 1.5 TB/s, SRAM ~19 TB/s)

(readers of my first article the free-normalization article will recognise this cast — warehouse, bench, courier. If you are new: that article derives why this economy works the way it does; here we just use it.)

Standard attention runs as three separate kernels — three separate programs, and nothing survives on the bench between programs. Watch the courier:

1. Haul $Q$ and $K$ in from HBM to compute $S$ — in: $2Nd$
2. Write $S$ to HBM — out: $N^2$
3. Load $S$ back from HBM to compute $P$ — in: $N^2$
4. Write $P$ to HBM — out: $N^2$
5. Haul $P$ and $V$ in from HBM — in: $N^2 + Nd$
6. Write $O$ to HBM — out: $Nd$

(why wasn't $V$ fetched back in step 1? Because kernel 3 is a different program — the bench was wiped clean twice in between. That wipe IS the crime, as you are about to see.)

Total movement bill:

$$T_{\text{std}} = 4Nd + 4N^2$$

(this is the paper’s Theorem 2, §3.2: standard attention requires Θ(Nd + N²) HBM accesses, while FlashAttention requires Θ(N²d²M⁻¹) — Dao et al., 2022)

Now — how much of this is IRREDUCIBLE? The inputs must come in ($Q + K + V = 3Nd$) and the answer must go out ($O = Nd$). So $4Nd$ is the fundamental cost of the operation, non-negotiable.

Everything else — the entire $4N^2$ — is ghost traffic: four trips that exist for no mathematical reason. Look at them again: $S$ was written out (step 2) only to be immediately read back (step 3). $P$ likewise (steps 4–5). The questionnaire matrix takes two full round-trips through slow memory purely because the pipeline is chopped into three programs that cannot hand objects to each other across the bench. The courier is shuttling paperwork between departments of the same office.

And at real scale the ghost dominates utterly. GPT-2 scale: $N = 1024$, $d = 64$. The ratio of ghost to legitimate is $N^2/Nd = N/d = 16$. out of every 17 numbers the courier hauls, 16 are ghost traffic — roughly 94% of the entire bill. the biggest thing the machine builds, it builds only to throw across the slow wire. Twice.

So the sec1 paradox half-resolves: now we know where the 40.3 GB lives. The questions that remain: can this traffic be reduced at all? And if so — what is the tradeoff we must accept? The next two sections find out.

The two walls

The dream is obvious, so let's state it: fuse the three kernels into one. compute $S$, launder it, pour $O$ — all in a single visit to the bench, so $S$ and $P$ live and die on the bench and never touch HBM. Traffic collapses to the irreducible $4Nd$. Why hasn't everyone always done this? Two walls.

Wall 1 — the bench is too small

Say the bench holds $M = 100{,}000$ numbers (an A100-class budget — roughly one compute unit's SRAM in half precision; treat it as an order-of-magnitude figure, not a datasheet number). Now look at the size of $S$. At GPT-2 scale, $N = 1024$, so $S$ has $N^2 \approx 1.05$ million entries — roughly ten times the entire bench, before $Q$, $K$, $V$, $P$, $O$ ask for their seats. The biggest object in the pipeline cannot visit the workshop whole. So clearly you can't fit $S$ in — forget computing softmax on it and multiplying by $V$.

Fine — so we chunk. Store only a piece of $S$ at a time. And here wall 2 rises.

Wall 2 — the denominator wants the whole row

Think about what softmax needs. Its denominator is $\sum_k e^{s_k}$ — it touches every score in the row — and even the max-shift (our fireproofing license from earlier) needs the row's max, which you cannot know until you have seen the whole row.

And here is the brutal part: the chunking that wall 1 forces runs along the rows. When the machine tiles, each asker-row's scores come into existence in installments — a row of $S$ is that token's questionnaire, 1024 scores long, and it gets minted a few hundred scores at a time as different chunks of $K$ visit the bench. At no moment do all 1024 coexist anywhere. It is like hearing a sentence one phrase at a time, with no replays allowed.

(careful with what gets chopped: a token's own vectors — its $q$, $k$, $v$ rows, each $d$ wide — always travel whole; no one splits a token in half. What arrives in installments is the row of scores. $d$-wide things are representations; $N$-wide things are score-lists — and it is the score-lists that get chopped.)

The interlock

So here is the contradiction, stated plainly:

wall 1 forces chopping. Wall 2 forbids it. the bench cannot hold a row's scores whole; the softmax cannot finish on a fragment. Size forbids wholeness; the denominator demands it.

Every attempt to fuse the kernels smashes into this interlock — which is exactly why the three-kernel crime survived for years. The next section dissolves the contradiction, and it does so in a way I genuinely find beautiful: the machine will commit early, and repair itself.

Online softmax — the game

In the last section we stopped at a contradiction: we could not store whole rows, but softmax needed whole rows. So what do we do? As always when we hit a roadblock — go back to fundamentals, find the single atomic problem, and solve it in its simplest form. The hardware flexing its muscles at scale is just show. We own the brain.

Let's play a game.

$$s = [\,101,\; 100,\; 102\,]$$

(Look familiar? It should — this is row 1 of our anchor, $[1, 0, 2]$, shifted up by 100. Shift invariance says the two rows have identical softmax. Keep that in your back pocket.)

You will be given this vector one element at a time, and your goal is — of course — to compute its softmax (what did you guys expect? 😄). But here is the thing: this is a game, and games have rules:

1. each number arrives once. whatever you want to do with it, do it while its in your hands — you will never see her again!!
2. you may keep a small notebook — a few running scalars of your choosing, updated as you go. But never the raw elements themselves; no stockpiling the vector in disguise. (derived quantities of the number currently in your hands are fair game — you just can't warehouse it.)
3. no time travel. you will not know the later elements in advance. THIS MOMENT IS ALL You HAVE → MAKE IT COUNT.

And the stakes, so you know this is not a puzzle-page filler: if a winning notebook exists, the whole-row requirement dissolves, chunking becomes legal, and the quadratic ghost dies. the entire article hangs on this game being winnable.

One hint before you begin: on a per-row basis the bottleneck is never COMPUTE — you have unlimited freedom to compute, or recompute rather (well, that's a giveaway in itself.. But its fine). So: 101 arrives, and you will never see her again — what do you do? Then 100 arrives — what changes? Then 102 — and why is 102 special? Or is it special at all?

The citizenship laws

What did the game teach us? Two laws:

1. Every number in the notebook is measured against the reigning queen — the max is the currency, not a gatekeeper. A coup re-denominates everything.
2. You cannot admit a newcomer into $r$ while the past sits in dead currency. coup first demands repair; admission comes second.

In mathematical form — the update run at every arrival $s$:

$$m_{\text{new}} = \max(m_{\text{old}},\, s) \qquad\qquad r_{\text{new}} = e^{\,m_{\text{old}} - m_{\text{new}}}\; r_{\text{old}} \;+\; e^{\,s - m_{\text{new}}}$$

Verify for yourself: when there is NO coup, $m_{\text{old}} = m_{\text{new}}$ and the factor is $e^0 = 1$ — the repair gracefully does nothing, and the law collapses to plain admission. One law, no special cases.

and now open the FlashAttention paper to Algorithm 1, line 11. that's this equation — their $\ell$ is our $r$, their $e^{m - m^{\text{new}}}$ is our repair factor. you did not read that line. You just derived it. this is the entire thesis of this article, cashing out for the first time.

The pour joins the game

But — are we done? What was our final goal anyway? You bet we still have $O$ to think about. Our old friend $O = PV$. The ultimate pour.

Recall the game's rule 1 — and realise it applies to $V$'s rows too. One rule for all: each score arrives holding hands with its ingredient $v$, and the ingredient is present now and gone after. You cannot stockpile the $v$'s and pour at the end — stockpiling is the crime we are escaping. So the pour must happen AS scores arrive, with weights that are not final yet. we refuse to wait. the notebook grows a third resident: $o$, the running blend, poured provisionally and — you guessed it — repaired when regimes change.

So what is the blend's repair factor? Here is the derivation, and it is three lines:

notebox — the two gears of $c$, and a trap. the scalar $c$ has two factors doing two different jobs. The conversion gear $e^{\,m_{\text{old}}-m_{\text{new}}}$ fires only at a coup (idles at $e^0 = 1$ otherwise). The dilution gear $r_{\text{old}}/r_{\text{new}}$ fires at every admission — because every newcomer's deposit grows the room, so every incumbent's share must shrink. The trap every reader falls into (I did): "no coup → no repair." wrong — no coup means no conversion, but the past still dilutes, because the room still grew. Dilution is not punishment; it is justice — each incumbent keeps exactly the share the full-row softmax will eventually grant it.

So the full notebook is the triplet $\langle m, r, o \rangle$ — max, running sum, running blend. Let's run the whole film, with ingredients $v_1 = [1,0]$, $v_2 = [0,1]$, $v_3 = [1,1]$ riding alongside our scores:

And the moment of truth. Compute the same row the ordinary, full-row way — softmax of $[101,100,102]$ gives weights $[0.245,\, 0.090,\, 0.665]$, pour them onto the three $v$'s — and you land on

$$o_{\text{ground truth}} = [\,0.910,\; 0.755\,]$$

to the digit. the notebook never saw the whole row at once, committed early, repaired itself — and arrived exactly where the patient full-row computation arrives. (and one more delight: those weights $[0.245, 0.090, 0.665]$ — go look at row 1 of $P$ back in section 2. Same numbers. Of course: same gaps.)

and Algorithm 1, line 12: the blend update you just derived — the $c \cdot o + \text{new pour}$ — is that line, in the paper written per-row in matrix clothing ($\mathrm{diag}(\ell)^{-1}$ and friends). Both of the algorithm's famously cryptic lines are now things you invented at a bench with three numbers.

REPAIR IS THE PRICE OF REFUSAL TO WAIT.

The scaffold — tiles, rent, and who commutes

In the last section we found the theme: REPAIR IS THE PRICE OF REFUSAL TO WAIT. We illustrated it on a single row, one number at a time. In reality, hardware is batch-o-philic — it always yearns for batches. If you have read my previous article on GEMM vs free-norm the free-normalization article, you are familiar with blockwise matmul: $Q$ and $K$ are giants, so what GPUs actually compute is

$$S_{\text{tile}} = Q_{\text{tile}}\, K_{\text{tile}}^\top$$

A whole tile of scores minted at once — several askers' scores against one section of keys.

Does our notebook survive group arrivals? Absolutely yes — because all said and done, softmax is and will always be a row-wise operation, and a tile is nothing but several rows stacked together. Each asker-row of the tile is one asker's installment of scores, and each such row-strip does its own private bookkeeping the moment its minted: a local queen $\tilde{m}_i$ (the strip's max) and a local sum $\tilde{r}_i$ (the strip's mini-denominator, shifted by its local queen). The strip then presents itself to that asker's notebook as one composite newcomer, and the deposit converts by the same law we already own — currency conversion, second costume:

$$m_{\text{new}} = \max(m_{\text{old}},\, \tilde{m}) \qquad\qquad r_{\text{new}} = e^{\,m_{\text{old}} - m_{\text{new}}}\, r_{\text{old}} \;+\; e^{\,\tilde{m} - m_{\text{new}}}\, \tilde{r}$$

Elegance check: a "block" of one element has $\tilde{m} = s$ and $\tilde{r} = e^0 = 1$, and the formula collapses to exactly last section's law. one rule for all — the game we played was the general case in disguise. And the repair scalar $c$ for the running blends? Unchanged, applied per row. Tiles arrive as rectangles; queens rule rows.

notebox — where is P? careful: what the notebook holds are provisional weights under the current regime — what P would be if the row ended here. The finished P is never minted anywhere, forward pass or otherwise. It exists only implicitly, inside $\langle m, r \rangle$. Remember this; it becomes the hero of the backward-pass story.

The rent arithmetic

Tiles solved the arrival question. Now the size question — how big can a tile be? Not a style choice: the bench decides.

During one visit, the bench must simultaneously hold four block tenants:

$$Q_{\text{tile}}: B_r \times d \qquad K_{\text{tile}}: B_c \times d \qquad V_{\text{tile}}: B_c \times d \qquad O_{\text{tile}}: B_r \times d$$

Start with the square toy — one knob, $B_r = B_c = B$. Total rent $4Bd$, bench capacity $M$:

$$4Bd = M \qquad \Rightarrow \qquad B = \left\lceil \frac{M}{4d} \right\rceil$$

(why ceiling and not floor? Don't read generosity into it — these are sizing heuristics, not physical caps. $M$ itself is an approximate budget, real kernels round block sizes to friendly powers of two, and the last ragged block is handled separately. The ceiling just guarantees a nonzero block.)

Plug in the A100-class numbers: $M = 100{,}000$, $d = 64$ → $B \approx 391$ rows per block. Do the division yourself — its your first rent calculation.

and now the audit that breaks the square plan. we forgot a fifth tenant: the minted score tile itself, $B_r \times B_c$. At $B = 391$: $391 \times 391 \approx 153{,}000$ numbers — bigger than the entire bench, before the four slabs even sit down. Bankrupt.

So something must shrink — and here is the crucial point: nothing gets truncated. We never cut a minted tile; we choose the block shapes before minting. And the side that gets capped is the asker side:

$$B_c = \left\lceil \frac{M}{4d} \right\rceil \approx 391 \qquad\qquad B_r = \min\!\left(\left\lceil \frac{M}{4d} \right\rceil,\; d\right) = 64$$

Only 64 askers shop at a time, while the shelf of keys stays long at 391. Why cap the askers and not the shelf? Two reasons, and the second is the punchline. First: the shelf wants to be long — a long shelf means fewer sections, fewer rounds, and (as the schedule below shows) fewer re-hauls. Second: at $B_r = d$, look at the tile's rent —

$$B_r \times B_c = 64 \times 391 = 391 \times 64 = B_c \times d$$

the tile's rent exactly equals one slab's rent — the same as a K-block. The quadratic tenant, tamed to slab class. Five tenants, one rent tier, everything fits. That's not "smaller" — that's the principled size.

The schedule — who commutes?

Now imagine the machine running. Toy world: $N = 4$ tokens, $d = 2$, blocks of 2 — so two shelf sections (K,V rows 1–2, then rows 3–4) and two asker blocks (Q rows 1–2 with their O rows and notebooks, then rows 3–4).

The paper ships one itinerary (Algorithm 1 — outer loop over K,V blocks, inner over Q blocks): shelf sections take residence one at a time, and the asker blocks parade past each resident section.

And here is the anatomy that decides everything. Watch one asker visit: the block hauls in its Q rows (the questionnaires), hauls in its O rows, works, and — the crucial part — hauls its O rows back OUT. Why does O ride both ways? Because O is wet paint: half-finished blends, mid-repair, that must survive eviction (the next asker block needs the bench) and return next round for more repair. An asker visit costs three trucks: Q in, O in, O out. A shelf section, by contrast, is read-only merchandise — K in, V in, nothing to write back. two trucks.

Tally the toy world's bill — the comic at the end of this section lays this schedule out page by page, with the timetable to count from, if you want to do it yourself — and the O-traffic dominates: the wet paint's round-trips are the single biggest line item.

Now the question the paper didn't ask: what if the roles swapped? let an asker block take residence, and parade the shelf sections past it:

Count again — same comic, second timetable — and the bill drops. Because now the wet paint never travels: a resident asker's blend stays on the bench through every parade, dries, and ships out exactly once, finished. The O round-trips vanish. The verdict, one sentence: the light traveler should commute. the shelf carries two trucks a visit; the askers carry three; FA-1 sat the light party and commuted the heavy one.

And if you just felt clever deriving that — you re-derived history. This exact flip was first implemented by Phil Tillet in Triton and then formalised as FlashAttention-2 (Dao, 2023 — FlashAttention-2; the flip was first implemented by Phil Tillet in Triton). One honesty note though: the traffic saving you just counted is real but constant-factor; FA-2's headline motive was parallelism — with askers outer, different blocks run on different compute units sharing nothing. That story is next-article territory.

So the machine is now fully assembled: tiles minted, notebooks running per row, blends repairing in place — and the full $S$ never exists anywhere. Next section: count what this machine actually hauls, and collect the prize.

The theorem — and where it stops working

Recall from the crime scene the crime scene: standard attention's movement bill was

$$T_{\text{std}} = 4Nd + 4N^2$$

In the last section we built the tiled machine and its schedule. Moment of truth — did all these trials and tribulations actually reduce the bill? Let's measure.

the shelf's haul: every K,V row visits the bench exactly once → $2Nd$. Cheap, irreducible-class.

the askers' commute (the inner loop, in programming terms (confirmed against Algorithm 1 — Dao et al., 2022)): one full parade of Q costs $Nd$. And how many times must the parade run? Once per shelf section — $N/B_c$ rounds. So

$$\text{commute} = \frac{N}{B_c} \times Nd = \frac{N^2 d}{B_c}$$

Now recall the rent: $B_c = M/4d$. Substitute — do it yourself, its one line:

$$\text{commute} = \frac{N^2 d}{M/4d} = \frac{4\,N^2 d^2}{M}$$

Hola!! We got the cost down. Who is the saviour? Of course — its $M$. We could not kill the quadratic term, but we divided it by the bench. A bigger bench now means cheaper postage — a lever that simply did not exist before. And this expression — $\Theta(N^2 d^2 / M)$ — is the paper's Theorem 2 (Theorem 2). you did not read that theorem. You just manufactured it from rent arithmetic and a timetable.

(what about O? It commutes alongside Q — your own visit anatomy from last section, O in and O out — which multiplies the constant, not the shape. Same $N^2d^2/M$ species; the paper's $\Theta$ absorbs both. We derive with Q for cleanliness and bill honestly in the footnote.)

The two ledgers meet

At GPT-2 scale ($N = 1024$, $d = 64$, $M = 100{,}000$), set the bills side by side — and be honest about the ladder:

• counting only flash's commute, against standard's whole bill: ~26× less traffic
• counting flash's shelf haul too, with the honest constants: ~15×
• the paper's measured number on real silicon: 9× (Fig. 2)

26 → 15 → 9: each step toward reality adds overhead — and the ordering itself is the lesson. Our pencil count brackets the measurement from above, exactly as a first-principles count should.

The breakeven — where the bench stops saving you

So is flash attention always a win? Form the ratio of flash's commute against standard's crime term:

$$\frac{4N^2d^2/M}{4N^2} = \frac{d^2}{M}$$

The $N$'s cancel — read that again. The ratio is independent of sequence length. Whether flash wins is a pure property of head size versus bench size, the same at every context length. Flash wins when $d^2/M < 1$, I.e.

$$d < \sqrt{M} \approx 316$$

That's the wall. Now walk the real world along the curve:

• GPT-2 heads, $d = 64$: ratio 0.04 — flash hauls 4% of the crime it replaced
• Llama-class heads, $d = 128$: ratio 0.16 — still a thumping win
• the Gemma family, $d = 256$: ratio 0.66 — flirting with the wall (HF Gemma config; Google — Gemma architecture)
• Gemma 4's global layers, $d = 512$: ratio 2.6 — crossed the wall, and promptly lost flash attention entirely. every FA kernel caps at head_dim 256; those layers fall back to slower attention (flash-attention issue #2427 — head_dim=512 for Gemma 4 global layers; HF: FlashAttention supports head dimension at most 256). The wall we derived with a pencil is roughly where the real ecosystem stops building kernels.
• and the one-fat-head fantasy, $d = 3072$: ratio ≈ 94 — flash would haul ninety-four crimes' worth. Nobody builds this. Now you know why.

Which lands the closing insight, and its my favourite in this article: multi-head attention is flash attention's secret ally. the 2017 design split $d_{\text{model}}$ into many small heads (for expressiveness reasons); the 2022 trick wins precisely and only in the small-$d$ regime. Neither team planned for the other. Deep structure compounds — no prophets required.

The backward pass — resurrection

Thus far, everything we discussed was the forward pass. But what makes neural networks what they are? Correct — the backward pass, the thing that actually trains these models. And the backward pass has a demand that threatens everything we just built.

A quick primer on backprop

Consider a product: $c = a \times b$, with $a = 2$, $b = 3$.

I nudge $a$ a little. How much does $c$ change — how big is $c$'s move relative to $a$'s nudge? Take your time.

So the derivative with respect to one input is the other input's value. Which means: to compute gradients at backward time — which happens later, after the forward pass is long finished — the machine must have kept its inputs. Every framework calls this "saved for backward," and it is exactly why training eats so much more memory than inference.

The ghost at the door

Now scale it up. $O = PV$ is nothing but a grid of these little products, and the same law scales wholesale. The gradient recipe for a matmul — the golden rule, quoted with shapes checked (full derivation lives in my companion notes — it deserves, and will get, its own article):

$$C = AB \;\;\Longrightarrow\;\; dA = dC\, B^\top \qquad\qquad dB = A^\top dC$$

(the transposes are positioned so the shapes close — that is the whole mnemonic)

Apply it: $dV = P^\top\, dO$.

But hold on a second. do you have P with you?

What is P's shape? $N \times N$. So — clear no-go, you can't store it; storing it resurrects the exact quadratic crime this whole article exists to kill. But can you compute $dV$ without P? No way — the recipe demands it. So you want to compute something, you have not stored its ingredients, and you are not permitted to store them.

All you CAN do is compute. Recompute. Recompute... N times — nobody is denying you that. memory is the premium; compute is not. so what do you do?

Correct: you recompute. simple as that. What is P anyway? Softmax of the score tiles. Do you have the raw scores stored? No — but you DO have $Q$ and $K$, pristine in HBM. Re-mint the tile: one cheap matmul, $S_{\text{tile}} = Q_{\text{tile}} K_{\text{tile}}^\top$. And the shift and the denominator? You kept them: the final queen $m_i$ and the final room total $r_i$ — two scalars per row, $2N$ numbers total ($N$ maxes plus $N$ sums). At GPT-2 scale that's 2,048 numbers kept, versus the 1.05 million we refused to store. Peanuts.

So at backward time, any weight resurrects by its own definition:

$$P_{ij} = \frac{e^{\,S_{ij} - m_i}}{r_i}$$

Not an approximation — the definition, evaluated late. The film was never stored; it is re-shot on demand, from actors who never left the building.

And here is the receipt, from our own clinic row. The film ended with $m = 102$, $r = 1.5032$. Pretend you are the backward pass arriving now — re-mint the middle score ($s_2 = 100$) and resurrect its weight:

$$P_2 = \frac{e^{\,100 - 102}}{1.5032} = \frac{0.1353}{1.5032} = 0.090$$

Check it against the film's ground-truth weights $[0.245,\, 0.090,\, 0.665]$ — to the digit. one exponential, one division, and a burned frame walks again.

The trade, named — and the paradox closed

So name what the machine is doing: paying FLOPs to avoid bytes. deliberately redundant arithmetic — every backward tile re-minted — purchased because arithmetic is the cheap currency and HBM traffic is the dear one. And now go back to section 1's receipt, the paradox that opened this article: flash attention does 75.2 GFLOPs where standard does 66.6 — more work — and finishes six times faster. the extra GFLOPs are the recomputation. The missing gigabytes are the never-stored S and P. the paradox was never a paradox; it was a currency exchange, and flash attention simply knows the exchange rate.

Training is memory-intensive, and memory is the premium. So:

RECOMPUTATION IS THE PRICE OF REFUSING TO STORE.

Hold it next to its twin from the forward pass — repair is the price of refusal to wait — and you are holding the entire soul of flash attention: the same trade, struck twice. Refuse to hold the big thing; pay in the cheap currency; keep the mathematics exact to the digit. Two laws, one machine.

Receipts, confessions, and the road ahead

We have come to the end of the derivations. How about we put these ideas to the test and see them live?? That would be pretty cool!!

Let's implement our algorithm in plain python — the $\langle m, r, o \rangle$ loops, literally the film wearing a for-loop. And watch two things happen: it matches standard attention to the digit (our film table is literally the test fixture)... And — honest surprise, I expected it to lose — it already comes out ahead even in numpy, somewhere between 1.0x and 1.6x depending on the run (it is a noisy measurement on a CPU with other things going on; run it a few times and watch it wobble). Why? Because CPU caches are a memory hierarchy too, and tiling is rewarded at every level of the hierarchy: the tiled version never sweeps the N x N ghost through slow RAM. The GPU story is the same physics with far bigger stakes — SRAM control that numpy simply cannot express.

→ flash_attention_numpy.py — run it: exactness PASS at the anchor and GPT-2 scale, and ~10× faster even on CPU (tier-1 numpy implementation — to be added)

So you will ask: it matches, but where is the saving you promised? numpy cannot show you, because numpy cannot see the bench. We need infrastructure that gives us control of the memory hierarchy. ENTER TRITON — a python-embedded language where you decide what lives in SRAM. (fun fact from this article's own story: Triton's creator, Phil Tillet, is the same person who first flipped the flash attention loops — the flip we re-derived in section 6.)

the public kaggle notebook (triton kernel + benchmark — run the speedup urself; notebook file also in code/)

Reading our own receipt

Everything above was measured on a Kaggle T4, $d = 64$, fp16 in and fp32 on the accumulator. Three receipts, in rising order of brutality.

Receipt one — it is still exactly attention. Largest disagreement with the materialise-everything version, across the whole sweep: $3.7 \times 10^{-4}$. That is fp16 rounding, not a different answer. We reordered the arithmetic and changed nothing about the result — which was the whole promise in section 1.

Receipt two — the courier bill, and this one is the article. At $N = 32{,}768$ the naive version needs 10,261 MiB to hold its ghost. Ours needs 25 MiB. That is a 410× difference, and look at the shape of it: naive climbs quadratically, ours crawls up linearly — 10 MiB at $N=1024$, 25 MiB at $N=32{,}768$. This is the $N^2$ versus $N^2d^2/M$ story of section 7, showing up on a real machine, in megabytes you can read off a meter.

Receipt three — the wall, exactly where we said it was. At $N = 65{,}536$ the naive version does not get slower. It dies — the GPU refuses to allocate $S$ at all. Our kernel answers the same question in 882 ms using 49 MiB. Section 4's wall 1 is not a metaphor; it is a hardware error message.

And now the confession, because the clock is the one receipt I did not get

Wall-clock, ours versus torch: 0.25× at $N=1024$, climbing to 0.88× at $N=32{,}768$. Read that honestly — our kernel is still slower, by about 14% at the long end. The line climbs beautifully and never quite crosses.

I could have hidden this by benchmarking against a deliberately clumsy baseline. I would rather tell you what the gap actually is, because the gap is interesting.

Two things are in it, and I measured which is which. The first is an empty GPU. With one head, our grid is $N/B_r$ programs on 40 multiprocessors — at $N=1024$ that is 16 programs on a 40-lane road. Run 12 heads instead, as a real transformer does, and the ratio at $N=1024$ jumps from 0.25× to 0.46×. But do the same at $N = 16{,}384$ and it moves by $-0.008$ — nothing. So the empty grid is a small-$N$ story only, and it is gone by the time the sequences get long.

The second is simply that I am one person and cuBLAS is not. The thing we are racing is not a naive baseline in the compute sense: Q @ K.T is a hand-tuned assembly GEMM and torch.softmax is a fused kernel, both polished for a decade. When I let the autotuner pipeline our loads — num_stages = 4, so the courier fetches the next shelf section while the bench is still working on the current one — our kernel got 2.4× faster without a single line of the algorithm changing. The algorithm was never the slow part. The engineering was, and that engineering is a career, not a section 9.

So: the memory receipts are ours, measured, and they are the argument. The clock belongs to the library engineers, and on production kernels the paper's own Fig. 2 shows what the same idea does when those engineers get hold of it — 41.7 ms down to 7.3 ms. What you derived with a pencil is real. What you have not derived is a decade of kernel tuning, and you should not pretend otherwise.

And one last honesty note, this one about me rather than about the kernel. This is the first GPU kernel I have ever written. So everything in this section is the measurement talking, not experience. I can tell you what the gap is and I can isolate what sits inside it, because those are things you can measure. I cannot yet tell you how to close it — that means tensor-core fragment layouts, shared-memory bank conflicts, instruction scheduling, and a dozen things I would only be paraphrasing at you right now. That is a dedicated sprint, not a subsection at the end of somebody else's article. When I have actually done it with my own hands, it gets its own piece. Until then I would rather hand you a measured 0.88× and name the thing I don't know than perform an expertise I haven't earned.

Honesty time — the confession list

This article walked past several things ON PURPOSE, and you deserve to know what and why:

1. the backward pass, fully — we derived recomputation (the soul), and I quoted the gradient formulas with shapes checked, but the complete derivation (the softmax jacobian and all four gradient flows) lives in my notes. It deserves its own article, and it will get one.

2. the paper's formal proofs — we manufactured theorem 2 with a pencil, but the paper also proves a lower bound: no exact attention algorithm can beat this asymptotically across all SRAM sizes. We counted; they proved. Worth reading in the original.

3. block-sparse flash attention — the paper's own approximate extension (its section 3.3). We stayed exact-only, deliberately.

4. causal masking and dropout — algorithm 2's extra lines. Mechanical once you own the forward pass, so I leave them as an exercise to the reader 😄

5. flash attention 2, 3 and 4 — remember the loop flip we discovered ourselves? That is literally FA-2's opening move. Next stop on this staircase.

6. MQA, GQA and the economics of inference — that is the KV-cache article's territory. And its coming.

7. kernel engineering, the whole discipline — the reason our clock lost. I deliberately did not teach it here, because I have not done it. Worth knowing how deep it goes: the official FlashAttention repo does not support Turing GPUs like the T4 at all — Ampere, Ada and Hopper only (the repo’s own support matrix) — so on the very card this article was measured on, the reference implementation declines to run. And FlashAttention-2 exists because FA-1 reached "only 25-40% of the theoretical maximum FLOPs/s"; FA-2 got to "50-73%", roughly a 2x gain, which its abstract attributes to "suboptimal work partitioning between different thread blocks and warps on the GPU, causing either low-occupancy or unnecessary shared memory reads/writes" (FlashAttention-2, abstract). Low occupancy and shared-memory waste — the same two things our control run found. The man who invented the algorithm needed a second paper to make it fast. That is the honest measure of the distance between deriving this and shipping it.

The staircase

This article stands on the first one — matmuls, couriers, and the roofline the free-normalization article — and borrowed its whole vocabulary. The next flight is already visible: KV cache, then mixture of experts. Same method every time: smallest countable case, pencil first, no borrowed formulas.

This was a great journey personally — highly rewarding in terms of learning. See you at the next one!!