Multi-Query Attention (MQA) & Grouped Query Attention (GQA)

Overview

Multi-Query Attention (MQA) and Grouped Query Attention (GQA) are architectural variants that reduce the size of Key-Value (KV) cache by using fewer K/V heads than Query heads, achieving 8-32x KV cache reduction while maintaining model quality.

• MQA Paper: "Fast Transformer Decoding: One Write-Head is All You Need" (Shazeer, 2019)
• GQA Paper: "GQA: Training Generalized Multi-Query Transformers" (Ainslie et al., 2023)
• Key Innovation: Separate number of K/V heads from Q heads
• Impact: 8-32x KV cache reduction, up to 4x faster inference
• Adoption: Llama 2 70B (MQA), Mistral 7B (MQA), Falcon (MQA)

---

The Problem: KV Cache Scaling

Standard Multi-Head Attention

Input shape: (batch, seq_len, hidden_dim)
Num heads: 32
Head dim: 128

Standard Attention:
  - Q: (batch, seq_len, 32, 128)
  - K: (batch, seq_len, 32, 128)  ← KV cache stores this
  - V: (batch, seq_len, 32, 128)  ← KV cache stores this

KV Cache for one token:
  - Per token: 32 heads × 128 dim × 2 (K+V) = 8,192 values
  - For 4096 tokens: 4096 × 8,192 = 33.5M values
  - In float16: 33.5M × 2 bytes = 67MB per token
  - Problem: KV cache grows with sequence length!

For 32 concurrent requests (batch):
  - 32 × 67MB × 4096 tokens ≈ 8GB KV cache
  - Huge memory footprint!

Why This is a Problem

GPU Memory Breakdown (Llama 2 70B inference):

Model weights:     140GB (quantized to 35GB)
KV cache (batch):  8GB (with 32 requests @ 4K tokens)
Activations:       2GB
────────────────────────────────────
Total:             45GB

Solution needed: Reduce KV cache from 8GB!

Observation:
"Do we really need 32 heads for Keys and Values?"
  - Q heads: Generate 32 different query perspectives
  - K/V heads: Store 32 different key/value perspectives
  - Intuition: Maybe fewer K/V heads are sufficient

---

Multi-Query Attention (MQA): The Radical Approach

Core Concept: One K/V Head for All Queries

Standard Multi-Head Attention:
- ┌─────────────────────────────────────────┐
    - 32 Query Heads                          │
    - Head 1: Q₁ @ K₁^T
    - Head 2: Q₂ @ K₂^T
    - ...
    - Head 32: Q₃₂ @ K₃₂^T
  - 32 Key/Value Heads (separate!)
  - ┘

MQA (Multi-Query Attention):
- ┌─────────────────────────────────────────┐
    - 32 Query Heads                          │
    - Head 1: Q₁ @ K^T (shared!)
    - Head 2: Q₂ @ K^T (shared!)
    - ...
    - Head 32: Q₃₂ @ K^T (shared!)
  - 1 Key/Value Head (shared by all!)
  - ┘

Key insight:
All query heads attend to the SAME K/V!
Massive reduction in cache size!

Memory Savings

MQA vs Standard Attention (Llama 2 7B):

                    Standard    MQA         Reduction
────────────────────────────────────────────────────
Q heads            32          32          0%
K heads            32          1           97%
V heads            32          1           97%
KV cache size      1GB         31MB        97%
(per request)

For batch of 32:
  - Standard: 32 × 1GB = 32GB
  - MQA: 32 × 31MB ≈ 1GB
  - Savings: 31x reduction! 🎉

How It Works: Mathematical View

Standard Attention:
For each query head i:
  S_i = Q_i @ K_i^T / sqrt(d)      (N × N matrix)
  P_i = softmax(S_i)
  O_i = P_i @ V_i

MQA:
For each query head i:
  S_i = Q_i @ K^T / sqrt(d)        (N × N matrix, K is shared!)
  P_i = softmax(S_i)
  O_i = P_i @ V                    (V is shared!)

Difference: K and V are NOT head-specific, shared across all heads!

Implementation

