← All Posts
Deep Learning · Popular Videos · Umar Jamil· Part 4 · Chapter 6

The KV Cache & Arithmetic Intensity

Generating a token from a 70B model moves 140 GB of weights to do 140 GFLOPs of arithmetic. That is one FLOP per byte, against hardware built for roughly 295. Decoding is not slow because the model is big; it is slow because almost nothing is computed per byte read. Every optimization in this note follows from that ratio.

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.

The KV cache stores $\mathbf k$ and $\mathbf v$ for every past position, in every layer. With it, a decoding step projects only the new token into $\mathbf q,\mathbf k,\mathbf v$, appends the new $\mathbf k,\mathbf v$ to the cache, and attends the single new query against the whole cache. Per-token work drops from $O(S)$ layer passes to $O(1)$.
Prefillwhole prompt at once, fills the cache, big matmuls
→
Decodeone token at a time, reads the cache, matrix-vector
→
Cache growsone entry per layer per step, forever

Sizing the cache exactly

Per token, per layer, we store one key and one value for each key-value head:

$$\text{bytes per token}=2\times L\times h_{kv}\times d_h\times \text{bytes per element}.$$

The leading $2$ is for K and V, not for the byte width — that is the last factor. Multiplying by sequence length and batch:

$$\boxed{\text{KV cache}=2\,L\,h_{kv}\,d_h\,S\,B\cdot\text{bytes}.}$$

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 token8k context, 1 sequence128k context, 32 sequences
Multi-head (MHA)642560 KiB20.0 GiB10 TiB
Grouped-query (GQA)8320 KiB2.5 GiB1.25 TiB
Multi-query (MQA)140 KiB0.31 GiB160 GiB
The cache is per sequence, and it does not shrink. Model weights are shared by every request in a batch; the cache is not. At long context and moderate batch it dwarfs the model itself, which is why serving systems spend so much effort on paging and eviction, and why the next note's whole purpose is to compress it.

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

$$I=\frac{W}{Q}\quad\text{FLOPs per byte.}$$

Given peak compute $F_{\text{peak}}$ and peak bandwidth $\beta_{\text{peak}}$, achievable throughput is

$$\text{FLOP/s}\le\min\big(F_{\text{peak}},\ I\cdot\beta_{\text{peak}}\big).$$

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,

$$I^\star=\frac{989\times10^{12}}{3.35\times10^{12}}\approx 295\ \text{FLOPs per byte}.$$
You must do about 300 FLOPs on every byte you read just to keep an H100 busy. Anything below that leaves the tensor cores idle waiting for memory, no matter how the kernel is written.

Decoding sits at intensity equal to the batch size

Now apply it. Generating one token with batch size $B$:

1

Bytes. Every weight must be read once, regardless of $B$. In BF16 that is $2N$ bytes.

2

FLOPs. Each of the $B$ sequences performs $2N$ FLOPs, so $2NB$ in total.

3

Intensity.

$$I=\frac{2NB}{2N}=B.$$

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:

$$t_{\min}=\frac{2N}{\beta_{\text{peak}}}=\frac{1.4\times10^{11}}{3.35\times10^{12}}\approx 42\ \text{ms per token},$$

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.

KV cache size
decode intensity
peak compute reachable
min ms / token
A 70B-parameter, 80-layer model with $d_h=128$ in BF16, on hardware with 989 TFLOP/s and 3.35 TB/s. Both axes are logarithmic. The ridge sits at 295 FLOPs per byte.

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.

This is the whole motivation for MQA, GQA and MLA. They do not reduce FLOPs meaningfully. They reduce the bytes of cache that must be streamed per token, which is the quantity that actually sets decoding speed.

MQA and GQA

The observation behind both: nothing forces the number of key-value heads to equal the number of query heads.

SchemeKey-value headsCache versus MHAQuality
MHA$h_{kv}=h$baselinebaseline
MQA$h_{kv}=1$$h\times$ smallermeasurable degradation; can be unstable to train
GQA$1$h/h_{kv}\times$ smallerclose 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.

Takeaway

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