The KV Cache & Arithmetic Intensity
Why a cache exists at all
Autoregressive generation appends one token at a time. Naively, producing token $t+1$ means running the whole prefix $1..t$ through the model again — $O(S^2)$ total work to emit $S$ tokens, most of it recomputation.
But causal attention has a convenient property: the key and value vectors of past tokens never change. Token $j$'s key depends only on token $j$, and the causal mask means it cannot be influenced by anything after it. So compute each $\mathbf k_j,\mathbf v_j$ once and keep them.
Sizing the cache exactly
Per token, per layer, we store one key and one value for each key-value head:
The leading $2$ is for K and V, not for the byte width — that is the last factor. Multiplying by sequence length and batch:
Take an 80-layer model with $d_h=128$ in BF16. Per token per layer, one head costs $2\cdot128\cdot2=512$ bytes.
| Attention variant | $h_{kv}$ | Per token | 8k context, 1 sequence | 128k context, 32 sequences |
|---|---|---|---|---|
| Multi-head (MHA) | 64 | 2560 KiB | 20.0 GiB | 10 TiB |
| Grouped-query (GQA) | 8 | 320 KiB | 2.5 GiB | 1.25 TiB |
| Multi-query (MQA) | 1 | 40 KiB | 0.31 GiB | 160 GiB |
Arithmetic intensity and the roofline
A kernel's speed is capped by whichever of two resources runs out first. Let $W$ be the FLOPs it performs and $Q$ the bytes it moves between memory and the chip. Its arithmetic intensity is
Given peak compute $F_{\text{peak}}$ and peak bandwidth $\beta_{\text{peak}}$, achievable throughput is
The two regimes meet at the ridge point $I^\star=F_{\text{peak}}/\beta_{\text{peak}}$. For an H100 with roughly 989 TFLOP/s of dense BF16 and 3.35 TB/s of HBM,
Decoding sits at intensity equal to the batch size
Now apply it. Generating one token with batch size $B$:
Bytes. Every weight must be read once, regardless of $B$. In BF16 that is $2N$ bytes.
FLOPs. Each of the $B$ sequences performs $2N$ FLOPs, so $2NB$ in total.
Intensity.
The arithmetic intensity of decoding is, to a first approximation, just the batch size.
At $B=1$ the intensity is $1$, roughly 300 times below the ridge. The consequence is a hard latency floor:
about 24 tokens per second for a 70B model on one H100-class device, even with a perfect implementation. No kernel fusion helps; the bytes have to move.
Prefill is the opposite. Processing $S$ prompt tokens at once reads the same weights but does $S$ times the arithmetic, so $I\approx S$. At $S=2048$ prefill is comfortably compute-bound. This is why prompt processing and generation have completely different performance characteristics and are increasingly scheduled as separate phases.
Batching helps the weights, not the cache
Since $I\approx B$, the obvious fix is a bigger batch. It genuinely works for the weight matmuls: one read of $W_Q$ serves all $B$ sequences. But it does not work for attention, and the reason is worth stating carefully.
Weights: shared
Bytes read stay at $2N$ while FLOPs grow as $2NB$. Intensity grows linearly with batch. This is the part batching fixes.
KV cache: private
Each sequence has its own cache, so both bytes and FLOPs grow with $B$. Intensity is stuck at roughly one FLOP per byte no matter how large the batch gets.
Concretely, attending one query against $S$ cached keys reads $2h_{kv}d_hS$ bytes and does about $4h\,d_hS$ FLOPs. The ratio does not contain $S$ or $B$ at all — it is a small constant near $h/h_{kv}$. Attention during decoding is permanently memory-bound, and at long context it becomes the dominant cost.
MQA and GQA
The observation behind both: nothing forces the number of key-value heads to equal the number of query heads.
| Scheme | Key-value heads | Cache versus MHA | Quality |
|---|---|---|---|
| MHA | $h_{kv}=h$ | baseline | baseline |
| MQA | $h_{kv}=1$ | $h\times$ smaller | measurable degradation; can be unstable to train |
| GQA | $1| $h/h_{kv}\times$ smaller | close to MHA | |
Under GQA the $h$ query heads are partitioned into $h_{kv}$ groups, and all heads in a group share one key and one value head. Queries stay fully expressive; only the cached tensors are shared. With $h=64$ and $h_{kv}=8$ that is an eightfold cache reduction for a small quality cost, which is why nearly every recent open-weight model uses it.
The obvious next question is whether the cache can be compressed further without giving up per-head keys and values at all. That is exactly what multi-head latent attention does, and it is the subject of the next note.
Decode intensity is roughly the batch size, the ridge point is roughly 300, and the KV cache is the one tensor batching cannot amortize. Size the cache with $2Lh_{kv}d_hSB$ bytes, and treat every attention variant as an answer to the question "how do we make that number smaller?"
Check yourself
- Compute the KV cache for $L=32$, $h_{kv}=8$, $d_h=128$, $S=32768$, $B=16$ in BF16. calculation
- Derive $I=B$ for decoding and state the assumption that makes it approximate. derivation
- Explain why batching raises intensity for the weight matmuls but not for attention. reasoning
- Compute the minimum per-token latency for a 13B model at 3.35 TB/s and compare it with a 70B model. calculation
- At what context length does the KV cache exceed the model weights for the GQA row of the table above? calculation