class MultiQueryAttention(nn.Module):
    def __init__(self, num_q_heads, hidden_dim):
        super().__init__()
        self.num_q_heads = num_q_heads  # e.g., 32
        self.head_dim = hidden_dim // num_q_heads

        # Standard multi-head
        self.q_proj = nn.Linear(hidden_dim, num_q_heads * self.head_dim)
        self.k_proj = nn.Linear(hidden_dim, self.head_dim)      # Only 1 head!
        self.v_proj = nn.Linear(hidden_dim, self.head_dim)      # Only 1 head!
        self.o_proj = nn.Linear(num_q_heads * self.head_dim, hidden_dim)

    def forward(self, x, kv_cache=None):
        batch, seq_len, hidden_dim = x.shape

        # Project Q, K, V
        q = self.q_proj(x).view(batch, seq_len, self.num_q_heads, self.head_dim)
        k = self.k_proj(x).view(batch, seq_len, 1, self.head_dim)        # 1 head
        v = self.v_proj(x).view(batch, seq_len, 1, self.head_dim)        # 1 head

        # Expand k, v to match query heads (broadcasting)
        k = k.expand(-1, -1, self.num_q_heads, -1)
        v = v.expand(-1, -1, self.num_q_heads, -1)

        # Standard attention computation
        scores = q @ k.transpose(-2, -1) / sqrt(self.head_dim)
        weights = softmax(scores, dim=-1)
        output = weights @ v

        # Merge heads
        output = output.view(batch, seq_len, -1)
        output = self.o_proj(output)

        return output

# Memory benefit:
# Standard: K has (batch, seq_len, 32, 128) → cache size = seq_len × 32 × 128
# MQA:      K has (batch, seq_len, 1, 128)  → cache size = seq_len × 1 × 128
# Savings:  32x reduction! ✓

---

Grouped Query Attention (GQA): The Balanced Approach

Problem with MQA

MQA Pros:
✅ 32x KV cache reduction
✅ Much faster inference
✅ Massive memory savings

MQA Cons:
❌ Quality drops by 5-10% (significant!)
❌ Information bottleneck: all heads share single K/V
❌ Less expressiveness than standard attention

Trade-off: Need more cache, less reduction, but better quality

GQA: Middle Ground

Standard Attention:    32 K/V heads (no reduction)
MQA:                   1 K/V head (32x reduction, quality loss)
GQA:                   8 K/V heads (4x reduction, minimal loss)

Grouped Query Attention:
- ┌─────────────────────────────────────────┐
    - 32 Query Heads                          │
    - Heads 1-4:   Query group 1 → K₁, V₁
    - Heads 5-8:   Query group 2 → K₂, V₂
    - ...
    - Heads 29-32: Query group 8 → K₈, V₈
  - 8 Key/Value Heads (grouped!)
  - ┘

Key insight:
Group query heads, each group shares K/V!
Multiple groups allow expressiveness while reducing cache.

Memory and Accuracy Trade-off

Comparison: Llama 2 7B

                    Standard    GQA         MQA
────────────────────────────────────────────────
K/V heads           32          8           1
KV cache (per req)  1GB         256MB       31MB
Reduction vs Std    1x          4x          32x
Quality loss        0%          0.5-1%      5-10%

MMLU Benchmark:
Standard:           45.9%       (baseline)
GQA:                45.7%       (-0.2%, negligible!)
MQA:                43.9%       (-2.0%, significant!)

Recommendation:
GQA is sweet spot: 4x reduction with minimal quality loss!

Implementation

