Grouped-Query Attention: How Sharing KV Heads Shrinks Inference Memory

Where the memory actually goes

During decode, a transformer produces one token at a time, and attention at each layer needs the keys and values of every earlier token in the sequence. Recomputing them each step would be quadratic, so servers cache them. The size of that cache per token is a short formula worth memorising:

kv_bytes_per_token = 2 (K and V) x layers x kv_heads x head_dim x bytes_per_value

# an 8B-class decoder: 32 layers, 32 query heads, head_dim 128, FP16 cache
MHA   (32 KV heads): 2 x 32 x 32 x 128 x 2 = 524,288 B  = 512 KiB per token
GQA-8 ( 8 KV heads): 2 x 32 x  8 x 128 x 2 = 131,072 B  = 128 KiB per token
MQA   ( 1 KV head ): 2 x 32 x  1 x 128 x 2 =  16,384 B  =  16 KiB per token
GQA-8 with FP8 cache:                                      64 KiB per token

Multiply by context length and the picture changes fast. A single 8,000-token conversation needs 4 GB of cache with MHA but 1 GB with eight KV heads. On a 48 GB card that reserves 90% of memory for the server, roughly 25.7 GB is left once 16 GB of weights and a working reserve for activations are subtracted. That holds about 49,000 tokens of MHA cache, or about 196,000 tokens with GQA-8. The same hardware, the same parameter count, and four times as many tokens in flight.

Notice what the formula does not contain: the number of query heads. Queries are computed fresh at every step and thrown away. Only keys and values persist, which is why sharing them is such an effective lever.

Decode is a bandwidth problem, and the cache is part of the bill

Memory capacity is the obvious win; bandwidth is the quieter one. A decode step is memory-bound: the GPU streams every weight matrix once and, for attention, reads the cached keys and values of every running sequence. On an 8B model the weights alone are 16 GB per step. The KV read is the sum over sequences of context length times bytes per token, and at long contexts it overtakes the weights.

Take 16 running sequences that each hold about 2,300 tokens. With MHA that is 16 x 2,300 x 512 KiB, about 19.3 GB, on top of 16 GB of weights. At 0.86 TB/s, a step costs roughly 41 ms before any overhead. With GQA-8 the KV read drops to about 4.8 GB and the step to roughly 24 ms. Every user streaming tokens sees that difference directly as time per output token (TPOT). These are calculated, illustrative figures, but the proportion holds on real hardware because attention kernels for decode are bandwidth-bound.

This is also why GQA helps more as contexts get longer: the weight read is fixed, while the KV read grows with every token. At 500-token chats the gain is modest; at 30,000-token retrieval prompts it dominates.

How grouped-query attention works

Standard multi-head attention projects the hidden state into H query heads, H key heads and H value heads; head i of the query attends over head i of the keys and mixes head i of the values. Multi-query attention (MQA), proposed by Shazeer in 2019, keeps H query heads but only one key head and one value head shared by all of them. That cuts the cache by a factor of H, but in practice it can cost quality and training stability.

GQA, introduced by Ainslie and colleagues in 2023, interpolates between the two. Query heads are divided into G groups; each group shares one key head and one value head. With 32 query heads and G = 8, four query heads read the same K and V. The attention computation itself is unchanged; the kernel simply maps query head h to KV head floor(h / (H / G)). Everything else about the layer stays the same, so the parameter count drops only slightly (the K and V projection matrices shrink) while the cache shrinks by H / G.

Many widely used open models ship with GQA. Llama 3 8B and Mistral 7B, for example, use 32 query heads with 8 KV heads, and Llama 2 70B used 8 KV heads for 64 query heads. In Hugging Face model configs the field is num_key_value_heads; when it equals num_attention_heads the model is plain MHA. Check the config of the exact checkpoint you deploy rather than assuming from the model family.

Why not always go to one KV head?

If eight KV heads beat thirty-two, one should beat eight. In memory terms it does: MQA has the smallest cache possible for a given depth and head size. The cost is representational. Different query heads learn to look for different things, and when they are forced to read the same keys, they can only differ in their queries. Some of that can be compensated during training, but the GQA paper reports that MQA models lose more quality than GQA models with a moderate number of groups, while GQA with a handful of groups gets close to MHA quality at a speed close to MQA.

