KV cache 如何佔用 GPU 記憶體，以及 MQA、GQA 如何降低需求

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:

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.