KV-cache optimization: Exercises — AI-ML deep dive

Definition Symbols and unit conventions used on this page

Before any problem, let us pin down every symbol and every unit so nothing is ever used undefined.

$L$ — sequence length (number of tokens).
$h$ — number of attention heads.
$d_{k}$ — head dimension: the length of one head's key/value/query vector.
$d$ — model width (a.k.a. hidden size): the full embedding dimension fed into the projection matrices. It relates to the heads by $d = h \times d_{k}$ . Whenever you see $d^{2}$ it comes from a projection matrix $W \in R^{d \times d}$ (see Ex 2.3 / 4.1 for why).
precision — bytes per stored number (fp16 $= 2$ bytes).
Unit convention (used everywhere below): we use binary units. $1 KB = 2^{10} = 1024$ bytes, $1 MB = 2^{20}$ bytes, $1 GB = 2^{30}$ bytes. So "16 KB" here means exactly $16 \times 1024 = 16, 384$ bytes. (Strictly these are KiB/MiB/GiB; we write KB/MB/GB for readability but the arithmetic is always powers of two.)

Before we start, one shared picture of the two worlds — "recompute everything" vs "remember and append":

Reading the figure. The top row is generation without a cache at step 4: all four key boxes $K_{1}, K_{2}, K_{3}, K_{4}$ are drawn in black because the model recomputes every one of them from scratch. The bottom row is generation with a cache: $K_{1}, K_{2}, K_{3}$ are black (already stored, just reused — follow the "reuse" arrow) and only $K_{4}$ is drawn in red, because it is the single new projection the cache must compute this step. That contrast — three reused vs one new — is the whole idea of the cache, and every memory/FLOP problem below is just counting these boxes.

Level 1 — Recognition

Exercise 1.1 (L1)

State which of the following are stored in the KV-cache during autoregressive generation, and which are recomputed every step: (a) key projections $K_{i}$ , (b) value projections $V_{i}$ , (c) query projections $Q_{t}$ , (d) the final softmax attention weights.

Recall Solution

(a) $K_{i}$ — cached. A past token's key never changes once computed, so we store it (the black boxes in Figure s01).
(b) $V_{i}$ — cached, same reason.
(c) $Q_{t}$ — recomputed each step. The query is "what the current token is looking for," and there is a brand-new current token every step, so there is nothing to reuse (the red box is new every step).
(d) softmax weights — recomputed. They depend on the new $Q_{t}$ , so they change every step.

The cache stores exactly the things that are fixed for past positions: $K$ and $V$ . Hence the name KV-cache.

Exercise 1.2 (L1)

Recognition of shape. A single layer has $h = 8$ heads and head dimension $d_{k} = 64$ . After generating 5 tokens, what is the shape of $Cache_{K}$ (as tokens × heads × head-dim)?

Recall Solution

Each token contributes one key per head: shape per token is $(h, d_{k}) = (8, 64)$ . Stacking 5 tokens gives $Cache_{K} \in R^{5 \times 8 \times 64} .$ $Cache_{V}$ has the identical shape. The first axis (here $5$ ) is the only one that grows as we generate.

Level 2 — Application

Exercise 2.1 (L2)

Per-token cache memory. A layer has $h = 32$ heads, $d_{k} = 128$ , stored in fp16 (2 bytes each). How many bytes does one token add to the KV-cache for that one layer?

Recall Solution

Where the formula comes from (short derivation). For one token we must store its keys and its values. A single token's key set is one $d_{k}$ -vector per head, so it has $h \times d_{k}$ numbers; the same holds for the value set. That is two tensors of $h \times d_{k}$ numbers, and each number occupies "precision" bytes. Multiply: $bytes/token/layer = K and V 2 \times numbers per tensor h \times d_{k} \times bytes/number precision .$ That is the whole reason for each factor — nothing is cited blindly. Plug in: $= 2 \times 32 \times 128 \times 2 = 16, 384 bytes = 16 KB .$ (Recall the unit convention: $16 KB = 16 \times 1024 = 16, 384$ bytes exactly.)

Exercise 2.2 (L2)

