Everyone talks about attention as a compute problem. At inference it's actually a memory problem, and grouped-query attention is the trick that makes long-context LLMs practical. When a transformer generates text one token at a time, each new token has to attend to the keys and values of every token before it. Recomputing that history every step is O(seq²) waste, so you cache it — and that KV cache is the real bottleneck. I built a demo where you pick the number of query heads, drag the number of KV heads from 1 to H, and watch a real head-grouping map redraw while the cache bytes, savings and compression are computed live from the genuine formula. Here's what it taught me.
The KV cache is huge — and re-read every single token
The cache is a growing tensor [batch, n_kv_heads, seq, d_head], for both K and V, at every layer. Its size:
def kv_cache_bytes(n_kv, d_head, seq, layers, dtype_bytes):
return 2 * n_kv * d_head * seq * layers * dtype_bytes
# 7B-ish model, 8k context, fp16:
kv_cache_bytes(32, 128, 8192, 32, 2) # MHA → 4.00 GB
At long context or large batch this dwarfs the model weights. But the sharper pain is bandwidth: generating each token requires reading the entire cache back from GPU memory. Modern accelerators have far more compute than memory bandwidth, so attention at inference is memory-bound — the GPU sits waiting on the cache. Shrink the cache and every token gets faster.
Only one term in that formula is an architectural knob
Look at the formula: d_head and layers come with the architecture, seq and batch come with the request, dtype you can quantize a bit. The one term you can cut hard without touching model width — because the query heads (and so d_model = H · d_head) stay put — is n_kv_heads. Memory is dead linear in it. Halve the KV heads, halve the cache. That linearity is the whole lever.
MHA and MQA are the two ends; GQA is the dial
Classic Multi-Head Attention gives every query head its own KV head — best quality, biggest cache. Multi-Query Attention goes to the extreme: all H query heads share one KV head — smallest possible cache, but a real quality hit and shakier training. GQA sits in between: split H query heads into G groups, each sharing one KV head.
G = 8 # KV heads / groups (a divisor of H)
n_rep = H // G # query heads per KV head
k = x @ Wk # → [B, G, T, d_head] only G cached
# G=1 → MQA , G=H → MHA , else → GQA
The elegant part is that MQA and MHA aren't separate designs — they're just G=1 and G=H of the same family. One knob spans the whole space.
repeat_kv — the grouping in one function
You store only G KV heads (small cache), then repeat each one H/G times so the shapes line up with the H query heads, only transiently, inside the matmul:
def repeat_kv(x, n_rep):
B, G, T, d = x.shape
if n_rep == 1: return x # MHA: nothing to do
x = x[:, :, None, :, :].expand(B, G, n_rep, T, d)
return x.reshape(B, G * n_rep, T, d)
This is the exact helper Llama and Mistral use. The memory win lives in the cache, not the math — the expansion to H never gets stored.
The tradeoff, and why 8 is the magic number
Here's the payoff the GQA paper found and my demo's spectrum chart makes visual: memory is a straight line in G, but quality rises fast and then saturates. You don't need one KV head per query head — you need a handful. Around 8 KV heads recovers nearly all of MHA's quality while the cache stays a fraction of its size. That gap between a linear memory line and a curve that flattens early is the entire idea. It's why Llama-2 70B (64 query heads, 8 KV heads → 8× cut), Llama-3 and Mistral 7B all ship GQA-8. And you don't even have to train it from scratch: mean-pool an MHA checkpoint's KV heads within each group and uptrain ~5% of the original steps to recover.
Pick H, drag G from MQA to MHA, and watch the cache bytes and savings move live:
https://dev48v.infy.uk/dl/day46-grouped-query-attention.html
Top comments (0)