Foundations — Flash attention and efficient attention

Before you can understand why Flash Attention is clever, you have to be fluent in the symbols the parent note throws at you. This page builds every one of them from nothing — plain words, a picture, and the reason the topic needs it. Read top to bottom; each symbol is earned before it is used.

1. A sequence of tokens, and the number $N$

Picture it. Line up the words of a sentence in boxes, left to right. If the sentence is "the cat sat", then $N=3$.

Figure 1 — A sequence: $N$ tokens in order. (Alt text: three boxes labelled "the", "cat", "sat" for tokens 0, 1, 2, with a red double-headed arrow underneath marking the span as "N = 3 tokens".)

Why the topic needs it. Everything expensive in attention grows with $N$. The parent note's whole complaint is that memory grows like $N^2$ — a table with one row and one column per token. So $N$ is the dial that turns a cheap problem into an impossible one.

2. A vector, and the dimension $d$

Picture it. Think of a vector as an arrow, or more simply as a short strip of numbered slots. If $d=2$, each token becomes 2 numbers, e.g. "cat" $\to [0.4,-1.1]$.

Why the topic needs it. The parent's memory math ($2Nd$, $Nd+N^2$) counts how many numbers move between memory chips. To count numbers you need to know both how many tokens ($N$) and how many numbers per token ($d$). When $N\gg d$ (many tokens, few numbers each), the $N^2$ terms win — that is the entire motivation.

3. A matrix and its shape ${\mathbb{R}}^{N\times d}$

Picture it. A spreadsheet: $N$ rows (one per token) and $d$ columns (one per feature). The whole sentence lives in one grid.

Figure 2 — Matrix $X\in {\mathbb{R}}^{N\times d}$: a spreadsheet, one row per token. (Alt text: a grid of numbers with 3 rows and 4 columns; one middle row is outlined in red and labelled "one token's vector (d numbers)"; black double arrows mark "N rows" down the side and "d columns" along the bottom.)

Why the topic needs it. The parent writes $X\in {\mathbb{R}}^{N\times d}$ and $S\in {\mathbb{R}}^{N\times N}$. The shape tells you the memory cost at a glance: an $N\times N$ grid holds $N^2$ numbers. Reading shapes is how you spot the memory explosion.

4. Matrix multiplication and the transpose $K^T$

Picture it. Row $i$ of $Q$ slides across every column of $K^T$; each landing spot gets one number. Row $i$ (length $d_k$) times column $j$ (length $d_k$) collapses to a single score.

Figure 3 — Dot product: a $d_k$-long row and column collapse to ONE score. (Alt text: the row of Q and the column of K^T shown side by side; text "1·1 + 0·0 = 1"; a red arrow points down to a single red-outlined cell labelled "one entry S[i,j]".)

Why we need the transpose (why this tool, not another). $Q$ is $N\times d_k$ and $K$ is $N\times d_k$. You cannot multiply $N\times d_k$ by $N\times d_k$ — the inner sizes clash. Transposing $K$ to $d_k\times N$ makes the inner sizes match ($d_k$ and $d_k$), and the output comes out $N\times N$: exactly one score for every pair of tokens. That $N\times N$ output is the scores matrix $S$ the parent worries about.

5. Query, Key, Value: $Q$, $K$, $V$ and the weights $W_Q,W_K,W_V$

Picture it (library metaphor). Your query is the topic you want. Each book has a key on its spine (its subject). You compare your query to every spine; where they match, you take that book's value (its contents).

Why the topic needs it. $Q{K}^T$ compares queries to keys → scores. Then those scores weight the values $V$ to build the output. Flash Attention is a smarter way to do exactly this comparison-and-weighting, so you must know what each letter is before the algorithm makes sense.

6. The exponential $e^x$

Picture it. A curve hugging zero on the left, passing through $1$ at $x=0$, then shooting up on the right. Two key facts you will use:

• $e^0=1$ (a score equal to the max becomes weight $1$).
• $e^{a+b}=e^a{e}^b$ (exponents split into products — this is the engine of the tiling trick).

