MHA vs MQA vs GQA: How Modern LLMs Shrink the KV Cache

Multi-head attention gave every head its own view of a sentence. In the original design, every query head also has its own key and value head. That is called multi-head attention, or MHA. It works well, but its KV cache can be large.

Query heads can remain separate even when some K and V heads are shared.

MHA: One KV Head Per Query Head

With 8 query heads, MHA has 8 key heads and 8 value heads. Query head 0 reads key and value head 0. Query head 1 reads head 1, and so on. There is no sharing between heads.

For one request with three cached positions and head size 128, the key tensor has shape $[1,8,3,128]$. The value tensor has the same shape. The head axis contains eight different learned views of the same three positions.

Each view uses a different projection matrix. For a query in head 6, matching keys come from key head 6, not all eight heads at once. The weighted result uses value head 6. Other heads perform their own matching and mixing, and the attention output projection combines their results afterward.

That head-to-head pairing explains the storage count. Eight heads do not mean eight separate token sequences. They mean eight sets of features for each position in one sequence.

MHA has the largest KV cache of the three designs when query-head count is fixed.

If our model has 8 query heads and uses MHA, then H_kv = 8 in the Day 22 cache formula. The cache stores 8 sets of keys and values at every layer and token position.

MQA: One Shared KV Head

Multi-query attention, or MQA, keeps many query heads but uses only one shared key head and one shared value head. All query heads compare against the same keys and mix the same values, though their queries remain different.

In the same three-position example, K and V now have shape $[1,1,3,128]$, while a newest query has shape $[1,8,1,128]$. The eight query vectors remain distinct. A compatible implementation maps each of them to the single saved key/value head.

Suppose query head 0 focuses on the first token and head 1 focuses on the last. Both can still do that with shared K and V because their different queries create different score rows. But they cannot request independently projected key or value features: that freedom was removed by sharing.

This distinction separates MQA from using just one attention head. A one-head model would also have only one query view. MQA keeps several query views and their separate output mixtures.

Many questions can read the same stored reference cards.

For 8 query heads, MQA makes H_kv = 1. Compared with MHA, the KV-cache part becomes 8 times smaller. This can greatly reduce memory bandwidth during decoding.

The eightfold cache reduction assumes the same number of layers, cached positions, feature width, and bytes per entry. If an architecture changes any of those alongside the head count, count the new arrays instead of quoting eightfold as a universal result.

The saved bytes matter most when storage or reads of the cache are a meaningful part of the workload. A short prompt may spend more time elsewhere in the model. A lower cache count is an exact shape result; a lower response time is a workload-dependent measurement.

Why Not Always Use MQA?

Sharing all K and V information is a strong constraint. Different query heads may benefit from separate key-value views. MQA saves memory, but a model designer must check whether the quality change is acceptable for the model and training setup.

We therefore want a middle choice between 8 KV heads and one KV head.

GQA: Share Within Groups

Grouped-query attention, or GQA, divides query heads into groups. Every group shares one key head and one value head. With 8 query heads and 2 KV heads, query heads 0 to 3 can use KV head 0, while query heads 4 to 7 use KV head 1.

GQA keeps more than one KV view without giving every query head a private copy.

This design has H_q = 8 and H_kv = 2. The cache has two KV-head rows instead of eight, so its KV portion is four times smaller than MHA.

See the Cache Formula Change

Only one factor changes in the Day 22 formula: H_kv.

For 8 query heads, the relevant counts are MHA H_kv=8, GQA H_kv=2 for groups of four, and MQA H_kv=1.

For the same three cached positions, GQA stores two KV views, each covering all three tokens. It does not split the sequence so that one group sees the first half and the other sees the second half. Every query group can still use the whole allowed prefix.

Keep these axes separate when interpreting a diagram. The head group selects which representation to use. The position axis selects which past token to read. Mixing the two would change the attention computation, not merely its storage.

Query-head count stays at 8 while cache-head count changes from 8 to 2 to 1.

If every other quantity stays fixed, GQA with 2 KV heads uses one quarter of MHA's KV-cache memory. MQA uses one eighth.

What Happens at Attention Time?

The query tensor still has one row per query head. Before the dot product, the selected KV head is used for all query heads in its group. This is sharing, not averaging. Each query head still forms its own scores because its query vector is different.