There is also a systems reason. With tensor parallelism, attention heads are split across GPUs. If a model has 8 KV heads and you shard it over 8 GPUs, each GPU owns one KV head. With MQA on 8 GPUs, the single KV head has to be replicated on every GPU, so the per-GPU cache saving relative to GQA largely disappears. Servers such as vLLM handle a tensor-parallel size larger than the KV head count by replicating KV heads; check your version's documentation for the exact behaviour. As a rule of thumb, a KV head count that is a multiple of your tensor-parallel degree avoids waste.

The diminishing returns are visible in the arithmetic too. Going from 32 to 8 KV heads saves 384 KiB per token; going from 8 to 1 saves only 112 KiB more. Once the KV read is small next to the weight read, further cuts barely change step time.

Converting an MHA checkpoint: mean-pooling and uptraining

You rarely choose KV heads for a model you did not train, but you may meet the question when a team wants a cheaper variant of an existing MHA model. The GQA paper describes the recipe: build each group's key and value projection by mean-pooling the projections of the heads in that group, then continue pretraining (uptraining) for a small fraction of the original compute so the model adapts to the shared heads. The authors report that around 5% of the original pretraining compute recovered most of the quality.

The pooled checkpoint without uptraining is a different model. Heads that used to look at different features are averaged together, attention patterns blur, and the damage is larger the more heads you merge. The pseudo-code makes it plain how blunt the operation is:

# W_k, W_v: [num_heads, head_dim, hidden] from the MHA checkpoint
group = num_heads // num_kv_heads
W_k_gqa = W_k.reshape(num_kv_heads, group, head_dim, hidden).mean(axis=1)
W_v_gqa = W_v.reshape(num_kv_heads, group, head_dim, hidden).mean(axis=1)
# then: set num_key_value_heads = num_kv_heads in the config and uptrain

In production the lesson is simple: never ship a pooled checkpoint on the strength of its memory numbers. Run your task evaluations, compare against the source model, and budget the uptraining compute as part of the project. Mission three in the lab starts from exactly this mistake.

FP8 KV cache: the other factor of two

The formula has one more factor you can change without touching the model: bytes per value. Most servers can store the cache in an 8-bit floating-point format while keeping weights and activations at higher precision. In vLLM this is --kv-cache-dtype (values such as fp8, fp8_e4m3 or fp8_e5m2, depending on version and hardware); SGLang and TensorRT-LLM expose equivalent options. Some formats use per-tensor scaling factors that are calibrated or shipped with the checkpoint. Check your version's documentation for the supported formats and whether scales are required.

FP8 halves both the cache footprint and the KV bytes read per step, and it composes with GQA: GQA-8 with FP8 is 64 KiB per token on the example model, an eighth of the MHA FP16 figure. The quality cost is usually small but not zero, and it depends on the model and on how outliers in keys are handled, so evaluate it like any other quantization. Because it is a serving flag rather than a new model, it is often the fastest fix when a traffic change suddenly needs more cache. Mission two in the lab asks you to absorb a day of three-times-longer prompts on an MHA model you are not allowed to swap.

Serving knobs that interact with KV heads

GQA changes how many tokens fit, but the server decides how that capacity is used. The knobs below are the ones to revisit whenever the KV footprint changes:

TGI and TensorRT-LLM have equivalents under different names, and defaults change between releases. On startup, most servers log how many KV blocks or tokens of cache they allocated; record that number for each deployment, because it is the real capacity figure and it changes when you change model variant, precision or memory share.

A worked example: Belgavi's Book Store assistant

Suppose the hypothetical Belgavi's Book Store runs a catalogue assistant on one 48 GB card with an 8B-class model. Prompts average about 2,000 tokens because each one carries retrieved product descriptions, answers average 256 tokens, and traffic is about one request per second. With an MHA variant, each live conversation of roughly 2,300 tokens holds about 1.15 GB of cache, so 25.7 GB fits around 22 conversations. That sounds like plenty, until you add the residency time: at about 70 ms per token, a 256-token answer keeps its cache for 18 seconds, so a steady one request per second already keeps 18 or more conversations resident, and any burst spills into queueing and preemption.

