Read num_key_value_heads before treating GQA as a capacity win.
The Mon/Wed/Fri briefing is AI news.
Every model card on Hugging Face has a config file full of cryptic fields. One of them, num_key_value_heads, quietly decides how many conversations a GPU can hold at once and how long they can be. Five minutes with this lesson and you can do the math yourself.
Background: the model’s notes
While a model writes, it keeps notes on everything said so far, called the KV cache (K for keys, V for values). These notes live in GPU memory, and they grow with every token in every conversation.
Inside the model, attention heads read those notes. Each layer has many query heads (the “askers”) and some number of key/value heads (the “note-keepers”).
- Multi-head attention (MHA): every asker has its own note-keeper.
- Multi-query attention (MQA): all askers share one note-keeper.
- Grouped-query attention (GQA): askers split into groups, and each group shares a note-keeper.
The formula
Only the note-keepers show up in memory. Bytes of KV cache per token:
2 x num_layers x num_kv_heads x head_dim x bytes_per_element
The 2 is keys plus values. The number of query heads doesn’t appear at all. That’s the point of GQA and MQA: keep lots of askers while shrinking what you store.
Worked example
A Llama 3-class 8B model: 32 layers, 32 query heads, 8 KV heads, head_dim 128, BF16 (2 bytes):
2 x 32 x 8 x 128 x 2 = 131,072 bytes per token, about 128 KiB.
- At 8K tokens of context, that’s roughly 1 GiB for one conversation.
- The same shape with full MHA (32 KV heads) would be 4x that.
- A 32K-token chat on Llama-3.1 8B is about 4 GiB of cache; the MHA twin would be about 16 GiB.
- Llama 3.1 70B (80 layers, 64 query heads, 8 KV heads) lands around 320 KiB per token.
“Uses GQA” isn’t a fixed discount, either. Qwen2.5 7B uses 4 KV heads for a steeper cut; Gemma 2 often uses only a 2x cut. Read the field.
How we got here
Shazeer’s MQA (2019) showed that one shared note-keeper kills most of the memory-speed tax when writing. Ainslie et al. (2023) then showed a middle group count gets quality close to MHA while staying near MQA speed. Their recipe converted an existing MHA model by averaging heads, then continued training for about 5% of the original training compute. Labs later trained GQA from the start. For current open models, GQA with a small number of KV heads (often 8) is the default.
The other half: PagedAttention
Notes grow and shrink as conversations start and end. Older systems reserved one big contiguous chunk per request for its maximum length, and ended up using only about 20 to 38% of reserved memory, per vLLM’s profiling. PagedAttention splits notes into fixed-size blocks (commonly 16 tokens) stored wherever there’s room, tracked by a table. Waste drops to the last partial block, under about 4%. Sharing blocks also cut memory for parallel samples and beam search by up to about 55%.
An analogy: the library study room
Thirty-two students (query heads) are working on a group project. If each student keeps their own full binder of notes, the shelf fills fast. If they form groups of four and each group shares one binder, you need a quarter of the shelf space and they still get their work done. That’s GQA. The number of binders, not the number of students, decides how many projects fit in the room.
PagedAttention is the librarian who lets you store binder pages in any free slot and keeps an index, instead of reserving a whole empty shelf “just in case.”
One gotcha for multi-GPU setups
When a model is split across GPUs (tensor parallelism, TP), the KV heads get split too. As long as TP is at most the number of KV heads and divides it evenly, each GPU’s cache shrinks. Past that, engines have to copy whole KV heads. On Llama-3.1 70B, going from TP=8 to TP=16 can leave per-GPU cache per token flat (about 40 KiB) while total cache doubles.
Try this
Pick any open model on Hugging Face, open its config.json, and find num_hidden_layers, num_key_value_heads, and head_dim (or hidden_size divided by num_attention_heads). Plug them into the formula with 2 bytes per element. Multiply by 8,000 tokens. You now know roughly how much GPU memory one 8K-token chat costs on that model. Try two models and compare.
Use it at dinner
“How many people a GPU can serve at once depends on one line in the model’s config: how many note-keepers it has. Fewer note-keepers, more conversations.”
Sources
Subscribe
One lesson a week. Unsubscribe anytime.