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
# MQA
# Savings

-

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
# GQA
# Savings

-

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