Switching to a GQA-8 variant of the same size changes three numbers at once (all calculated by the lab's model, not measured): cache per conversation falls to about 0.29 GB, the decode step falls because each step reads a quarter of the KV bytes, and p95 time to first token drops from several seconds to under one. Cost per million output tokens falls too, because the same card produces more tokens per second. The price is a small loss on the quality proxy, which in a real project is a set of task evaluations that someone must run and sign off.

The procedure transfers: compute bytes per token from the config, compute resident tokens from arrival rate times residency time times context, compare with the KV budget the server reports, and only then decide between a different model variant, an FP8 cache, a bigger card or more replicas.

Failure modes and how to spot them

Most KV-capacity incidents show the same signatures, and they are easy to misread as compute problems:

• Preemption storms. When decode needs a new cache block and none is free, the server evicts a running sequence and recomputes it later. Watch the preemption counter your server exports; a steady non-zero rate means the cache, not the GPU's arithmetic, is the bottleneck.
• TTFT climbs while GPU utilization looks fine. Requests are waiting for cache blocks, not for compute. Queue depth and the number of waiting requests tell the story better than utilization.
• Startup refusal after a model swap. A new variant with more KV heads, or a larger max-model-len, may no longer fit one maximal request; the server fails to start or caps the context.
• Silent quality regressions. A pooled or aggressively quantized cache can pass smoke tests and still degrade long-context retrieval or multi-step reasoning. Keep a regression suite that includes long prompts.
• Tensor-parallel surprises. Doubling GPUs on a model with fewer KV heads than ranks may not double cache capacity, because heads are replicated.

Common mistakes cluster around assuming parity: copying a capacity plan from an MHA model to a GQA one or the reverse, sizing for average context instead of the tail, and treating an FP8 cache as free without an evaluation.

Run the lab: KV heads on a budget: fit more users per GPU with GQA

The lab simulates sixty seconds of synthetic traffic for an 8B-class model with 32 query heads. The top of the canvas shows how query heads map to KV heads; the chart tracks KV cache in use against the capacity left after weights; the bar at the bottom shows the GPU's memory split at the moving cursor. Mission one asks you to fit a chat load on a small card, mission two adds a long-document day, mission three repairs a badly converted checkpoint, and mission four hunts for the cheapest card that survives a burst.

Prefer a full-width view? Open the standalone lab. Your stars are saved in this browser and shown on the series hub.

What to do next

1. Read num_attention_heads, num_key_value_heads, layer count and head dimension from each deployed model's config and compute KV bytes per token.
2. Record the KV cache size your server reports at startup for every deployment, and alert when a config change reduces it.
3. Estimate resident tokens as arrival rate x residency time x mean context at peak, and compare with the reported KV budget.
4. Export and alert on preemption counts and waiting-queue length rather than GPU utilization alone.
5. Trial an FP8 KV cache on a staging replica with your long-context evaluations before relying on it for capacity.
6. If you plan to convert an MHA checkpoint to GQA, budget uptraining compute and a full evaluation run before any traffic sees it.

Sources and further reading

• GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (Ainslie et al., 2023)
• Fast Transformer Decoding: One Write-Head is All You Need (Shazeer, 2019)
• Attention Is All You Need (Vaswani et al., 2017)
• Efficient Memory Management for Large Language Model Serving with PagedAttention (SOSP 2023)
• vLLM documentation: engine arguments
• DeepSeek-V2: multi-head latent attention as another way to shrink the KV cache

Related in this series:

• RoPE and Long Context: Why Models Fail Past Their Trained Length
• Why INT4 Breaks Some Layers: Outliers, Group Size and Activation-Aware Scaling
• KV Cache Capacity Planning: How Many Concurrent Users Fit on One GPU

Diagrams, examples and the lab model are original educational material. Lab numbers are calculated by a synthetic model, not measured on real hardware or services.

Views are my own and do not represent any employer, past or present.