The KV cache is eating your GPU. MQA and GQA are how models fight back

Quick recap from the last post: the KV cache saves you from recomputing old tokens, and it's the reason generation is fast.

But it has a bill. And the bill is memory.

I wanted to put an actual number on it, because "the cache is big" doesn't mean much until you see the GB.

The formula

Per token, you store a Key and a Value for every layer and every head:

bytes per token = 2 x layers x kv_heads x head_dim x bytes_per_number

The 2 is K and V. fp16 is 2 bytes per number. Here it is as code:

def kv_bytes(layers, kv_heads, head_dim, tokens, nbytes=2):
    return 2 * layers * kv_heads * head_dim * nbytes * tokens

Now plug in a Llama-2 7B style model: 32 layers, 32 heads, head dim 128.

That's 512 KB for every single token. At 8k context, 4 GB. At 128k context, 64 GB. For one request. Weights not included.

If you're serving many users at once, each one has their own cache. That's the thing that actually fills your GPU, not the model.

Look at the formula again

Layers and head_dim are baked into how good the model is. Bytes per number you can shrink with quantization, but that's a different post.

The interesting variable is kv_heads.

In standard multi-head attention (MHA), every query head has its own key head and value head. 32 query heads, 32 KV heads.

But do the heads really need separate Keys and Values? Turns out, mostly no.

MQA: one KV head for everybody

Multi-Query Attention keeps all the query heads but shares a single K head and V head across all of them. 32 query heads, 1 KV head.

Memory goes down by 32x. The catch is quality: forcing every head to read from one shared set of keys and values can hurt, and training gets a bit touchier.

GQA: the compromise

Grouped-Query Attention sits in between. Split the query heads into groups, and each group shares one KV head.

Llama-3 8B has 32 query heads but only 8 KV heads. So 4 query heads share each K/V pair. That's a 4x cut in cache size, and the published results say quality stays close to full MHA.

Same model shape, three attention styles:

Style	KV heads	KB per token	8k context	128k context
MHA (Llama-2 7B)	32	512	4 GB	64 GB
GQA (Llama-3 8B)	8	128	1 GB	16 GB
MQA (hypothetical, same shape)	1	16	0.12 GB	2 GB

The MQA row is me taking the same shape with one KV head, to show the extreme. It's not a model I measured.

What changes in the code

Less than you'd think. The only real difference is that you repeat the KV heads to line up with the query heads at attention time:

# k, v: (kv_heads, t, head_dim)   q: (q_heads, 1, head_dim)
group = q_heads // kv_heads
k = np.repeat(k, group, axis=0)   # now (q_heads, t, head_dim)
v = np.repeat(v, group, axis=0)

Important detail: you store the small version in the cache, and expand it only when computing. That's where the memory saving lives. (Real implementations avoid the copy, but the idea is the same.)

Why this matters if you're building agents

You probably aren't training a model. But:

• When a provider says a model supports 128k or 1M context, GQA-style tricks are a big part of why that's affordable to serve at all.
• Long contexts are cheaper per token on models built this way, which shows up in latency and sometimes in price.
• If you self-host, KV memory is the number that decides how many concurrent users fit on one GPU. Do this arithmetic before buying hardware.

Where i'd be careful

• Those numbers are for one sequence at fp16. Batching multiplies them, quantizing the cache divides them.
• "Little quality loss" comes from papers and model cards. I haven't run my own quality evals here, and you shouldn't take it from a blog post either. Check the benchmarks for the specific model you use.
• The GB figures are straight arithmetic from the formula, not profiler output. Real servers add overhead.

Next up: why the same sentence can cost very different amounts depending on the tokenizer. Slightly less hardware, more wallet.

Comments

[Max Kaya]

Nice writeup, and the byte-per-token framing is the part people usually skip.

One number I'd add for the agent case specifically: a coding agent that burns 40k tokens of context at 8 KV heads (Llama-3 style) is carrying ~5 GB of KV for a single request. Two of those on one GPU and you're already fighting the driver for memory, and that's before any weights. It's the reason we ended up truncating tool output aggressively rather than trying to be clever about caching.

Small correction on the expand-at-attention-time note: np.repeat there is the toy version, but the practical trap is that people then keep the expanded copy alive in the cache too, which gives the memory back. Worth stating explicitly that the un-expanded K/V has to be the thing that lives in the cache object, not just what you compute into.

Also curious whether you measured the actual hit rate effect of GQA on your own traffic. The 4x is the theoretical ceiling; in practice KV pressure mostly shows up as more concurrent sessions per card, and I'd be curious whether that translated into throughput or just more idle memory between bursts.