Full-model cache for a whole sequence. Same layer specs ( $h = 32$ , $d_{k} = 128$ , fp16), now $n_{layers} = 40$ and a sequence of $L = 1024$ tokens. Total cache size in gigabytes?

Recall Solution

One token in one layer costs the $2 h d_{k} \times precision$ from Ex 2.1. We now have $L$ tokens (first cache axis grows) and $n_{layers}$ independent layers (each keeps its own cache), so we multiply by both: $Memory = 2 \times n_{layers} \times L \times h \times d_{k} \times precision .$ Plug in: $= 2 \times 40 \times 1024 \times 32 \times 128 \times 2 bytes .$ Step it: $2 \times 40 = 80$ ; $\times 1024 = 81, 920$ ; $\times 32 = 2, 621, 440$ ; $\times 128 = 335, 544, 320$ ; $\times 2 = 671, 088, 640$ bytes. Convert (binary GB): $671, 088, 640/ 2^{30} = 0.625$ GB. Answer: $0.625$ GB ( $= 640$ MB, since $0.625 \times 1024 = 640$ ).

Exercise 2.3 (L2)

Speedup factor. Recall $d = h d_{k}$ is the model width, and each of $W_{Q}, W_{K}, W_{V} \in R^{d \times d}$ projects a token's $d$ -vector into a $d$ -vector, costing $O (d^{2})$ multiply-adds per token. Without caching, at step $t$ the model re-projects all $t$ tokens, and summing over the whole generation makes the projection work scale as $O (L^{2} d^{2})$ ; with caching it re-projects only the one new token per step, giving $O (L d^{2})$ . First derive why the no-cache total is $O (L^{2} d^{2})$ , then compute the projection speedup for $L = 256$ .

Recall Solution

Why the no-cache cost is $O (L^{2} d^{2})$ (the missing "why"). One token's three projections cost $O (d^{2})$ each — because multiplying a $d$ -vector by a $d \times d$ matrix is $d^{2}$ multiply-adds — so $O (d^{2})$ per token. Without a cache, at generation step $t$ the model treats the sequence as fresh and re-projects all $t$ tokens seen so far, costing $O (t d^{2})$ that step. Summing over every step $t = 1 \dots L$ : $\sum_{t = 1}^{L} O (t d^{2}) = O (d^{2}) \sum_{t = 1}^{L} t = O (d^{2} \cdot \frac{L ( L + 1 )}{2}) = O (L^{2} d^{2}) .$ The quadratic in $L$ comes purely from the triangular sum $\sum t = \frac{L ( L + 1 )}{2}$ . With a cache, only the one new token is projected per step: $\sum_{t = 1}^{L} O (d^{2}) = O (L d^{2})$ . Speedup: $\frac{L ^{2} d ^{2}}{L d ^{2}} = L = 256.$ The $d^{2}$ cancels — the speedup is purely the sequence-length factor $L$ . Longer generations benefit more.

Level 3 — Analysis

Exercise 3.1 (L3)

MQA vs standard. A model uses $h = 32$ query heads, $d_{k} = 128$ . Compare the KV-cache size of standard multi-head attention against Multi-Query Attention (MQA), which shares a single K/V head. Give the ratio and interpret.

Recall Solution

Standard cache (per token, per layer, ignoring the constant factor): $2 h d_{k} = 2 \times 32 \times 128.$ MQA keeps one K and one V head: $2 \times 1 \times d_{k} = 2 \times 128.$ Ratio: $\frac{MQA}{standard} = \frac{2 d _{k}}{2 h d _{k}} = \frac{1}{h} = \frac{1}{32} .$ MQA uses $32 \times$ less KV-cache. Interpretation: query heads still specialize (each has its own $W_{Q}$ ), but they all read from the same keys and values, so we only store one set. Cost: a small quality drop (~1–2% perplexity).

Exercise 3.2 (L3)

GQA middle ground. Llama-2 70B uses $h = 64$ query heads grouped into $g = 8$ K/V groups. (a) What is the cache reduction versus standard attention? (b) Where does GQA sit between standard and MQA?

