Tokens Become Directions in Space
Before attention can do anything, meaning has to become geometry.
The one idea the whole series rests on
A transformer never sees words. It sees vectors — long lists of numbers, one per token, living in a space of a few thousand dimensions. The single most important fact about that space is that direction encodes meaning: tokens that behave alike point alike. Everything attention does is manipulate these directions.
We can't draw 4096 dimensions, so throughout this course we project down to 3 and animate the shadow. The intuition survives the projection.
What "an embedding" actually is
An embedding is a learned lookup table \(E \in \mathbb{R}^{V \times d}\): one row per vocabulary entry \(V\), each row a \(d\)-dimensional vector. Token id \(i\) becomes row \(E_i\). The rows start random and are shaped by gradient descent until the geometry above emerges — no one hand-places "king" near "queen"; the training objective does.
That table \(E\) can be enormous — a 128k-token vocab at \(d=8192\) is a billion parameters — but the idea is simple: every token is a point in a high-dimensional space, and which way it points carries its meaning. Attention is entirely about comparing those directions. Hold that thought.
Position has to be injected
A bare embedding is order-blind: "dog bites man" and "man bites dog" produce the same set of vectors. Transformers add positional information — classically sinusoidal encodings, in modern models rotary embeddings (RoPE) that rotate Q and K by an angle proportional to position. The upshot: each token's vector now carries both what it means and where it sits.
With tokens as positioned directions, we're ready for the question attention answers: which other tokens should this one listen to?
Every Token Asks, Advertises, and Offers
Q, K, V are three learned projections of the same vector — a query, a key, a value.
Three roles from one vector
Take a single token's vector \(x\). Attention multiplies it by three learned weight matrices to produce three new vectors:
$$ q = W_Q\,x \qquad k = W_K\,x \qquad v = W_V\,x $$
Read them as roles:
Query — what am I looking for?
The direction this token uses to probe every other token.
Key — what do I advertise?
The direction a token exposes so others can decide if it's relevant.
Value — what do I contribute?
The information actually passed on if a token is attended to.
It's three matmuls, batched over the sequence
In practice you don't do this one token at a time. Stack all \(T\) tokens into \(X \in \mathbb{R}^{T \times d}\) and the projections become three dense matrix multiplies:
$$ Q = X W_Q,\quad K = X W_K,\quad V = X W_V \qquad W_\bullet \in \mathbb{R}^{d \times d_{\text{head}}} $$
Notice the shapes: \(Q\) and \(K\) share dimension \(d_{\text{head}}\) so we can dot them together. And here's a foreshadow — the whole MHA→MQA→GQA→MLA story in Module 5 is about how many copies of \(K\) and \(V\) you actually need to keep around. For now, on to the dot product itself.
A token's Q, K, and V vectors are…
One Attention Head, End to End
Dot-product to score, softmax to weight, weighted sum to mix. This is the whole engine.
Score by alignment
A query \(q\) scores a key \(k\) by their dot product \(q \cdot k\). Geometrically that's alignment: keys pointing the same way as the query score high, orthogonal keys score near zero, opposed keys go negative. The scene below draws one query (violet) against three keys (blue) — watch which key aligns.
The formula, one symbol at a time
$$ \text{Attention}(Q,K,V) = \underbrace{\text{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d_{\text{head}}}}\right)}_{\text{weights, one row per query}} V $$
- \(QK^{\top}\) — every query dotted with every key: a \(T \times T\) score matrix.
- \(/\sqrt{d_{\text{head}}}\) — the scaling that keeps scores from exploding as \(d_{\text{head}}\) grows, so softmax doesn't saturate. This is the "scaled" in scaled dot-product attention.
- softmax — turns each score row into a probability distribution: non-negative, sums to 1.
- \(\cdot V\) — each output is that distribution's weighted average of value vectors.
Why the \(T \times T\) matrix is the villain of this whole course
That score matrix is quadratic in sequence length. At \(T=128\text{k}\) it's 16 billion entries per head — it cannot live in memory naively. This single fact drives every variant we meet from Module 5 on: shrinking the KV cache (MQA/GQA/MLA), skipping cells (local/sparse), reordering the matmul (linear attention), and never materializing the grid at all (FlashAttention). We build the grid explicitly in the next module, then spend the rest of the course beating it.
In a decoder, a token may only attend to itself and earlier tokens. We mask the upper triangle of \(QK^{\top}\) to \(-\infty\) before softmax so future positions get weight zero. Masking also shapes how work splits across devices.
Many Heads — and the T×T Matrix That Costs You
Multi-head is cheap. The quadratic score matrix is the bill every attention variant is trying to shrink.
Parallel heads, each with its own subspace
One head can only track one kind of relationship. So a block runs \(h\) heads in parallel, each with its own \(W_Q, W_K, W_V\) projecting into a smaller \(d_{\text{head}} = d/h\) subspace. One head might track subject–verb agreement, another coreference, another local syntax — learned, not assigned.
$$ \text{MHA}(X) = \text{Concat}(\text{head}_1,\dots,\text{head}_h)\,W_O,\qquad \text{head}_i = \text{Attention}(XW_Q^i, XW_K^i, XW_V^i) $$
Splitting \(d\) into \(h\) heads of size \(d/h\) is no more expensive than one big head: the FLOPs are \(h \cdot n^2 \cdot (d/h) = n^2 d\) either way. Multi-head buys you diversity for free.
The real cost is that grid in the middle
Look at one head's core: \(\text{softmax}(QK^{\top}/\sqrt{d})\,V\). The \(QK^{\top}\) term is an \(n \times n\) matrix — one score for every (query, key) pair. That grid is the whole story of attention's cost.
The grid has \(T^2\) cells. Double the context, quadruple the score matrix — in both compute and memory. At \(T = 16{,}384\) one head's score matrix is 512 MB in FP16. That wall is why the rest of this course exists.
Four ways to fight the T×T bill
Every attention variant is an answer to "the \(T \times T\) matrix is too expensive." They attack it from different directions:
Shrink the KV cache
MHA → MQA → GQA → MLA: keep the full grid but store far fewer keys/values (Module 5).
Skip cells
Local & sparse attention: only compute a band or a pattern of the grid (Module 6).
Reorder the matmul
Linear attention: never form the grid at all — use \(Q(K^{\top}V)\) (Module 7).
Never store it
FlashAttention: compute the exact grid in tiles, never write it to memory (Module 8).
Shrinking the KV Cache: MHA → MQA → GQA → MLA
One story, four steps — how much of the keys and values do you actually need to keep?
Why keys and values pile up
When a model generates text, it produces one token at a time. To attend, each new token needs the keys and values of every previous token. Recomputing them each step would be wasteful, so they're stored — the KV cache. Its size is the memory bottleneck of inference: it grows with sequence length and with the number of KV heads. The variants below are one long effort to make it smaller.
The four steps
MHA — Multi-Head
Every query head has its own key and value head. Maximum expressiveness, maximum KV cache. The baseline (Vaswani 2017).
MQA — Multi-Query
All query heads share a single KV head. KV cache shrinks ~\(h\times\); decode gets much faster — but quality can drop.
GQA — Grouped-Query
An intermediate number of KV groups (more than 1, fewer than \(h\)). Uptrained GQA gets near-MHA quality at MQA-like speed. Used in Llama-2/3, Mistral.
MLA — Multi-head Latent
Compress K and V into one shared latent vector, decompressed on the fly. The smallest KV footprint (DeepSeek-V2).
Query heads never change — you always have \(h\) of them. What shrinks is how many distinct \(K\)/\(V\) you store: \(h\) → groups → 1 → a latent. Same attention math, progressively less memory to carry between decode steps.
Grouped-Query Attention (GQA) uses…
Local & Sparse Attention: Skip Most of the Grid
If a token only needs a few neighbors, don't compute the whole T×T matrix.
Most attention weight is local
In practice, a token's attention is dominated by nearby tokens. Full attention still pays for the whole \(T \times T\) grid anyway. Local and sparse attention exploit the observation: only compute the cells you expect to matter, and treat the rest as zero.
The two workhorse patterns
Sliding window (local)
Each query attends only to the last \(w\) keys — a diagonal band. Cost drops from \(O(T^2)\) to \(O(T\cdot w)\), i.e. linear in sequence length. Stacking layers still grows the effective receptive field. (Mistral, Longformer.)
Strided / global sparse
Combine a local window with a handful of global tokens every query can see (and that see everyone). Keeps long-range reach at a fraction of the cells. (Sparse Transformer, BigBird.)
Full causal attention computes \(\tfrac{T(T+1)}{2}\) cells. A window of \(w\) computes about \(T\cdot w\). At \(T=8192, w=256\) that's ~2M cells instead of ~34M — a 16× cut, and it grows linearly, not quadratically, as context lengthens.
Sparsity is approximate — you're betting the skipped cells didn't matter. That bet usually holds, but a genuinely long-range dependency outside the window is lost unless a global/strided path carries it. The pattern is a modeling choice, not a free lunch.
Linear Attention: Reorder the Matmul, Kill the Grid
The T×T matrix is optional. Associativity lets you never build it.
A trick hiding in the parentheses
Ignore softmax for a moment and look at the shape of attention: it's \((QK^{\top})V\). Matrix multiplication is associative — you get the same answer computing it as \(Q(K^{\top}V)\). But the two orders cost wildly different amounts:
Where softmax goes
The catch: real attention has \(\text{softmax}(QK^{\top})\) in the middle, and softmax does not factor across the reorder. Linear attention replaces the softmax similarity \(\exp(q\cdot k)\) with a kernel feature map \(\phi(q)\cdot\phi(k)\) that does factor:
$$ \text{out}_i = \frac{\phi(q_i)\left(\sum_{j\le i}\phi(k_j)^{\top} v_j\right)}{\phi(q_i)\left(\sum_{j\le i}\phi(k_j)\right)} $$
The inner sums are a running state of fixed size \(d\times d\) — update it token by token, no \(T\times T\) matrix, \(O(T)\) time and \(O(1)\)-in-\(T\) memory. This is the bridge to recurrent/state-space views (Performer, Linear Transformers, and the RNN-like reading of attention).
The feature map \(\phi\) only approximates softmax's sharp, selective weighting. You buy linear cost and a compact recurrent state; you pay in some of softmax's ability to focus hard on a single key. Great for very long sequences, not always a drop-in for quality-critical short-context work.
FlashAttention: Compute the Exact Grid, Never Store It
Same math as full attention — but tiled through fast on-chip memory so the T×T matrix is never written out.
Attention is memory-bound, not compute-bound
The surprise in the naive implementation: the GPU spends most of its time waiting on memory, not multiplying. It writes the giant \(T\times T\) score matrix out to slow HBM, reads it back for softmax, writes again, reads again. At \(T=16{,}384\) that intermediate is 512 MB per head, per batch item in FP16 (\(16384^2 \times 2\) bytes). Moving it dominates the runtime.
The obstacle: softmax isn't associative
Matmul tiles cleanly because addition is associative — you can accumulate partial sums in any order. Softmax can't, because it needs the max and the sum over the whole row for numerical stability. FlashAttention's key trick is online softmax: carry two running numbers per row and correct them as each new block arrives.
$$ m_i = \max(m_{i-1}, \tilde{m}),\qquad \ell_i = \ell_{i-1}\,e^{\,m_{i-1}-m_i} + \textstyle\sum e^{\,x - m_i} $$
The \(e^{m_{i-1}-m_i}\) factor rescales the old running sum whenever a bigger max appears. The result is the exact softmax — no approximation — computed without ever holding a full row, let alone the full matrix.
Why it scales
FlashAttention tiles \(Q, K, V\) into blocks that fit in on-chip SRAM and fuses the whole computation into one kernel: scores, softmax, and the value-weighted sum happen on-chip, and only the final output goes back to HBM. The SRAM footprint depends on block size and head dimension — not on sequence length — so it scales to 16k+ context without the memory blow-up. Result: 2–4× faster and dramatically less memory, enabling longer context.
It does not make attention sub-quadratic in compute — it still performs \(O(T^2)\) FLOPs. What it removes is the \(O(T^2)\) memory traffic and storage. It's an exact, IO-aware reorganization, not an approximation like sparse or linear attention. That's why it composes with them: FlashAttention for the kernel, sparsity/linearity for the algorithm.
The family, in one table
| Variant | Attacks | Cost vs full | Exact? |
|---|---|---|---|
| MHA | — | O(T²), full KV | exact (baseline) |
| MQA / GQA / MLA | KV-cache size | same compute, far less KV memory | exact |
| Local / sliding window | which cells compute | O(T·w), linear | approximate |
| Sparse / strided | which cells compute | sub-quadratic | approximate |
| Linear attention | the matmul order | O(T) | approximate |
| FlashAttention | memory traffic | O(T²) compute, O(1)-in-T memory | exact |
Every row of that table is a different answer to the same question from Module 4: the \(T\times T\) score matrix is too expensive. Shrink what you store, skip what you compute, reorder to avoid the grid, or compute it exactly but never write it down. Understand the grid, and the whole family falls into place.