Fast Transformer Decoding: One Write-Head is All You Need

The Problem

In standard Multi-Head Attention (MHA), during the autoregressive decoding phase of a Transformer, generating each new token requires loading the previously computed Keys (K) and Values (V) for all past tokens from memory. As sequences get longer and batch sizes increase, this "KV cache" grows massively. The process becomes bottlenecked by memory bandwidth—the time it takes to move these large tensors from memory to the compute units—rather than raw computing power.

The Idea

The paper proposes a simple architectural tweak to standard Multi-Head Attention: instead of each attention head having its own separate Key and Value projections, all attention heads share a single Key projection and a single Value projection. The queries remain separate for each head. This approach is called Multi-Query Attention (MQA).

How It Works

In a standard Transformer, if you have $H$ attention heads, you compute $H$ queries, $H$ keys, and $H$ values for each token.

Shared Projections: In MQA, you still compute $H$ distinct query vectors. However, you only compute $1$ key vector and $1$ value vector.

The Attention Calculation: During the attention step, each of the $H$ query heads computes its attention scores using the exact same shared key vector, and then aggregates the exact same shared value vectors.

Memory Reduction: Because there is only one set of keys and values, the size of the KV cache stored during decoding is reduced by a factor of $H$. This drastically reduces the memory bandwidth required to fetch the KV cache for every new token, speeding up generation significantly.

Why It Mattered

At the time of publication, the memory bottleneck of autoregressive generation wasn't the most pressing issue for most researchers, so MQA was initially overlooked. However, as Large Language Models (LLMs) scaled to billions of parameters and deployment became a massive challenge, memory bandwidth became the absolute limiting factor for serving speeds. MQA provided a crucial optimization that allowed LLMs to be served much faster and to much larger batches of users simultaneously, with minimal degradation in model quality.

What Came After

MQA became a standard feature in many major open-weight and proprietary models, most notably Google's PaLM. Later, a compromise between MQA and MHA emerged called Grouped-Query Attention (GQA), which shares keys and values among groups of heads rather than all heads. GQA was popularized by Llama 2, offering a sweet spot between the high quality of MHA and the extreme speed of MQA.