Reading the figure. Three head-sharing schemes are drawn side by side. The top boxes are the query heads; the bottom boxes are the stored KV heads; each line shows which queries read which KV. Left ("standard $g = h$ "): every query has its own private KV — no sharing, maximal cache. Right ("MQA $g = 1$ "): all queries funnel into one shared KV — minimal cache. Middle ("GQA $g = 2$ ", drawn in red): queries are bundled into groups, each group sharing one KV — the tunable middle ground. As you slide $g$ from $h$ down to $1$ , you slide the cache from "standard" down to "MQA"; the red middle is exactly the case part (b) asks about.

Recall Solution

(a) GQA stores $g$ K/V heads instead of $h$ : $\frac{GQA}{standard} = \frac{2 g d _{k}}{2 h d _{k}} = \frac{g}{h} = \frac{8}{64} = \frac{1}{8} .$ So 8× reduction. (b) Read the figure left-to-right: standard has one K/V head per query head ( $g = h$ , no reduction); MQA collapses to a single shared K/V head ( $g = 1$ , maximal $h \times$ reduction); GQA (the red grouping in the middle) sits between with $1 < g < h$ . It recovers most of MQA's memory savings while keeping more of standard attention's quality by giving each group of queries its own keys/values.

Exercise 3.3 (L3)

Sliding-window memory. For a 100K-token context, compare full-cache memory to a sliding-window cache of width $W = 4096$ . Assume $n_{layers} = 32$ , $h = 32$ , $d_{k} = 128$ , fp16. Give both sizes and the ratio.

Recall Solution

Per-token-per-layer bytes: $2 \times 32 \times 128 \times 2 = 16, 384$ B $= 16$ KB. Full cache ( $L = 100, 000$ ): $32 \times 100, 000 \times 16 KB = 51, 200, 000 KB = 48.83 GB .$ (Binary conversion: $51, 200, 000/ 2^{20} = 48.83$ GB.) Sliding-window ( $W = 4096$ ): $32 \times 4096 \times 16 KB = 2, 097, 152 KB = 2.0 GB .$ Ratio: $100000/4096 \approx 24.4 \times$ smaller. Trade-off: the window keeps memory constant at $O (W)$ regardless of context length, but the model literally cannot attend to anything older than $W$ tokens — long-range dependencies past the window are lost.

Level 4 — Synthesis

Exercise 4.1 (L4)

FLOP accounting, cache vs no-cache. For GPT-2 small ( $n_{layers} = 12$ , model width $d = 768$ ), estimate projection FLOPs to generate $L = 100$ tokens with and without caching. Report both and the speedup.

Recall Solution

Where the per-step cost comes from. As established in Ex 2.3, projecting one token through a $d \times d$ matrix costs $\approx d^{2}$ multiply-adds; we treat the three projections $W_{Q}, W_{K}, W_{V}$ as $O (d^{2})$ and drop the small constant factor (the "3" and any $+ d$ terms) because we only want the scaling with $L$ and the ratio — constants cancel in the ratio anyway. So:

No-cache, step $t$ , per layer: re-project all $t$ tokens $\Rightarrow\approx t d^{2}$ .
Cache, step $t$ , per layer: project only the 1 new token $\Rightarrow\approx d^{2}$ .

Without cache — sum the triangular series over all steps, times layers: $12 \times (\sum_{t = 1}^{100} t) \times 76 8^{2} = 12 \times 5050 \times 589, 824 \approx 3.57 \times 1 0^{10} FLOPs .$ (Used $\sum_{t = 1}^{100} t = \frac{100 \cdot 101}{2} = 5050$ — the same triangular sum as Ex 2.3.) With cache — one new K/V/Q per step, so $d^{2}$ per step: $12 \times 100 \times 76 8^{2} = 12 \times 100 \times 589, 824 \approx 7.08 \times 1 0^{8} FLOPs .$ Speedup: $\frac{3.57 \times 1 0 ^{10}}{7.08 \times 1 0 ^{8}} \approx 50.5 \times .$ This matches the expected $\sim \frac{L + 1}{2}$ factor: the sum-of-integers gives $\frac{L ( L + 1 )}{2}$ vs $L$ , i.e. $\frac{L + 1}{2} = 50.5$ . The dropped constant factor never mattered because it appears identically in both numerator and denominator.

Exercise 4.2 (L4)

