Multi-Head, Grouped-Query, and Multi-Head Latent Attention

In large language model (LLM) inference, the bottleneck is rarely compute; it is memory bandwidth. As models scale, the time spent waiting for data to move from High Bandwidth Memory (HBM) to the compute units dominates latency. This constraint has driven the evolution of attention mechanisms from standard multi head attention (MHA) to Grouped-Query Attention (GQA) and, most recently, Multi-Head Latent Attention (MLA). These variants are not merely academic exercises; they are engineering solutions to the KV cache explosion.

Part 2 of the Transformers Under the Hood series.

The Baseline: Multi-Head Attention (MHA)

Standard MHA, introduced by Vaswani et al., divides the query, key, and value vectors into $h$ independent heads. Each head computes attention independently. During training, all parameters are updated. During inference, the model generates tokens autoregressively. For each new token, the model must compute the attention weights against all previous tokens.

The critical artifact here is the KV cache. To avoid recomputing keys and values for every new token, the model stores the key and value vectors for all previous positions in memory. If a model has $d_{model}$ dimensions, $h$ heads, and a context window of $L$ tokens, the KV cache size is approximately:

$\text{Size} = 2 \times n_{layers} \times h \times d_{head} \times L \times \text{batch\_size} \times \text{precision\_bytes}$

where $d_{head} = d_{model} / h$ and $n_{layers}$ is the number of transformer layers.

For a 70B parameter model with 64 heads and a 32K context window, using 16-bit precision, the KV cache alone can consume over 100GB of VRAM. This creates a severe limitation: the available context window is constrained by the GPU’s memory, not the model’s architecture. As $L$ grows, the memory footprint grows linearly, eventually forcing a reduction in batch size or context length.

The First Optimization: Grouped-Query Attention (GQA)

Grouped-Query Attention, proposed by Ainslie et al., "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints" (2023), addresses the KV cache size by sharing key and value heads across multiple query heads. Instead of $h$ separate key/value projections, GQA uses $g$ groups ($g < h$). Each group shares a single key and value vector for all $h/g$ query heads within that group.

Mechanistically, GQA computes:

1. $Q = X W_Q$ (standard projection)
2. $K, V = \text{reshape}(X W_K), \text{reshape}(X W_V)$ where the output is reshaped to $g$ heads.
3. Attention is computed between each query head and the shared key/value pairs for its group.

The KV cache size is reduced by a factor of $h/g$. If $h=64$ and $g=8$, the memory usage drops by 8x. This allows for significantly larger batch sizes or longer context windows on the same hardware.

However, GQA introduces a precision trade-off. By sharing keys and values, the model loses the ability to represent fine-grained attention patterns that distinct heads might capture. In practice, GQA with larger group counts (i.e., $g$ closer to $h$) can approach MHA performance, but extreme grouping (small $g$, such as $g=1$) can lead to degradation in perplexity and reasoning tasks. The optimal group size is often determined empirically, balancing memory savings against accuracy.

The Frontier: Multi-Head Latent Attention (MLA)

Multi-Head Latent Attention, introduced by DeepSeek-AI in "DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model" (2024), takes a different approach. Instead of reducing the number of heads, MLA compresses the dimensionality of the KV cache itself. The core insight is that the KV cache contains redundant information that can be compressed using low-rank factorization.

MLA works by:

1. Compression: The high-dimensional KV states are compressed into a low-dimensional "latent" vector $z$ using a projection matrix $W_D$.
2. Decoding: During attention computation, the latent vector $z$ is projected back to the full KV dimension using decoding matrices $W_R$ and $W_R^V$ directly within the attention kernel matmuls, avoiding intermediate memory materialization.

Mathematically, instead of storing $K \in \mathbb{R}^{L \times d_k}$, MLA stores $z \in \mathbb{R}^{L \times d_{latent}}$, where $d_{latent} \ll d_k$. The attention mechanism then effectively computes: $\text{Attention}(Q, K, V) \approx \text{Attention}(Q, z W_R, z W_R^V)$

This results in a massive reduction in KV cache size, often by an order of magnitude compared to MHA. Unlike GQA, which shares heads, MLA compresses the data itself, preserving the number of heads and their individual expressiveness. This allows MLA to maintain higher accuracy than GQA at equivalent memory footprints, particularly for long-context tasks.

The trade-off is increased computational complexity during the decoding step. The multiplication by $W_R$ and $W_R^V$ adds FLOPs to each attention step. However, since memory bandwidth is the primary bottleneck, this computational overhead is often acceptable if it significantly reduces memory transfers.

Comparative Analysis

Feature	MHA	GQA	MLA
KV Cache Size	High ($h \times d$)	Medium ($g \times d$)	Very Low ($d_{latent}$)
Compute Overhead	Low	Low	Moderate (Decoding matrices)
Accuracy	Highest	Good (depends on $g$)	High (preserves head diversity)
Implementation	Standard	Simple modification	Requires rank decomposition

MHA remains the gold standard for accuracy but is impractical for long contexts due to memory constraints. GQA offers a practical balance, widely adopted in open-source models like Llama 3 and Mistral, providing significant memory savings with minimal accuracy loss, a key goal of llm inference optimization. MLA represents the next step, offering aggressive compression while maintaining head diversity, making it ideal for ultra-long context applications.

Conclusion

The evolution from MHA to GQA to MLA reflects the industry’s shift from pure model capacity to system efficiency. For ML engineers, understanding these mechanisms is crucial for deploying LLMs at scale. GQA is likely to remain the default for general-purpose models due to its simplicity and effectiveness. MLA, however, is poised to become the standard for specialized long-context applications, where memory bandwidth is the ultimate limiting factor. As hardware evolves, the interplay between memory capacity and compute speed will continue to drive innovations in attention architectures.