Attention Variants: MHA, MQA, GQA — Speed vs Quality at Scale

The attention mechanism is the heart of every transformer. It's the reason GPT can write poetry, Llama can code, and Gemini can understand images. But the original design — Multi-Head Attention (MHA) — has a dirty secret.

It's a memory hog.

For every token the model generates, MHA stores a key-value pair for every head in every layer. At 128K context, that's gigabytes of memory just for the cache. Not the model weights — just the conversation history.

So researchers asked: do we really need all those key-value pairs? The answer created two variants that every modern model now uses: MQA and GQA.

Quick Refresher: How Attention Works

If you need the full deep dive, revisit Episode 2. Here's the essential mechanic:

For each token, attention computes three vectors:

Query (Q): "What am I looking for?"
Key (K): "What do I contain?"
Value (V): "What information do I carry?"

Attention score = how well a Query matches each Key. The output = weighted sum of Values based on those scores.

Attention(Q, K, V) = softmax(QK^T / √d) × V

Multi-Head Attention (MHA)

Instead of one Q/K/V per token, MHA uses multiple heads — each looking at different aspects of the relationship:

Head 1: Q₁, K₁, V₁ → focuses on syntax ("subject-verb agreement")

Head 2: Q₂, K₂, V₂ → focuses on semantics ("related concepts")

Head 3: Q₃, K₃, V₃ → focuses on position ("nearby words")

...

Head 32: Q₃₂, K₃₂, V₃₂ → focuses on something else

Each head has its own Q, K, and V projection matrices. The outputs are concatenated and projected back.

For a model with h = 32 heads and d = 128 dimensions per head:

Each token generates:

• 32 Query vectors (32 × 128 = 4096 dims)
• 32 Key vectors (32 × 128 = 4096 dims)
• 32 Value vectors (32 × 128 = 4096 dims)

The KV Cache Problem

During inference, the model generates tokens one at a time. Each new token needs to attend to all previous tokens. To avoid recomputing K and V for every previous token, we cache them — the KV cache.

Token 1: compute K₁, V₁ → store in cache
Token 2: compute K₂, V₂ → store in cache, attend to [K₁,K₂], [V₁,V₂]
Token 3: compute K₃, V₃ → store in cache, attend to [K₁,K₂,K₃], [V₁,V₂,V₃]

...

Token 128000: attend to ALL 128K previous K,V pairs

How Big Is This Cache?

For Llama 3 70B (MHA, if it used it):

KV cache per token = 2 (K+V) × num_layers × num_heads × head_dim × bytes_per_param

= 2 × 80 × 64 × 128 × 2 (FP16)
                   = 2,621,440 bytes per token
                   = 2.5 MB per token

At 128K context:

= 2.5 MB × 128,000

= 320 GB just for KV cache!
That's 320 GB — more than the model weights themselves. For a single conversation. You literally can't serve it.

This is the problem MQA and GQA solve.

🔧

Multi-Query Attention (MQA)

Paper: "Fast Transformer Decoding" (Noam Shazeer, 2019) Used by: Falcon, PaLM, StarCoder

The idea is brutal in its simplicity: share one K and V across all query heads.

MHA (original):

Head 1: Q₁, K₁, V₁ ← each head has its own K, V

Head 2: Q₂, K₂, V₂

Head 3: Q₃, K₃, V₃

...

Head 32: Q₃₂, K₃₂, V₃₂

Total KV: 32 K's + 32 V's = 64 vectors per token

MQA:

Head 1: Q₁, K_shared, V_shared ← ALL heads share ONE K, V

Head 2: Q₂, K_shared, V_shared

Head 3: Q₃, K_shared, V_shared

...

Head 32: Q₃₂, K_shared, V_shared

Total KV: 1 K + 1 V = 2 vectors per token

The Impact

MQA gives you a 32x reduction in KV cache memory. That's the difference between "impossible to serve" and "fits on one GPU."

But there's a cost. Having only one K/V pair means all heads see the same keys and values — they lose the ability to attend to different aspects of the input. Quality drops, especially on complex reasoning tasks.

📊

Grouped-Query Attention (GQA)

Paper: "GQA: Training Generalized Multi-Query Attention" (Ainslie et al., 2023, Google) Used by: Llama 3, Mistral, Gemma 2, Phi-3, Qwen 2, DeepSeek V3

GQA is the Goldilocks solution: group the heads and share K/V within each group.

MHA (32 heads, 32 KV groups):

Group 1: Q₁, K₁, V₁

Group 2: Q₂, K₂, V₂

...

Group 32: Q₃₂, K₃₂, V₃₂

MQA (32 heads, 1 KV group):

Group 1: Q₁-Q₃₂, K_shared, V_shared

GQA (32 heads, 8 KV groups):

Group 1: Q₁-Q₄, K₁, V₁ ← 4 query heads share 1 KV

Group 2: Q₅-Q₈, K₂, V₂ ← 4 query heads share 1 KV

Group 3: Q₉-Q₁₂, K₃, V₃

...

Group 8: Q₂₉-Q₃₂, K₈, V₈

Total KV: 8 K's + 8 V's = 16 vectors per token

GQA: The Numbers

GQA-8 (8 groups for 32 query heads) gives you 4x memory reduction with virtually no quality loss. That's the sweet spot — and it's why every major model released in 2024-2025 uses GQA.

Visual Comparison

MHA (Multi-Head Attention) — Every head has its own K,V

┌───────────────────────────────────────────────┐

│ Q₁→K₁,V₁ Q₂→K₂,V₂ Q₃→K₃,V₃ Q₄→K₄,V₄ │

