DEV Community

Shrijith Venkatramana
Shrijith Venkatramana

Posted on

Multi-Head vs. Multi-Query vs. Grouped-Query Attention: The LLM Engineer's Guide to the KV Cache

Hello, I'm Shrijith Venkatramana. I'm building git-lrc, an AI code reviewer that runs on every commit. Star Us to help devs discover the project. Do give it a try and share your feedback for improving the product.


If you have ever wondered why an LLM can have billions of parameters yet inference becomes painfully constrained by memory bandwidth, attention is a good place to look.

The real bottleneck is often the amount of state the model must move through memory while generating one token at a time.

That observation led to a small but consequential evolution in Transformer architecture:

Multi-Head Attention (MHA) -> Multi-Query Attention (MQA) -> Grouped-Query Attention (GQA).

The three mechanisms perform essentially the same conceptual operation—queries look at keys and retrieve values—but they make very different choices about how many copies of keys and values need to exist.

1. First: what is a "head" actually doing?

The Transformer introduced by Vaswani et al. in 2017 replaced recurrence with attention. The paper, Attention Is All You Need, introduced the architecture that became the foundation for modern LLMs.

The basic attention operation is:

Attention(Q, K, V) = softmax((Q * K^T) / sqrt(d_k)) * V
Enter fullscreen mode Exit fullscreen mode

A useful developer interpretation is:

  • Q (query): "What information am I looking for?"
  • K (key): "What kind of information do I contain?"
  • V (value): "Here is the information itself."

Suppose the model is processing:

The server rejected the request because it lacked a valid certificate.

When processing it, attention can discover that server is a useful antecedent.

A single attention mechanism gives the model one learned way of performing this lookup.

Multi-head attention gives it several.

Instead of one Q, K, V, we have:

Q1, K1, V1
Q2, K2, V2
...
Qh, Kh, Vh
Enter fullscreen mode Exit fullscreen mode

and concatenate their outputs.

The original Transformer used 8 heads in its base model and 16 in the larger model. Each head operates in a learned subspace, allowing different heads to specialize in different relationships.

Think of it as having several specialists reading the same document:

  • one might learn syntactic relationships,
  • another positional relationships,
  • another semantic associations,
  • another long-range dependencies.

The model learns what those specialists should actually represent.

2. The catch: autoregressive generation changes the economics

During training, Transformers are extremely parallel.

Given a sequence of 4,000 tokens, the model can process essentially the whole sequence simultaneously.

Generation works differently.

Suppose you ask:

Explain how TCP congestion control works.

The model generates something like:

TCP
TCP congestion
TCP congestion control
TCP congestion control works
...
Enter fullscreen mode Exit fullscreen mode

At every generation step, the model needs to attend to everything it has already generated.

Recomputing the keys and values for all previous tokens would be extremely wasteful, so inference systems maintain a KV cache.

For every previous token, each Transformer layer stores its:

K and V
Enter fullscreen mode Exit fullscreen mode

Then, when generating the next token, the model computes the new query and attends against the cached keys and values.

This is where the engineering problem appears.

With conventional multi-head attention, every attention head gets its own K and V.

If there are 32 heads, you effectively maintain:

K1 V1
K2 V2
...
K32 V32
Enter fullscreen mode Exit fullscreen mode

for every token and every layer.

And the cache grows linearly with context length.

This is the crucial observation behind Noam Shazeer's 2019 paper Fast Transformer Decoding: One Write-Head is All You Need. Shazeer identified the memory-bandwidth cost of repeatedly loading the large K and V tensors as a major bottleneck in incremental decoding.

That observation produced the next architectural step.

3. Multi-Query Attention: keep all the questions, share the answers

Shazeer's idea was remarkably simple.

Keep multiple query heads:

Q1, Q2, ..., Qh
Enter fullscreen mode Exit fullscreen mode

and use one shared K and one shared V:

K, V
Enter fullscreen mode Exit fullscreen mode

So instead of:

Q1 -> K1,V1
Q2 -> K2,V2
Q3 -> K3,V3
Q4 -> K4,V4
Enter fullscreen mode Exit fullscreen mode

