MHA vs GQA vs MQA Visualization

The idea: share the keys and values, keep the queries

An attention layer has several heads. Each head has its own query, key and value projections, so it can look for a different pattern in the text (see Self-Attention). When a model writes text, it keeps every past token's keys and values in the KV cache, so the size of that cache depends on how many key/value heads there are. The three designs differ only there. Multi-head attention (MHA) gives every query head its own K/V head. Grouped-query attention (GQA) splits the query heads into groups, and each group shares one K/V head. Multi-query attention (MQA) has a single K/V head for all of them.

The query heads stay the same; only the number of K/V heads they read changes: 8, 2 or 1.

The canvas runs the same 8 query heads on the same sentence in all three lanes. Each new token writes one K,V pair per K/V head into the cache, and every query head then reads the whole cache row of its K/V head. Under each lane are the real sizes for Llama 3 and a context length you pick.

What the cache stores per token

For every token, every layer stores one key vector and one value vector per K/V head:

bytes per token = 2 × layers × K/V heads × head_dim × bytes per number

Llama 3 70B has 80 layers, 64 query heads and head_dim 128, and it uses GQA with 8 K/V heads. The MHA and MQA rows below are the same model with 64 or 1 K/V heads, to compare. Numbers are fp16 (2 bytes). (Demo: one token — 8 vs 2 vs 1 K/V pairs written)

model	MHA	GQA (8 K/V heads)	MQA
Llama 3 70B, per token	2.5 MiB	320 KiB	40 KiB
Llama 3 70B, 4,096 tokens	10 GiB	1.25 GiB	160 MiB
Llama 3 8B, per token	512 KiB	128 KiB	16 KiB
Llama 3 8B, 4,096 tokens	2 GiB	512 MiB	64 MiB

Sequences that fit, and bytes read per token

The cache lives in GPU memory next to the weights. Llama 3 70B in fp16 is about 140 GB, so on two 80 GB GPUs about 20 GiB is left for the cache. The page divides that by one sequence's cache to get how many sequences a server can run at once, which is the batch size of continuous batching.

model, context	MHA	GQA	MQA
70B, 4K (20 GiB)	2	16	128
70B, 32K	0	2	16
70B, 128K	0	0	4
8B, 32K (60 GiB)	3	15	120

Memory size is half the story. Each new token reads the whole cache of its sequence, and decoding is limited by how fast the GPU reads memory, not by arithmetic. With 8 times fewer K/V heads each token reads 8 times fewer bytes. (Demo: Llama 3 70B at 4K — 2 vs 16 vs 128 sequences, Demo: 128K context on 70B — only MQA fits)

Quality vs memory

Sharing has a price. With one K/V head, all query heads must find their patterns in the same keys and values, and the model gets a bit worse. The GQA paper measured this on T5-XXL: GQA with 8 groups almost matches MHA's quality while running almost as fast as MQA. (Demo: quality vs memory — what each step down costs)

A chart of time per sample against average score from the GQA paper on T5-XXL: MHA 1.51 seconds and 47.2, GQA-8 0.28 seconds and 47.1, MQA 0.24 seconds and 46.6. The smaller MHA-Large model takes 0.37 seconds and scores 46.0. — GQA-8 keeps MHA's score (47.1 vs 47.2) at close to MQA's speed (0.28 s vs 0.24 s).

T5-XXL variant	K/V heads	time per sample	average score
MHA	64	1.51 s	47.2
GQA-8	8	0.28 s	47.1
MQA	1	0.24 s	46.6

Turning an MHA model into GQA

A model does not have to be trained with GQA from scratch. The GQA paper takes an MHA checkpoint, mean-pools the key and value projections of the heads in each group into one, and then uptrains the model for about 5% of its original training steps so it adapts to the shared heads. MQA is the same step with a single group.

Before: an MHA checkpoint with 8 key/value projections. The projections of heads 1 to 4 are averaged into one, and those of heads 5 to 8 into another. After: a GQA model with 2 K/V heads, then uptrained for about 5 percent of the original steps. — Converting MHA to GQA: average the K/V projections inside each group, then train a little more.

Who uses what

MQA came first (Shazeer, 2019) and is used by PaLM and Falcon. GQA is used by Llama 2 70B, Llama 3, Mistral and most newer open models, usually with 8 K/V heads. It also makes tensor parallelism simple: 8 K/V heads split cleanly over 8 GPUs, while one MQA head has to be copied to all of them. GQA does not change how attention is computed, so it combines with FlashAttention, paged caches and quantized caches. See also How a Large Language Model Works for the token-by-token loop.

What the page leaves out

The attention math itself (scores, softmax, the weighted sum), the query and output projections, the activations and other memory a server needs besides weights and cache, prefill, paged and quantized caches, multi-head latent attention (DeepSeek), which compresses K and V into a smaller latent vector, sliding-window attention, and tensor-parallel splits of the heads.