MHA → GQA → MQA — why modern LLMs share their KV cache

With ordinary Multi-Head Attention (MHA), each attention head gets its own Q, K, and V projections.

Suppose there are 4 heads:

Head 1: Q₁ K₁ V₁
Head 2: Q₂ K₂ V₂
Head 3: Q₃ K₃ V₃
Head 4: Q₄ K₄ V₄

Every previous token therefore contributes four K vectors and four V vectors to the cache.

Why might that be wasteful?¶

Imagine the heads have learned different query behaviors:

Head 1 → people
Head 2 → grammatical subjects
Head 3 → objects
Head 4 → ownership relationships

The queries can be different without necessarily needing four completely separate stored memories.

Grouped Query Attention¶

Keep four Q heads but use two KV heads:

Q₁ ─┐
    ├── KV-A
Q₂ ─┘

Q₃ ─┐
    ├── KV-B
Q₄ ─┘

Even though Q₁ and Q₂ share K/V, their attention patterns can differ because:

\[ Q_1K_A^T\neq Q_2K_A^T \]

The key insight is:

Sharing the searchable memory does not force queries to ask the same question.

Memory calculation¶

Return to: - 32 layers - 2048 tokens - head dimension 128 - BF16 - 32 query heads

With MHA: \(H_{KV}=32\)

\[ 2\times32\times2048\times128 \]

values per layer for K+V, about 33.6 MB/layer.

Across 32 layers:

\[ \approx1.07\text{ GB} \]

With GQA using only 8 KV heads:

\[ 2\times8\times2048\times128 \]

or about 8.4 MB/layer.

Across 32 layers:

\[ \approx268\text{ MB} \]

A 4× reduction.

Multi-Query Attention¶

MQA pushes sharing to the extreme:

\[ N_Q=32,\quad N_{KV}=1 \]

All query heads share one K and one V head.

That can slash KV-cache size by ~32× relative to 32-head MHA, though it constrains representation capacity more aggressively.

Beyond capacity: bandwidth¶

During decode, the model must read historical K/V data. Fewer KV heads means less HBM traffic as well as less memory occupancy.

That matters because decode is often bandwidth-bound.

Mental model¶

Think of each Q head as a researcher and K/V as a database.

MHA: every researcher keeps a private database copy.
MQA: everyone shares one database.
GQA: researchers are organized into teams, with one shared database per team.

The researchers can still ask different questions because their queries remain independent.

Why query heads are more valuable than duplicated KV heads¶

A query head controls how the current token asks for information. Different query heads can look for syntax, entity relationships, recency, delimiters, or other patterns.

The stored K/V cache is more like the representation of the historical context. Keeping 32 independent ways of querying that history may be more valuable than keeping 32 full copies of the history representation.

This is the intuition behind the asymmetry:

\[ H_Q > H_{KV} \]

in GQA.

Decode bandwidth calculation¶

Suppose the current context contains \(T=32,000\) tokens, the model has 32 layers, head dimension 128, and BF16 K/V values.

With 32 KV heads, the cache is roughly:

\[ 2\times32\times32000\times32\times128\times2 \]

bytes, about 16.8 GB for one sequence in this simplified setup.

With 8 KV heads:

\[ \approx4.2GB \]

That is not only 12.6 GB less capacity; each decode step also has fewer K/V bytes to scan.

The trade-off is representational freedom. With MHA, each attention head can build its own key/value representation of history. With GQA, several query heads must interpret a shared representation. With MQA, all query heads share one.

Empirically, a moderate amount of sharing often preserves much of model quality while making serving dramatically cheaper. Hence GQA has become a common design point.

The general lesson is recurring throughout LLM engineering:

\[ \boxed{\text{remove redundancy where it costs more than the extra flexibility is worth}} \]

A more exact GQA tensor picture¶

Suppose:

\[ H_Q=32,\qquad H_{KV}=8,\qquad d_{head}=128. \]

Then a prefill query tensor might be:

\[ Q:[B,32,T,128] \]

while:

\[ K,V:[B,8,T,128]. \]

Every block of four query heads shares one K/V head. The implementation maps query head \(h\) to a KV head such as:

\[ \left\lfloor h/4\right\rfloor. \]

The attention score computation still happens independently for every query head. What is shared is the historical representation being searched.

This explains why the quality loss from GQA can be much smaller than the memory reduction suggests: we are not collapsing 32 attention heads into eight attention heads. We still have 32 independently learned queries and 32 independently produced attention distributions.

During decode this distinction is especially valuable. The current token creates 32 query vectors, but the server has to fetch only eight sets of historical keys and values per layer.

So GQA attacks exactly the tensor dimension that grows persistently with context length.