Multi-Head Attention Optimization

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.

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.

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.

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.

Read 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.