For eight queries and two KV heads, query heads 0, 1, 2, and 3 share the first K/V view. Query heads 4, 5, 6, and 7 share the second. All eight heads run softmax separately for each query position.

We do not average the four query vectors before attention. Doing that would produce one score row for the group rather than four. We also do not average their output vectors before the normal concatenation step. The model still needs the separate head outputs to preserve those different mixtures.

Shared keys do not force identical attention weights.

This is a useful distinction. GQA reduces stored K and V rows. It does not reduce the number of query heads that produce output features.

Shared Keys Can Still Give Different Weights

Take one shared two-feature key set: $k_0=[1,0]$ and $k_1=[0,1]$. Query head A has $q_A=[2,0]$, while query head B has $q_B=[0,2]$. Their raw score rows are $[2,0]$ and $[0,2]$.

After scaling by $\sqrt{2}$ and applying softmax, the weight rows are approximately $[0.80443,0.19557]$ and $[0.19557,0.80443]$. If the shared values are $[10,0]$ and $[0,10]$, the outputs are approximately $[8.0443,1.9557]$ and $[1.9557,8.0443]$.

Both heads read identical stored keys and values. Their different queries produce different mixtures. What sharing removes is the ability to give those heads independently learned KV projections, not the ability to ask different questions.

A Concrete 8-Head Comparison

For one token and one layer with head size 128 in bfloat16, MHA stores 2 x 8 x 128 x 2 = 4,096 bytes. GQA with 2 KV heads stores 1,024 bytes. MQA with 1 KV head stores 512 bytes.

The stored K and V data drops from 4,096 bytes to 1,024 bytes to 512 bytes.

The savings repeat for every cached token, layer, and active request. That is why modern decoder design pays close attention to this seemingly small head-count choice.

Here is the storage arithmetic in plain Python. The cache datatype uses two bytes per number; one key and one value create the other factor of two.

query_heads = 8
head_size = 128
bytes_per_number = 2
kv_head_counts = {"MHA": 8, "GQA": 2, "MQA": 1}
bytes_per_layer_token = {
    name: 2 * count * head_size * bytes_per_number
    for name, count in kv_head_counts.items()
}

The dictionary contains 4,096, 1,024, and 512 bytes. This calculation holds query heads and feature width fixed. It is a cache-only comparison, not a prediction that total model memory or generation time falls by exactly those ratios.

All eight query heads still produce attention outputs. Some projection work shrinks with KV sharing, but the model still has attention scores for eight query heads and the rest of its Transformer blocks. End-to-end speed needs measurement.

The PyTorch Shape Clue

In standard MHA, Q, K, and V commonly have the same head axis length. In GQA, Q has more heads than K and V. A suitable attention implementation knows how to repeat or map the smaller KV head set across query groups without asking us to create a full duplicated cache.

The smaller K/V tensor is the source of memory savings.

Do not permanently duplicate the compact cache into eight full KV heads merely to match Q's shape. That would lose the intended persistent storage saving. A compatible kernel can map each query head to the right stored KV group.

For eight queries and two KV heads, the group size is four. A zero-based query head index maps to query_head // 4, so head 3 uses KV head 0 and head 4 uses KV head 1. Ordinary equal-size grouping requires the query-head count to be divisible by the KV-head count.

Changing a trained model from MHA to GQA is not just changing this integer. Its projection weights and training must support the new arrangement. The GQA paper studies a method for converting and further training models; it does not claim arbitrary head sharing is lossless.

To compare model quality fairly, hold the evaluation tasks and prompting procedure fixed and examine the trained models, not only their cache dimensions. A smaller cache can improve serving capacity while changing some answer qualities. Neither outcome follows from the memory formula alone.

A useful teach-back exercise is to draw Q with eight heads and K/V with two, then trace query head 5 through the whole operation. It must choose KV head 1, score every allowed token in that head, normalize its own score row, mix that head's values, and occupy its own slot in the concatenated output.

Recap

This is how MHA, MQA, and GQA change a cache. MHA gives every query head its own key and value head. MQA gives all query heads one shared KV head. GQA groups query heads so each group shares a KV head.

The cache formula depends on KV-head count, not simply query-head count. Day 27 goes further by storing a smaller learned latent representation instead of ordinary per-head keys and values.