you have:

Q1 --+
Q2 --+
Q3 --+--> K,V
Q4 --+
Enter fullscreen mode Exit fullscreen mode

This is Multi-Query Attention (MQA).

The queries remain independent. The model still gets multiple different ways of asking questions.

All of those questions search the same representation of the previous tokens.

That has an enormous consequence for the KV cache.

Imagine:

  • 32 query heads
  • head dimension = 128
  • 32 Transformer layers
  • FP16
  • 8,192-token context

With ordinary MHA, the KV cache is approximately:

2 * 32 * 32 * 128 * 2 * 8192 bytes
Enter fullscreen mode Exit fullscreen mode

That's roughly 4 GB per sequence.

With MQA, there is only one KV head:

2 * 32 * 1 * 128 * 2 * 8192 bytes
Enter fullscreen mode Exit fullscreen mode

which is approximately 128 MB.

Same model depth. Same query heads. Same context.

The KV cache has shrunk by 32x.

That changes the economics of serving.

During autoregressive decoding, the historical K/V tensors are repeatedly loaded. Reducing their size reduces memory traffic. Shazeer's experiments demonstrated substantially faster decoding with only modest quality degradation relative to conventional multi-head attention.

4. MQA's trade-off: sharing the representation of the past

In MHA, each head has:

K1,V1,K2,V2,...,Kh,Vh
Enter fullscreen mode Exit fullscreen mode

Those different K/V projections can encode different representations.

MQA collapses all of them into:

K,V
Enter fullscreen mode Exit fullscreen mode

The queries remain diverse, while the information they query is shared.

A useful mental model is:

MHA gives every specialist their own database.

MQA gives all specialists one database.

The shared database is dramatically cheaper to keep in memory. The specialists also have less representational independence.

The engineering question becomes:

How much representational independence are we willing to trade for memory efficiency?

That question led naturally to the third architecture.

5. GQA: the compromise that turned out to be extremely useful

In 2023, Joshua Ainslie and colleagues at Google introduced Grouped-Query Attention (GQA) in their EMNLP paper GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.

The idea is almost embarrassingly intuitive once MHA and MQA are on the table.

Consider three configurations:

32 query heads
32 KV heads     <- MHA
Enter fullscreen mode Exit fullscreen mode
32 query heads
1 KV head      <- MQA
Enter fullscreen mode Exit fullscreen mode

and:

32 query heads
8 KV heads      <- GQA
Enter fullscreen mode Exit fullscreen mode

Each KV head serves a group of query heads.

For example:

Q0 Q1 Q2 Q3       --> K0,V0

Q4 Q5 Q6 Q7       --> K1,V1

Q8 Q9 Q10 Q11     --> K2,V2

...
Enter fullscreen mode Exit fullscreen mode

Now the model has:

  • 32 independent query heads
  • 8 KV heads
  • each KV head shared by 4 query heads

GQA is therefore a continuum between MHA and MQA.

The number of KV heads becomes an architectural knob.

Ainslie et al. explicitly describe GQA as an interpolation between multi-head and multi-query attention, with a single K/V head for each subgroup of query heads.

This gives model designers a useful spectrum:

MHA                         MQA
 |---------------------------|
32 KV heads                1 KV head
        ^
        |
       GQA
Enter fullscreen mode Exit fullscreen mode

6. The really important number: KV-cache size

For inference engineers, this is probably the most useful formula in the entire discussion.

Ignoring minor implementation details, the KV cache per sequence is approximately:

Memory =
    2 * L * n_KV * d_head * T * bytes
Enter fullscreen mode Exit fullscreen mode

where:

  • 2 = K and V
  • L = number of Transformer layers
  • n_KV = number of KV heads
  • d_head = head dimension
  • T = number of cached tokens
  • bytes = bytes per element

Notice the important variable:

n_KV
Enter fullscreen mode Exit fullscreen mode

The number of query heads does not directly determine KV-cache size.

That is the architectural trick.

Consider a hypothetical model with:

32 query heads
128-dimensional heads
32 layers
FP16
8192-token context
Enter fullscreen mode Exit fullscreen mode