class GroupedQueryAttention(nn.Module):
    def __init__(self, num_q_heads, num_kv_heads, hidden_dim):
        super().__init__()
        self.num_q_heads = num_q_heads        # 32
        self.num_kv_heads = num_kv_heads      # 8 (or 4)
        self.head_dim = hidden_dim // num_q_heads
        self.num_groups = num_q_heads // num_kv_heads  # 4

        # Projections
        self.q_proj = nn.Linear(hidden_dim, num_q_heads * self.head_dim)
        self.k_proj = nn.Linear(hidden_dim, num_kv_heads * self.head_dim)  # Fewer heads
        self.v_proj = nn.Linear(hidden_dim, num_kv_heads * self.head_dim)  # Fewer heads
        self.o_proj = nn.Linear(num_q_heads * self.head_dim, hidden_dim)

    def forward(self, x):
        batch, seq_len, hidden_dim = x.shape

        # Project
        q = self.q_proj(x).view(batch, seq_len, self.num_q_heads, self.head_dim)
        k = self.k_proj(x).view(batch, seq_len, self.num_kv_heads, self.head_dim)
        v = self.v_proj(x).view(batch, seq_len, self.num_kv_heads, self.head_dim)

        # Repeat k, v to match query heads (group-wise expansion)
        # k shape: (batch, seq_len, 8, 128) → (batch, seq_len, 32, 128)
        k = k.repeat_interleave(self.num_groups, dim=2)
        v = v.repeat_interleave(self.num_groups, dim=2)

        # Standard attention
        scores = q @ k.transpose(-2, -1) / sqrt(self.head_dim)
        weights = softmax(scores, dim=-1)
        output = weights @ v

        # Merge heads
        output = output.view(batch, seq_len, -1)
        output = self.o_proj(output)

        return output

# Memory benefit:
# Standard: K has 32 heads → cache = seq_len × 32 × 128
# GQA:      K has 8 heads  → cache = seq_len × 8 × 128
# Savings:  4x reduction! ✓ (vs 32x for MQA but better quality)

---

Adoption in Modern Models

Models Using MQA:
  - Llama 2 70B: MQA (32→1 heads)
  - Mistral 7B: MQA (32→1 heads)
  - Falcon 40B: MQA
  - Phi-2: MQA

Models Using GQA:
  - Llama 3: GQA (48→8 heads)
  - Qwen: GQA variants
  - Newer models: Increasingly adopting GQA

Why the shift from MQA to GQA:
  - MQA had quality concerns in large models
  - GQA provides better quality/efficiency trade-off
  - Slight increase in memory vs huge quality improvement
  - Industry consensus: GQA is the future

---

Performance Impact

Inference Speedup

Benchmark: Llama 2 7B inference

Metric              Standard    GQA         MQA
──────────────────────────────────────────────
Batch size possible 8           32          32
KV cache memory     1GB         256MB       31MB
Max concurrent      8 req       32 req      32 req
Tokens/sec          200         750         800
(per GPU)

Effective throughput (with constraints):
  - Standard: Limited by memory to 8 concurrent
  - GQA: 32 concurrent, still limited by KV
  - MQA: 32 concurrent, minimal memory limit
  - GQA achieves 3-4x throughput improvement

Why GQA slower than MQA:
  - Larger KV cache (8 heads vs 1)
  - More memory bandwidth needed
  - But quality is similar to standard!

Training Impact

Training with MQA/GQA:

Standard:
  - Forward: Standard attention computation
  - Backward: Compute gradients for 32 K/V heads
  - Time: Baseline (1x)

GQA:
  - Forward: Reduced attention computation
  - Backward: Compute gradients for 8 K/V heads
  - Time: ~5-10% faster training!

MQA:
  - Forward: Minimal attention computation
  - Backward: Compute gradients for 1 K/V head
  - Time: ~10-15% faster training!

Training savings are modest (quality cost not worth it)

---

When to Use

MQA (Extreme Efficiency)

Use when:
✅ Inference speed is critical
✅ Mobile/edge deployment
✅ 1-2% quality loss acceptable
✅ Need absolute minimum cache size

Avoid when:
❌ Quality is paramount
❌ Training stability matters
❌ Production systems with SLAs

GQA (Balanced)

Use when:
✅ Want 4x cache reduction (2-3x throughput)
✅ Quality must be near-original (0.5-1% loss)
✅ Training stability important
✅ Production systems (recommended!)

Recommendation:
Default to GQA for new models.
MQA only if cache is absolute bottleneck.

---

Key Takeaways

🔑 MQA: 32x cache reduction but 5-10% quality loss
⚖️ GQA: 4-8x cache reduction with <1% quality loss
📊 GQA is becoming industry standard (better trade-off)
⚡ 4x throughput improvement with GQA possible
🎯 Critical for 70B+ models at scale

---

Related Notes

• Kv Cache - What MQA/GQA optimize
• Pagedattention - Works with MQA/GQA
• Flash Attention - Attention optimization
• Llm Inference Optimization - Complete stack