multi head attention

Multi-head attention (MHA), multi-query attention (MQA), and grouped-query attention (GQA) are three ways to wire the key and value projections of a Transformer's attention layer. They all keep the same set of query heads, each looking at the sequence from a different learned subspace; what changes is how many independent key/value heads those queries share. That single choice trades model quality against the size of the KV cache — the per-token memory that dominates the cost of generating long outputs — which is why nearly every recent large model has moved from MHA toward GQA.\n\n**Multi-head attention gives every query head its own keys and values.** Rather than computing one attention over the full model dimension, MHA splits the vectors into H heads, and each head runs its own scaled dot-product attention over its own query, key, and value projections. Different heads specialize — one tracks syntax, another long-range coreference — and their outputs are concatenated and mixed. The cost is memory: during generation the model must cache the keys and values of every past token for all H heads, so the KV cache scales with the head count and quickly becomes the binding constraint at long context lengths.\n\n**MQA shares one KV head; GQA shares a few.** Multi-query attention keeps all H query heads but collapses the keys and values to a single shared head, so the KV cache shrinks by a factor of H. That is a large memory and bandwidth win — decoding is memory-bound, and a smaller cache means more tokens and more concurrent requests fit — but forcing every query to read the same keys can cost accuracy and destabilize training. Grouped-query attention interpolates: the query heads are divided into G groups, each with its own KV head, so the cache shrinks by H/G. With, say, eight query heads in two groups, GQA recovers almost all of MHA's quality while still cutting the cache several-fold, which is why models like Llama 2/3 and Mistral adopt it.\n\n| | MHA | GQA | MQA |\n|---|---|---|---|\n| Query heads | H | H | H |\n| KV heads | H | G (1 Multi-Head Attention — MHA vs GQA vs MQA key/value head sharing reduces KV-cache memory without losing quality MHA (standard) every head has its own K, V Q heads: h=32 K heads: h=32 V heads: h=32 KV-cache: 32 × d_head × seq full KV per head GPT-2/3, BERT, original 100% KV memory Params: 4 × h × d² per layer best quality, most memory used in: GPT-3, BERT 32 Q × 32 K × 32 V GQA (grouped) groups of Q heads share K, V Q heads: h=32 K heads: h=8 V heads: h=8 KV-cache: 8 × d_head × seq 4 Q heads share 1 KV pair Llama 2/3, Mistral, Gemma 25% KV memory ~97% of MHA quality best speed/quality trade-off used in: Llama 3, Mistral 32 Q × 8 K × 8 V MQA (single KV) all Q heads share 1 K, 1 V Q heads: h=32 K heads: 1 V heads: 1 KV-cache: 1 × d_head × seq 32 Q heads → single KV PaLM, Falcon, StarCoder 3% KV memory ~95% of MHA quality fastest inference, min memory used in: PaLM, Falcon 32 Q × 1 K × 1 V KV-Cache Memory at 128K Context (70B model, fp16) MHA: 160 GB (won't fit 1 GPU) GQA-8: 40 GB (fits H100) MQA: 5 GB (fits any GPU) GQA is the 2024 consensus: near-MHA quality with 4–8× less KV-cache — the reason 128K context is affordable. ```\n\n**The whole point is the KV cache, so this is a serving decision.** Because autoregressive decoding is limited by memory bandwidth and by how many sequences' KV caches fit in GPU memory, shrinking the per-token KV footprint directly raises throughput and the maximum context length you can serve. GQA has become the default precisely because it sits at the sweet spot of that curve — most of the memory savings of MQA with almost none of the quality loss of MHA. It also composes with everything else in the stack: a smaller KV cache means PagedAttention has fewer blocks to manage, continuous batching can hold more requests, and Flash Attention still applies within each head. Multi-head latent attention (MLA) pushes the same idea further by caching a compressed latent instead of full keys and values.\n\nRead MHA/MQA/GQA through a quant lens rather than a 'number of heads' lens: the number they move is bytes of KV cache per token, which equals two times the KV-head count times the head dimension times precision, and that figure sets both decode bandwidth and how many sequences share a GPU. MHA fixes KV heads at H, MQA at 1, and GQA at a tunable G, so the design question is how far you can drop G before the shared keys stop giving each query enough distinct context — empirically a handful of groups keeps quality at MHA levels while capturing most of MQA's memory win.