KV Cache Math: Sizing the Dominant Cost

The KV cache is what makes decode affordable: instead of re-attending over the whole prompt at every step, each layer stores the key and value vectors it already computed and reads them back. That trade — recompute for memory — is the right call, but the memory is not free. This page makes it a number you can compute on paper, because until you can size the KV cache you cannot answer the throughline question for inference: what does this cost in memory, and how many requests can one GPU hold?

The formula

For one request, the total bytes the KV cache occupies are:

kv_bytes = 2 × n_layers × n_kv_heads × head_dim × seq_len × batch × bytes_per_elem
           │
           └─ the 2 is K and V: every token stores both a key and a value vector

Read it left to right: each of n_layers layers keeps, for each of n_kv_heads heads, a head_dim-sized key and value, for every one of seq_len tokens, for every request in the batch, at bytes_per_elem per number. Nothing exotic — it is just counting the vectors the cache must hold.

Worked example: Llama-2-13B

Llama-2-13B has n_layers = 40, n_heads = 40, head_dim = 128, and in FP16 bytes_per_elem = 2. It uses standard multi-head attention, so n_kv_heads = 40. Take seq_len = 2048, batch = 1.

First, the cost of a single token across the whole model:

per_token = 2 × 40 × 40 × 128 × 2  = 819,200 bytes ≈ 0.78 MB / token

That is the figure to memorize: ~0.78 MB of KV per token for this model. Now scale to the full context:

kv_bytes = 0.78 MB × 2048  ≈ 1,677,721,600 bytes  = 1.5625 GB  per request

So one 2048-token conversation costs ~1.56 GB of HBM in KV alone. That number drives everything below.

Try it yourself: the defaults below (32 layers, 32 KV heads, head_dim 128, 4096 ctx, bf16) come out to ~2 GiB. Drop num_kv_heads to 8 to see GQA’s effect, or flip the dtype to fp8/int8 to halve it.

It scales linearly in batch and context

Both seq_len and batch are simple multipliers, so the cache grows linearly in each:

Scenario	KV cache
batch 1, 2048 ctx	1.56 GB
batch 16, 2048 ctx	25 GB
batch 32, 2048 ctx	50 GB
batch 1, 8192 ctx	6.25 GB

Long context is expensive twice over — and notice it grows per request, so a long-context, high-batch workload is the worst case: batch 32 at 8192 tokens would want 32 × 6.25 = 200 GB, far past any single GPU.

Why the KV cache, not the weights, caps your batch

Weights are a fixed cost: 13B parameters × 2 bytes ≈ 26 GB, loaded once, shared by every request in the batch. The KV cache is a per-request cost that stacks. On an 80 GB A100:

80 GB total
− 26 GB weights (fixed, shared)
− a few GB activations / framework overhead
≈ 54 GB left for KV cache

max concurrent requests ≈ 54 GB ÷ 1.56 GB ≈ 35 requests @ 2048 ctx

That ~35 is your ceiling — and throughput is essentially proportional to how many requests you can batch (next page). So the KV cache is the throughput bottleneck: the weights set the floor on memory, but the KV cache sets the ceiling on concurrency, and concurrency is what pays the bills. Every later technique in this part is, in part, a fight to fit more KV into that 54 GB.

The weights are the cheap part

It is tempting to budget HBM as “does the model fit?” The real question is “how much is left for KV, and how many requests does that buy?” A model that fits comfortably can still deliver terrible throughput if its KV cache is fat.

GQA and MQA: shrink n_kv_heads directly

Look at the formula again: n_kv_heads is a multiplier. Multi-Query Attention (MQA) uses a single shared KV head for all query heads (n_kv_heads = 1); Grouped-Query Attention (GQA) uses a small number, say 8 KV heads feeding 64 query heads. The query side keeps its full head count for quality; only the cached K and V shrink.

If Llama-2-13B used GQA with n_kv_heads = 8 instead of 40:

kv_bytes = 2 × 40 × 8 × 128 × 2048 × 2  ≈ 0.3125 GB  per request   (5× smaller)
max concurrent ≈ 54 GB ÷ 0.3125 GB ≈ 172 requests   (vs 35)