The KV-cache comparison becomes:

Configuration Q heads KV heads Relative KV cache
MHA 32 32 1x
GQA 32 8 0.25x
MQA 32 1 0.031x

At 8K context:

MHA   ~= 4.0 GB
GQA   ~= 1.0 GB
MQA   ~= 0.125 GB
Enter fullscreen mode Exit fullscreen mode

The arithmetic is straightforward.

For MHA:

2 * 32 layers * 32 KV heads * 128 dimensions
* 8192 tokens * 2 bytes
~= 4.0 GB
Enter fullscreen mode Exit fullscreen mode

For GQA:

2 * 32 * 8 * 128 * 8192 * 2
~= 1.0 GB
Enter fullscreen mode Exit fullscreen mode

For MQA:

2 * 32 * 1 * 128 * 8192 * 2
~= 0.125 GB
Enter fullscreen mode Exit fullscreen mode

The cache scales linearly with sequence length:

8K   -> 4 GB
16K  -> 8 GB
32K  -> 16 GB
Enter fullscreen mode Exit fullscreen mode

for the hypothetical MHA configuration.

This is where attention architecture becomes an operations problem.

Suppose your GPU has 80 GB of usable memory.

Your model weights might consume 60 GB.

That leaves roughly 20 GB for:

  • KV cache
  • activations
  • CUDA workspace
  • batching overhead
  • runtime allocations

The attention architecture therefore influences how many simultaneous sequences the GPU can accommodate.

It also affects how much data the GPU needs to read during every decoding step.

That is why LLM serving can become constrained by memory capacity and bandwidth alongside compute.

7. Why GQA became the practical middle ground

Ainslie et al. were motivated by a practical problem.

MQA made decoding cheaper, while its quality trade-offs could matter. Training an entirely new model solely to obtain MQA inference characteristics also represented a substantial cost.

Their 2023 paper demonstrated a way to uptrain existing MHA checkpoints into MQA or GQA models using approximately 5% of the original pre-training compute. They found that GQA could achieve quality close to MHA while delivering inference speed comparable to MQA.

This is an important detail in the history.

The progression followed a practical engineering path:

Build a powerful architecture -> discover an inference bottleneck -> remove expensive redundancy -> measure the quality trade-off -> introduce a tunable sharing scheme.

The evolution looks like this:

2017
Transformer
   |
   |  Multiple independent K/V projections
   v
MHA
   |
   |  "Why are we carrying so many K/V tensors?"
   v
2019
MQA
   |
   |  "One shared KV representation can affect quality."
   v
2023
GQA
   |
   |  "Let's share K/V within groups."
   v
Modern LLM serving
Enter fullscreen mode Exit fullscreen mode

There is a broader systems lesson here.

Good systems engineering often consists of finding expensive redundancy and deciding how much of it you can safely remove.

Attention provides a particularly clean example because the redundancy has a direct relationship to GPU memory traffic.

8. The developer's mental model

If you are implementing or operating an LLM, I would remember the three architectures this way.

MHA — maximum independence

Q0 -> K0,V0
Q1 -> K1,V1
Q2 -> K2,V2
...
Q31 -> K31,V31
Enter fullscreen mode Exit fullscreen mode

Mental model: every attention head has its own memory.

You pay the most in KV-cache memory.


MQA — maximum sharing

Q0 --+
Q1 --+
Q2 --+--> K,V
... |
Q31-+
Enter fullscreen mode Exit fullscreen mode

Mental model: all attention heads share one memory.

KV-cache efficiency is excellent, with the strongest sharing of K/V representations.


GQA — grouped sharing

Q0 Q1 Q2 Q3       -> K0,V0
Q4 Q5 Q6 Q7       -> K1,V1
Q8 Q9 Q10 Q11     -> K2,V2
...
Enter fullscreen mode Exit fullscreen mode

Mental model: teams of attention heads share memory.

You choose the team size.

If:

n_Q = 32
n_KV = 8
Enter fullscreen mode Exit fullscreen mode

then each KV head serves:

32 / 8 = 4
Enter fullscreen mode Exit fullscreen mode