Design under a memory budget. You have a 24 GB GPU. Model weights take 14 GB, leaving 10 GB for KV-cache. Layer specs: $n_{layers} = 32$ , $h = 32$ , $d_{k} = 128$ , fp16. (a) What is the maximum context length $L$ you can cache with standard attention? (b) With GQA at $g = 8$ ?

Recall Solution

Per-token, all-layers, standard: $2 \times 32 \times 32 \times 128 \times 2 bytes = 524, 288 B = 512 KB/token .$ Budget $= 10$ GB $= 10 \times 2^{30} = 10, 737, 418, 240$ B (binary GB). (a) $L_{max} = \frac{10 , 737 , 418 , 240}{524 , 288} = 20, 480$ tokens. (b) GQA with $g = 8$ replaces the head count $32$ by $8$ in K/V → $\frac{8}{32} = \frac{1}{4}$ the per-token size $= 128$ KB/token: $L_{max} = 20, 480 \times 4 = 81, 920 tokens .$ Design conclusion: switching to GQA( $g = 8$ ) quadruples the affordable context length under the same memory budget.

Level 5 — Mastery

Exercise 5.1 (L5)

Crossover analysis. MQA saves memory but you worry about quality. Suppose standard attention gives perplexity $p_{0}$ and MQA gives $p_{0} (1 + 0.015)$ (a $1.5%$ increase). Meanwhile MQA lets you cache $h = 32 \times$ more tokens. If longer context reduces perplexity by roughly $2%$ when context grows $4 \times$ (empirical), does trading standard→MQA to buy $4 \times$ context help or hurt perplexity net?

Recall Solution

Two effects compose multiplicatively on perplexity:

MQA head-sharing penalty: factor $1.015$ (worse).
$4 \times$ longer usable context: factor $0.98$ (better). Net factor: $1.015 \times 0.98 = 0.9947.$ Since $0.9947 < 1$ , net perplexity drops by ~0.53% — the context gain outweighs the sharing penalty. Mastery point: memory optimizations are not "free lunch vs pure cost"; they enable longer context, which itself improves quality. You must weigh the direct penalty against the capability it unlocks.

Exercise 5.2 (L5)

Full stack estimate. Combine everything: $n_{layers} = 80$ , $h = 64$ , $d_{k} = 128$ , fp16, and you deploy GQA( $g = 8$ ) with a sliding window $W = 8192$ for a nominal 200K-token context. (a) What would the full standard cache have cost at 200K tokens? (b) What does the GQA + sliding-window cache actually cost? (c) Combined reduction factor?

Recall Solution

Standard per-token-all-layers: $2 \times 80 \times 64 \times 128 \times 2 = 2, 621, 440 B = 2.5 MB/token .$ (Binary: $2, 621, 440/ 2^{20} = 2.5$ MB.) (a) Full standard at $L = 200, 000$ : $200, 000 \times 2.5 MB = 500, 000 MB = 488.28 GB .$ (b) GQA cuts K/V heads $64 \to 8$ , factor $8/64 = 1/8$ → per-token $= 2.5 MB /8 = 320$ KB. Sliding window caps tokens at $W = 8192$ : $8192 \times 320 KB = 2, 621, 440 KB = 2.5 GB .$ (c) Reduction: $\frac{488.28}{2.5} \approx 195 \times .$ Interpretation: two orthogonal levers multiply — GQA gives $8 \times$ (head sharing) and the window gives $200000/8192 \approx 24.4 \times$ (bounded length), and $8 \times 24.4 \approx 195$ . This is why production long-context serving (7.2.1-llm-deployment, 6.2.3-inference-optimization) stacks multiple cache tricks rather than relying on one.

Recall Quick self-check summary

Cache holds K and V only ::: queries are fresh each step Memory formula leading constant ::: 2, for K and V Model width in terms of heads ::: $d = h \times d_{k}$ Standard→MQA cache ratio ::: $1/ h$ GQA( $g$ ) cache ratio ::: $g / h$ Projection speedup from caching ::: factor $L$ Sliding window memory order ::: $O (W)$ , constant in $L$ Combining independent cache savings ::: multiply the factors

Related: 6.1.11-sparse-attention, 6.3.1-model-quantization, 5.4.2-beam-search, 6.1.13 KV-cache optimization (Hinglish).