Foundations — KV-cache optimization

Before you can understand why we cache K and V (and not Q), you must know what every letter and picture in the parent note actually means. This page builds them in order, each one leaning on the previous. Nothing is assumed.

Related maps: 6.1.13 KV-cache optimization (Hinglish) · parent KV-cache optimization · prerequisites 6.1.1-transformer-architecture, 6.1.4-multi-head-attention.

1. A token and its embedding — the atom of everything

The picture: imagine a row of boxes, one per word: "The cat sat". Each box is replaced by an arrow of numbers pointing in "meaning space". Position $t=1$ is "The", $t=2$ is "cat", and so on.

Why the topic needs it: the whole cost story is "what happens per token". If you don't know that each token becomes a vector $x_t$, you can't count what gets computed for it.

The number of components in one embedding arrow is the model dimension $d$. Big models use $d=768$ (GPT-2 small) up to thousands. That single number $d$ shows up everywhere in the cost formulas, so pin it down now.

2. Q, K, V — three questions asked of every token

Attention (built in 6.1.4-multi-head-attention) turns each token's embedding $x_t$ into three new vectors by multiplying it by three fixed number-grids (matrices) called $W_Q$, $W_K$, $W_V$.

The picture — a library: every past book has a spine label (its Key) and contents (its Value). The new token walks in with a search slip (its Query). It compares its slip against every spine label, then pulls contents from the books that match best.

Why the topic needs it — and this is the heart of caching: a past token's label $K_i$ and contents $V_i$ never change once written. But every new token brings a fresh search slip $Q_t$. That asymmetry is exactly why we cache K and V but recompute Q.

3. The multiply that ranks the books — $Q{K}^T$ and the dot product

To measure how well a search slip matches a spine label, attention uses the dot product: multiply matching components and add them up. A big dot product = strong match.

Pinning down the shapes (this is where novices get lost): at generation step $t$ there is one new query, a row vector $Q_t\in {\mathbb{R}}^{1\times d_k}$ and there are $t$ cached keys stacked into a matrix, one key per row: $K=\left[\begin{array}{c}{K}_1\\ K_2\\ ⋮\\ K_t\end{array}\right]\in {\mathbb{R}}^{t\times d_k}.$

To dot the query against every key at once, we need each key laid out as a column. The transpose $K^T$ does exactly that — it turns the $t\times d_k$ stack of rows into a $d_k\times t$ block of columns: $K^T\in {\mathbb{R}}^{d_k\times t}.$

Now the shapes line up for a matrix multiply: $\underset{1\times d_k}{\underbrace{Q_t}}\underset{d_k\times t}{\underbrace{K^T}}=\underset{1\times t}{\underbrace{\text{scores}}},$ one score per past token. The inner $d_k$ dimensions must match and cancel — that is the whole reason we transpose.

The picture: the query arrow's shadow cast onto each key arrow — the longer the shadow, the higher the score.

In the parent note's single-layer worked example (generating the 3rd token, 4 heads, $d_k=64$), $Q_3(K_1,K_2,K_3{)}^T$ produces exactly 3 scores — one per past position — matching the row of length $t=3$ we derived above.

4. Softmax — turning scores into a choice

The picture — a pie chart of attention: bigger raw score → bigger slice, and all slices sum to one whole pie. The exponential $e^s$ exaggerates the leader so the strongest match dominates but nothing is fully ignored.

Why the topic needs it — and where V's shape enters: these weights multiply the cached Values to blend them into one output. Stack the $t$ Value vectors the same way we stacked Keys: $V=\left[\begin{array}{c}{V}_1\\ V_2\\ ⋮\\ V_t\end{array}\right]\in {\mathbb{R}}^{t\times d_k}.$ The softmax output is a row of $t$ weights, shape $1\times t$. Multiplying shapes: $\underset{1\times t}{\underbrace{\text{softmax scores}}}\underset{t\times d_k}{\underbrace{V}}=\underset{1\times d_k}{\underbrace{\text{output}}}.$ The inner $t$ cancels, so a $1\times t$ weight row times the $t\times d_k$ Value stack yields one length-$d_k$ output vector — a weighted blend of all the Values. So the full attention formula

reads left-to-right as: score the labels → turn into a pie → mix the contents.

5. Heads, $h$, and $d_k$ — splitting the work

One attention isn't enough, so the model runs several in parallel — each is a head.

What counts as a "layer"? A transformer stacks identical blocks on top of each other; each block is a layer (6.1.1-transformer-architecture builds these). One layer contains its own attention (with its own $W_Q,W_K,W_V$) plus a feed-forward network. GPT-2 small has $12$ such layers; a GPT-3-style model has $96$. Let $N$ = number of layers. Every layer keeps its own separate KV-cache, which is why full-model memory multiplies by $N$.

Why the topic needs it: cache memory is counted per head, per layer, per token. Now that $h$, $d_k$ and $p$ are all defined, the memory formula reads:

The $2$ is because we store two things: the Key stack and the Value stack.

6. Multi-Query Attention — where the $h$-fold saving comes from

The memory derivation, made explicit:

With $h=12$ heads, MQA stores $12\times$ less KV-cache, because K and V are each stored once (at length $d_k$) instead of $h$ times. Quality typically drops only 1–2%.

7. Autoregressive generation and the causal mask

The picture — a lower-triangular grid: row $i$ (the attending token) can fill in columns $j\le i$ (itself and the past) but the upper-right triangle (the future) is blacked out.

Figure 5 — an $L\times L$ grid indexed by attending position $i$ (rows) and key position $j$ (columns); green cells $j\le i$ hold $0$ (allowed), red upper-right cells $j>i$ hold $-\infty$ (future blocked).

Why the topic needs it: because the mask guarantees a token only ever attends to positions $\le t$, and those past K,V are frozen — that is the exact license to cache them. Even with a cache you must still add the mask, or the model would illegally read future positions.

8. Big-O, $L$, and the cost we're fighting

Why the topic needs it: the entire motivation is a cost comparison. The projection cost per token is $O(d^2)$ per layer (a $d_k\times d$ matrix times a length-$d$ vector, summed over $h$ heads gives $h{d}_{k}d=d\cdot d=d^2$). Big-O hides the constant factors $N$ (layers) and the per-head split, so the clean formulas below are per-layer projection costs; multiply by $N$ layers for the full model:

For a 100-token sequence, that's roughly $100\times$ less projection work. The parent note's GPT-2-small worked example (which does include the $N=12$ layer factor on both sides, so it cancels) measures a $\approx 50\times$ real speedup for $L=100$ — a bit under the ideal $L$ because attention itself still costs a little each step, but the same order of magnitude.

The two classic mistakes (why they're wrong, in one line each)

Prerequisite map

Each foundation feeds the next; together they justify why caching K and V (not Q) is safe and cheap. Deeper follow-ups: 6.2.3-inference-optimization, 6.1.11-sparse-attention, 6.3.1-model-quantization, 7.2.1-llm-deployment, 5.4.2-beam-search.

Equipment checklist

Test yourself — cover the right side.