A 5× cut in KV bytes buys roughly 5× the concurrency — which is exactly why every modern large model (Llama-2-70B, Llama-3, Mistral, and friends) ships GQA. It is the cheapest single lever on inference cost, paid once at training time.

Back-of-napkin rule

KV per token ≈ 2 × n_layers × n_kv_heads × head_dim × bytes. Compute it once for your model; multiply by context and batch to get HBM; subtract weights from GPU memory to get your concurrency ceiling. That three-step estimate predicts throughput better than any benchmark you haven’t run yet.

It actually happened — MQA (2019) waited four years for GQA (2023)

Multi-Query Attention was proposed by Noam Shazeer in 2019 (“Fast Transformer Decoding: One Write-Head Is All You Need”) — collapse all KV heads to one and shrink the cache n_query_heads-fold. It was right but largely ignored, because dropping to a single KV head cost noticeable quality and few teams were yet serving at a scale where KV memory hurt. It took until 2023 and Grouped-Query Attention (Ainslie et al., Google) — a middle ground of a few KV heads, plus a recipe to uptrain an existing multi-head checkpoint into GQA cheaply instead of retraining from scratch — for the idea to go mainstream. Llama-2-70B shipped with GQA, and from then on essentially every new large model did too. A four-year gap between a correct idea and its adoption, closed exactly when inference cost finally became the thing that mattered.

The architect’s lens

The lever this page introduces is GQA/MQA — shrinking n_kv_heads to lift the concurrency ceiling. The five questions:

• Why does it exist? Because n_kv_heads is a direct multiplier in the KV formula, and it is the KV cache — not the weights — that caps how many requests fit in HBM. Grouped/Multi-Query Attention cut that multiplier so the cache shrinks proportionally.
• What problem does it solve? The concurrency ceiling: Llama-2-13B with full 40-head KV holds ~35 requests on an 80 GB card; dropping to 8 KV heads (GQA) cuts KV ~5× and lifts that to ~172 — and throughput tracks concurrency.
• What are the trade-offs? Quality for memory: the query heads stay at full count, but MQA’s single shared KV head costs noticeable quality, so GQA’s handful of heads is the deployed middle ground. And it is baked in at training time — not an inference toggle.
• When should I avoid it? Push to MQA (n_kv_heads = 1) only where the quality hit is tolerable; you cannot add GQA to an existing model without uptraining it; and if KV isn’t your bottleneck (short context, small batch) the win is moot.
• What breaks if I remove it? KV reverts to full multi-head — ~5× larger, ~5× fewer concurrent requests, throughput collapses — and a long-context, high-batch workload (batch 32 at 8,192 tokens ≈ 200 GB) becomes impossible to hold on any single GPU.

Check your understanding

1. In the KV formula, what does the leading factor of 2 represent, and what does bytes_per_elem depend on?
2. Compute the KV cache for Llama-2-13B at batch 4, 4096 context, FP16. (Use ~0.78 MB/token.)
3. Why is it the KV cache rather than the model weights that limits how many requests you can batch?
4. How do GQA and MQA reduce the KV cache, and what do they deliberately not shrink?
5. A team doubles context from 2048 to 4096 and keeps batch fixed. What happens to KV memory, and what must give if the GPU was already near full?

Show answers

1. The 2 counts both the key and the value vector stored per token. bytes_per_elem depends on the numeric precision — 2 bytes for FP16/BF16, 1 for INT8/FP8, etc.
2. 0.78 MB/token × 4096 tokens × 4 requests ≈ 12.8 GB.
3. Weights are a fixed, shared cost loaded once; the KV cache is a per-request cost that stacks with the batch. So after the weights are loaded, remaining HBM ÷ KV-per-request sets the concurrency ceiling, and concurrency drives throughput.
4. They cut n_kv_heads — MQA to 1 shared KV head, GQA to a small group — which shrinks the cached keys and values proportionally. They keep the full number of query heads, so model quality is largely preserved.
5. KV memory doubles (it is linear in seq_len). If the GPU was near full, you must reduce the batch size — fewer concurrent requests — to fit, which lowers throughput.