Paper: "GQA: Training Generalized Multi-Query Attention Generalization..." (Ainslie et al., 2023) ArXiv: https://arxiv.org/abs/2305.13245 Key Insight: Multi-head attention creates K, V for every head (expensive). Share K, V across groups of query heads. Dramatically reduces KV cache size with minimal accuracy loss.
The Problem: KV Cache is Huge
Standard Multi-Head Attention:
num_heads = 32
head_dim = 64
For each token in context:
Store: K_all_heads = (32, 64) = 2,048 dims
Store: V_all_heads = (32, 64) = 2,048 dims
For 4K context: 4,000 × 2,048 × 2 = 16.4M values = 66MB just for KV cache
For 100K context: 1.6GB just for KV cache!
This KV cache grows with context length and limits inference throughput.
GQA: Grouped Query Attention
Instead of separate K, V for each head, share K, V across groups:
Standard MHA (32 heads):
- 32 query heads
- 32 key heads
- 32 value heads
GQA with groups=4:
- 32 query heads (8 per group)
- 4 key heads (shared across groups)
- 4 value heads (shared across groups)
Result: 8× fewer K, V to store!
Implementation
import torch
import torch.nn as nn
import torch.nn.functional as F
class GroupedQueryAttention(nn.Module):
def __init__(self, d_model, num_heads, num_kv_heads, head_dim=None):
super().__init__()
if head_dim is None:
head_dim = d_model // num_heads
self.num_heads = num_heads
self.num_kv_heads = num_kv_heads
self.head_dim = head_dim
# Query projection: creates all num_heads
self.W_q = nn.Linear(d_model, num_heads * head_dim)
# Key projection: creates only num_kv_heads (shared)
self.W_k = nn.Linear(d_model, num_kv_heads * head_dim)
# Value projection: creates only num_kv_heads (shared)
self.W_v = nn.Linear(d_model, num_kv_heads * head_dim)
# Output projection
self.W_o = nn.Linear(num_heads * head_dim, d_model)
def forward(self, x, kv_cache=None):
batch_size, seq_len, d_model = x.shape
# Project to Q, K, V
Q = self.W_q(x).reshape(batch_size, seq_len, self.num_heads, self.head_dim)
K = self.W_k(x).reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
V = self.W_v(x).reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
# Transpose for attention
Q = Q.transpose(1, 2) # (batch, num_heads, seq_len, head_dim)
K = K.transpose(1, 2) # (batch, num_kv_heads, seq_len, head_dim)
V = V.transpose(1, 2) # (batch, num_kv_heads, seq_len, head_dim)
# **Key difference: Repeat K, V for each query group**
# Maps num_kv_heads to num_heads
repeat_factor = self.num_heads // self.num_kv_heads
K = K.repeat(1, repeat_factor, 1, 1) # (batch, num_heads, seq_len, head_dim)
V = V.repeat(1, repeat_factor, 1, 1) # (batch, num_heads, seq_len, head_dim)
# Standard attention
scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)
attn_weights = F.softmax(scores, dim=-1)
attn_output = torch.matmul(attn_weights, V)
# Reshape and project
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.reshape(batch_size, seq_len, self.num_heads * self.head_dim)
output = self.W_o(attn_output)
return output
# Usage: Standard 32 heads, but only 4 key-value heads
gqa = GroupedQueryAttention(
d_model=768,
num_heads=32,
num_kv_heads=4 # Group size = 32/4 = 8
)
x = torch.randn(4, 512, 768)
output = gqa(x)
Benchmarks
LLaMA-2 70B Inference
Context: 4K tokens
Standard MHA:
- KV cache: 268MB per sequence
- Throughput: 100 tok/s
- Latency: 10ms per token
With GQA:
- KV cache: 33MB per sequence
- Throughput: 300 tok/s
- Latency: 3ms per token
3× speedup, 8× memory reduction, <1% accuracy loss!
Accuracy Impact
Benchmark: MMLU, HellaSwag, TruthfulQA
Standard MHA: 100%
GQA-4 (4 kv heads): 99.8% (minimal loss)
GQA-8 (8 kv heads): 99.9% (tiny loss)
Quality barely changes, but inference is dramatically faster.
Our Analysis: Why This Matters
GQA is elegant because it solves a real problem (KV cache bloat) with a minimal architectural change. The key insight—that you don't need separate K, V for every query head—is simple but powerful. This is exactly the kind of practical optimization that drives production LLMs.
Practical Tip: GQA is now standard in state-of-the-art models:
LLaMA-2: No GQA
LLaMA-2-Chat: GQA enabled (3× faster inference)
Falcon: Full GQA from the start
References
- Paper: GQA - Training Generalized Multi-Query Attention (Ainslie et al., 2023)
- Code: Available in HuggingFace (Falcon, LLaMA models)
Conclusion
GQA demonstrates that architectural innovations can dramatically improve inference efficiency without sacrificing quality. Understanding how different components of attention contribute to inference cost is essential for production LLM deployment. Next: we'll explore RLHF and how to align models with human preferences.