query heads.

That ratio:

n_Q / n_KV
Enter fullscreen mode Exit fullscreen mode

is a useful number to keep in your head when reading modern LLM architectures.

A model architecture that says:

Attention heads: 32
KV heads:        8
Enter fullscreen mode Exit fullscreen mode

immediately tells you that the model is using GQA with four query heads per KV group.

9. The deeper lesson: inference architecture is economics

There is a tendency to think about neural-network architecture in terms of accuracy and FLOPs.

For training, those metrics are fundamental.

For production inference, the system has another layer of economics:

Model
  |
GPU memory
  |
KV cache
  |
Memory bandwidth
  |
Batch size
  |
Tokens/sec
  |
Users/GPU
  |
$/million tokens
Enter fullscreen mode Exit fullscreen mode

Changing MHA to GQA can therefore affect:

  • maximum context you can fit,
  • maximum concurrent sequences,
  • batching efficiency,
  • tokens/sec,
  • GPU utilization,
  • number of GPUs required,
  • latency,
  • and ultimately cost per generated token.

Imagine two serving configurations with identical model weights.

Configuration A uses MHA:

32 KV heads
Enter fullscreen mode Exit fullscreen mode

Configuration B uses GQA:

8 KV heads
Enter fullscreen mode Exit fullscreen mode

At the same context length, Configuration B needs roughly one quarter of the KV-cache storage.

That can create room for more concurrent requests.

More concurrent requests can improve batch utilization.

Higher utilization can improve the economics of the GPU.

A model architecture choice has therefore propagated all the way into infrastructure cost.

This is also why Shazeer's 2019 paper is interesting historically. The paper framed MQA around the operational reality of incremental decoding: the model repeatedly loads K/V state, and memory bandwidth becomes a limiting resource.

The original Transformer demonstrated how attention could replace recurrence and make sequence modeling dramatically more parallel during training.

The subsequent work on MQA and GQA shows another stage of the story: once these models entered large-scale inference, the physical movement of model state became an architectural concern.

So when you encounter an LLM architecture diagram saying:

32 attention heads, 8 KV heads

read it as a systems decision:

The designers decided that four query heads can economically share one representation of the past.

That single line in a model card tells you something about the model's memory behavior, inference characteristics, and architectural trade-offs.

Conclusion: attention heads are also a memory architecture

The progression from MHA to MQA to GQA is a beautiful example of how modern ML systems evolve.

MHA: give every head its own K/V representation.

MQA: share K/V across all query heads and dramatically reduce inference memory traffic.

GQA: share K/V selectively and preserve more representational independence.

The attention equation remains essentially the same.

The architectural decision concerns who gets to own the keys and values.

Once you see that, the three mechanisms become much easier to reason about:

MHA -> maximum independence

MQA -> maximum sharing

GQA -> controlled sharing
Enter fullscreen mode Exit fullscreen mode

The fascinating question for LLM engineers is:

How much representational independence are you willing to buy when every byte of KV cache has to be stored, moved, and paid for?

And perhaps the more practical question is:

When you choose an LLM for production, how often do you look at its KV-head configuration?


*AI agents write code fast. They also silently remove logic, change behavior, and introduce bugs -- without telling you. You often find out in production.

git-lrc fixes this. It hooks into git commit and reviews every diff before it lands. 60-second setup. Completely free.*

Any feedback or contributors are welcome! It's online, source-available, and ready for anyone to use.

GitHub logo HexmosTech / git-lrc

Free, Micro AI Code Reviews That Run on Git Commit




GenAI today is a race car without brakes. It accelerates fast -- you describe something, and large blocks of code appear instantly. But AI agents silently break things: they remove logic, relax constraints, introduce expensive cloud calls, leak credentials, and change behavior -- without telling you. You often find out in production.

git-lrc is your braking system. It hooks into git commit and runs an AI review on every diff before it lands. 60-second setup. Completely free.

In short, git-lrc helps Prevent Outages, Breaches, and Technical Debt Before They Happen

At a glance: 10 risk categories · 100+ failure patterns tracked · every commit…

Top comments (0)