│ Q₅→K₅,V₅ Q₆→K₆,V₆ Q₇→K₇,V₇ Q₈→K₈,V₈ │

│ ... (32 KV pairs total) ... │

└───────────────────────────────────────────────┘

GQA (Grouped-Query) — Groups of heads share K,V

┌───────────────────────────────────────────────┐

│ Q₁ Q₂ Q₃ Q₄ → K₁,V₁ (Group 1) │

│ Q₅ Q₆ Q₇ Q₈ → K₂,V₂ (Group 2) │

│ ... (8 KV pairs total) ... │

└───────────────────────────────────────────────┘

MQA (Multi-Query) — ALL heads share one K,V

┌───────────────────────────────────────────────┐

│ Q₁ Q₂ Q₃ Q₄ Q₅ Q₆ ... Q₃₂ → K₁,V₁ │

│ (1 KV pair total) │

└───────────────────────────────────────────────┘

Real-World Model Configurations

Model Attention Query Heads KV Groups KV Heads Ratio

GPT-3 MHA 96 96 96 1:1

Notice: Llama 2's small model (7B) could afford MHA. The 70B model needed GQA. By Llama 3, even the 8B model uses GQA — because context lengths got longer (8K → 128K), making KV cache the bottleneck even for small models.

🔬

The Math: Why GQA Barely Hurts Quality

Intuitively, it seems like sharing K/V would lose information. But the research shows otherwise:

Many heads learn redundant patterns — when you analyze what different heads attend to, there's massive overlap

K/V contains less unique info than Q — the query is "what am I looking for" (unique per head), while K/V is "what's in this position" (more shared)

Training adapts — when you train a model from scratch with GQA, the remaining KV heads learn to be more general

The quality difference between MHA and GQA on benchmarks:

MMLU (knowledge): MHA: 79.5% GQA-8: 79.3% (−0.2%)

HumanEval (code):   MHA: 67.1%  GQA-8: 66.8%  (−0.3%)

GSM8K (math): MHA: 84.2% GQA-8: 83.9% (−0.3%)

Less than 0.5% drop for 4x memory savings. That's a no-brainer.

🛡️

Try It Yourself

python

import torch

import torch.nn as nn

class GroupedQueryAttention(nn.Module):

"""GQA implementation — the attention used in modern LLMs."""

def __init__(self, d_model=512, n_heads=8, n_kv_heads=2):
        super().__init__()
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.head_dim = d_model // n_heads
        self.n_rep = n_heads // n_kv_heads  # How many Q heads per KV group

self.wq = nn.Linear(d_model, n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(n_heads * self.head_dim, d_model, bias=False)

def forward(self, x):
        B, T, _ = x.shape

# Project Q, K, V
        q = self.wq(x).view(B, T, self.n_heads, self.head_dim)
        k = self.wk(x).view(B, T, self.n_kv_heads, self.head_dim)
        v = self.wv(x).view(B, T, self.n_kv_heads, self.head_dim)

# Repeat K, V to match number of Q heads
        # This is the GQA trick — expand KV to match Q heads
        k = k.repeat_interleave(self.n_rep, dim=2)  # (B, T, n_heads, head_dim)
        v = v.repeat_interleave(self.n_rep, dim=2)

# Standard attention
        q, k, v = [x.transpose(1, 2) for x in (q, k, v)]  # (B, n_heads, T, head_dim)
        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
        attn = torch.softmax(scores, dim=-1)
        out = torch.matmul(attn, v)

out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.wo(out)

Compare parameter counts

mha = GroupedQueryAttention(d_model=512, n_heads=8, n_kv_heads=8)  # MHA
gqa = GroupedQueryAttention(d_model=512, n_heads=8, n_kv_heads=2)  # GQA
mqa = GroupedQueryAttention(d_model=512, n_heads=8, n_kv_heads=1)  # MQA

for name, model in [("MHA", mha), ("GQA", gqa), ("MQA", mqa)]:

kv_params = sum(p.numel() for n, p in model.named_parameters() if 'wk' in n or 'wv' in n)
    print(f"{name}: KV parameters = {kv_params:,}")

MHA: KV parameters = 524,288

GQA: KV parameters = 131,072 (4x smaller)

MQA: KV parameters = 65,536 (8x smaller)

📦

Practical Takeaways

MHA gives every head its own KV — highest quality, highest memory, doesn't scale to long contexts

MQA shares one KV across all heads — 32x smaller cache, but measurable quality drop

GQA groups heads to share KV — the sweet spot, ~4x smaller cache, negligible quality loss

KV cache is the bottleneck at long contexts — at 128K tokens, it dwarfs the model weights

Every major 2024-2025 model uses GQA — Llama 3, Mistral, Gemma 2, Qwen 2, Phi-3

The ratio matters — 4:1 (Llama 3 8B) to 8:1 (Llama 3 70B) is the current sweet spot

🚀

What's Next?

Episode 75: FlashAttention — Even with GQA reducing KV cache size, the attention computation itself is a bottleneck. FlashAttention rewrites the algorithm to be IO-aware — understanding the GPU's memory hierarchy to run 2-4x faster while using 5-20x less memory. It's the single most impactful optimization in modern LLM inference.

← Previous

Ep 73: Model Formats

Next →

Ep 75: FlashAttention

Next: Episode 75 — FlashAttention

Standard attention wastes 90% of its time waiting for data to move between memory chips. Not computing. Waiting.

This is part of a 98-episode series covering AI engineering from tokens to production deployment.