Figure 4 — The exponential: always positive, big scores dominate. (Alt text: the curve $y=e^x$ in red on black axes, passing through the point $(0,1)$; annotations "shrinks toward 0" on the left and "shoots up fast" on the right.)

Why this tool and not another. We need a way to turn scores (which can be negative) into positive "weights" that never cancel out, and where bigger scores dominate smoothly. $e^x$ does exactly that: always positive, monotonically increasing. Its split rule $e^{a+b}=e^a{e}^b$ is precisely what lets Flash Attention rescale block sums to a common baseline.

7. Softmax

Picture it. Take the scores, exponentiate each (all become positive, big ones stretch away from small ones), then divide by their total so they sum to $1$. The result: "what fraction of attention goes to each token."

Why the topic needs it. The parent's incremental-softmax derivation is entirely about splitting this ${\sum }_j$ across blocks. If you don't know softmax is "exponentiate, then divide by the sum," none of that derivation is readable.

8. The max $m$ and the safe-softmax shift

Why this tool (why subtract the max). $e^x$ overflows the computer for even moderately large $x$ (e.g. $e^{1000}$ = infinity to a computer). Shifting by the max guarantees the biggest exponent is exactly $1$ and all others are $\le 1$ — no overflow. This is the "safe" in safe softmax, and it is why Flash Attention must carry a running max, not just a running sum.

9. Scaling by $\sqrt{d_k}$

Why this tool. A dot product of two $d_k$-long vectors is a sum of $d_k$ products; the more terms, the larger its typical size grows — roughly proportional to $\sqrt{d_k}$. If scores get huge, softmax becomes razor-sharp (all weight on one token) and gradients vanish. Dividing by $\sqrt{d_k}$ cancels this growth, keeping scores at a stable scale regardless of how big $d_k$ is. That is exactly the parent's note "keeps variance stable as dimensions grow." Note it is specifically $d_k$ (the query/key size), not $d$ or $d_v$, because it is the query–key dot product whose size we are taming.

10. The memory hierarchy: SRAM vs HBM, and bandwidth

Picture it. SRAM is the small desk right in front of you; HBM is the big filing cabinet across the room. Computing is instant once the paper is on your desk — the slow part is walking to the cabinet. Flash Attention keeps small tiles on the desk and avoids trips to the cabinet.

Why the topic needs it. The entire justification ("2–4× faster") comes from cutting round-trips to HBM, not from doing less arithmetic. Without this picture, "memory-bound" is just words.

11. Big-O notation: $O(N)$ vs $O(N^2)$

Picture it. Two curves: a straight-ish line ($O(N)$) and a steep upward bend ($O(N^2)$). At small $N$ they're close; at large $N$ the quadratic dwarfs the linear one.

Why the topic needs it. The headline claim is "$O(N^2)$ memory $\to$ $O(N)$ memory." This notation is the language of that win. See 4.13-Scaling-laws-for-transformers for how these growth curves shape whole-model budgets.

Prerequisite map

The arrows read "is needed before". Follow any path top-to-bottom and every symbol on it is already defined by the time you reach the next.

In words (read this if the diagram below does not render):

• Tokens & $N$ → vectors & $d$ → matrix shape $N\times d$ → matrix product $Q{K}^T$ → the $Q,K,V$ projections → Flash Attention.
• Separately, exponential $e^x$ → softmax → max-shift $m$ / safe softmax → Flash Attention.
• Scaling by $\sqrt{d_k}$, SRAM vs HBM bandwidth, and big-O growth each feed Flash Attention directly.

This page feeds directly into the parent Flash Attention topic. The same $Q,K,V$ machinery underlies 4.1.1-Self-attention-mechanism and 4.1.3-Multi-head-attention; the recomputation trick is a cousin of 5.27-Gradient-checkpointing; and these tokens gain order from 4.1.8-Positional-encoding and power 4.2.1-BERT-architecture and 4.2.3-GPT-architecture.

Equipment checklist

Cover the right side and answer before revealing.