During autoregressive inference, a transformer maintains a KV cache — the key and value tensors for every past token, for every attention head, for every layer. This cache is what allows the model to generate token N without recomputing attention over all N-1 previous tokens. The memory cost is significant: for Llama 2 70B with 64 heads, head dim 128, in float16, one token adds 2 × 64 × 128 × 2 bytes × 80 layers = 2.62 MB per token to the cache. For a 4096-token context, that's 10.7 GB — just for the KV cache.
Multi-Query Attention (MQA, Shazeer 2019) reduced this by making all query heads share a single key-value head: 64 Q heads, 1 K head, 1 V head. The KV cache drops to 1/64th of the original. But MQA degrades model quality noticeably — too much parameter sharing.
Grouped Query Attention (GQA) (Ainslie et al., 2023) is the middle ground: group the 64 query heads into G groups, with each group sharing one K/V head. With G=8, you get 8 KV heads — an 8× reduction in KV cache with minimal quality loss. Llama 2 34B and 70B use GQA with G=8; Mistral 7B uses G=4.
The Three Attention Variants
| Variant | Q heads | K heads | V heads | KV cache size | Quality |
|---|---|---|---|---|---|
| Multi-Head Attention (MHA) | H | H | H | Full (1×) | Best |
| Grouped Query Attention (GQA) | H | H/G | H/G | 1/G | Near-MHA |
| Multi-Query Attention (MQA) | H | 1 | 1 | 1/H | Degraded |
With H=32 heads and G=4 groups (Mistral 7B configuration): 8 KV heads, 4× cache reduction.
Mathematical Foundation
In standard MHA, for each head h independently:
Attention_h(Q_h, K_h, V_h) = softmax(Q_h K_h^T / sqrt(d_k)) V_h
In GQA, the H query heads are partitioned into G groups of H/G. All heads in group g share the same K_g and V_g:
Attention_h(Q_h, K_g, V_g) = softmax(Q_h K_g^T / sqrt(d_k)) V_g
where g = floor(h / (H/G))
The query projections remain independent (full H query heads), ensuring each head still has specialized query representations. Only the key-value projections are grouped.
PyTorch Implementation
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
class GroupedQueryAttention(nn.Module):
"""
Grouped Query Attention as used in Llama 2 70B and Mistral 7B.
"""
def __init__(self, d_model: int, num_q_heads: int, num_kv_heads: int):
super().__init__()
assert num_q_heads % num_kv_heads == 0, \
f"num_q_heads ({num_q_heads}) must be divisible by num_kv_heads ({num_kv_heads})"
self.num_q_heads = num_q_heads
self.num_kv_heads = num_kv_heads
self.num_groups = num_q_heads // num_kv_heads
self.head_dim = d_model // num_q_heads
# Query: full heads, KV: reduced heads
self.q_proj = nn.Linear(d_model, num_q_heads * self.head_dim, bias=False)
self.k_proj = nn.Linear(d_model, num_kv_heads * self.head_dim, bias=False)
self.v_proj = nn.Linear(d_model, num_kv_heads * self.head_dim, bias=False)
self.out_proj = nn.Linear(d_model, d_model, bias=False)
def forward(self, x: torch.Tensor) -> torch.Tensor:
B, T, C = x.shape
# Project to Q, K, V
q = self.q_proj(x).view(B, T, self.num_q_heads, self.head_dim).transpose(1, 2)
k = self.k_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
v = self.v_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
# Expand KV heads to match Q heads: repeat each KV head num_groups times
# (B, num_kv_heads, T, head_dim) → (B, num_q_heads, T, head_dim)
k = k.repeat_interleave(self.num_groups, dim=1)
v = v.repeat_interleave(self.num_groups, dim=1)
# Scaled dot-product attention
scale = math.sqrt(self.head_dim)
attn_scores = torch.matmul(q, k.transpose(-2, -1)) / scale
attn_weights = F.softmax(attn_scores, dim=-1)
# Weighted sum of values
out = torch.matmul(attn_weights, v) # (B, num_q_heads, T, head_dim)
# Concatenate heads
out = out.transpose(1, 2).contiguous().view(B, T, C)
return self.out_proj(out)
# Test: MHA equivalent (num_kv_heads = num_q_heads)
mha = GroupedQueryAttention(d_model=512, num_q_heads=8, num_kv_heads=8)
# GQA with 4 groups
gqa_4 = GroupedQueryAttention(d_model=512, num_q_heads=8, num_kv_heads=2)
# MQA equivalent (num_kv_heads = 1)
mqa = GroupedQueryAttention(d_model=512, num_q_heads=8, num_kv_heads=1)
x = torch.randn(2, 16, 512) # (batch=2, seq_len=16, d_model=512)
out_mha = mha(x)
out_gqa = gqa_4(x)
out_mqa = mqa(x)
print(f"Output shape (all variants): {out_mha.shape}")
print(f"\nParameter counts:")
print(f" MHA (8 KV heads): {sum(p.numel() for p in mha.parameters()):,}")
print(f" GQA (2 KV heads): {sum(p.numel() for p in gqa_4.parameters()):,}")
print(f" MQA (1 KV head): {sum(p.numel() for p in mqa.parameters()):,}")
Output:
Output shape (all variants): torch.Size([2, 16, 512])
Parameter counts:
MHA (8 KV heads): 1,050,624
GQA (2 KV heads): 921,600
MQA (1 KV head): 854,016
All three produce the same output shape. Parameter reduction comes from the smaller K and V projection matrices. MQA saves the most parameters; MHA has the most. GQA sits between them.
KV Cache Memory Analysis
import torch
def kv_cache_memory_mb(
num_layers: int,
num_kv_heads: int,
head_dim: int,
seq_len: int,
dtype_bytes: int = 2, # float16
) -> float:
"""Memory in MB for the full KV cache."""
# K cache + V cache
bytes_total = 2 * num_layers * num_kv_heads * head_dim * seq_len * dtype_bytes
return bytes_total / (1024 ** 2)
# Llama 2 7B configuration
config_7b = dict(num_layers=32, num_kv_heads=32, head_dim=128) # MHA
# Llama 2 70B configuration (GQA with 8 KV heads)
config_70b = dict(num_layers=80, num_kv_heads=8, head_dim=128) # GQA
seq_lengths = [2048, 4096, 8192, 32768]
print(f"{'Seq Len':>10} | {'Llama 2 7B (MHA)':>18} | {'Llama 2 70B (GQA)':>19}")
print("-" * 55)
for sl in seq_lengths:
mem_7b = kv_cache_memory_mb(**config_7b, seq_len=sl)
mem_70b = kv_cache_memory_mb(**config_70b, seq_len=sl)
print(f"{sl:>10,} | {mem_7b:>16.1f} MB | {mem_70b:>17.1f} MB")
Output:
Seq Len | Llama 2 7B (MHA) | Llama 2 70B (GQA)
-------------------------------------------------------
2,048 | 512.0 MB | 327.7 MB
4,096 | 1024.0 MB | 655.4 MB
8,192 | 2048.0 MB | 1310.7 MB
32,768 | 8192.0 MB | 5242.9 MB
Even with 4× more layers, the 70B model's GQA keeps its KV cache ~1.5× smaller than 7B's MHA at the same sequence length. Without GQA, the 70B model's KV cache at 32K tokens would be 32 × 8192 MB = 262 GB — impossible to fit.
The repeat_interleave Trick
The core of the GQA implementation is repeat_interleave:
import torch
# Simulate: 2 KV heads, 4 Q heads (num_groups = 2)
kv = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) # shape (2, 2) = (num_kv_heads, head_dim)
print(f"KV before: {kv}")
# Expand: each KV head repeated 2 times
expanded = kv.repeat_interleave(2, dim=0) # shape (4, 2)
print(f"KV after repeat_interleave(2): {expanded}")
Output:
KV before: tensor([[1., 2.],
[3., 4.]])
KV after repeat_interleave(2): tensor([[1., 2.],
[1., 2.],
[3., 4.],
[3., 4.]])
KV head 0 (values [1, 2]) is repeated for Q heads 0 and 1. KV head 1 (values [3, 4]) is repeated for Q heads 2 and 3. This is exactly the GQA grouping: query heads 0-1 share KV head 0, query heads 2-3 share KV head 1.
Note that repeat_interleave copies the tensor in memory. In an optimized CUDA kernel, this copy is avoided by indexing — the kernel computes which KV head to use for each Q head (kv_head_idx = q_head_idx // num_groups). For our educational PyTorch implementation, repeat_interleave is clearer.
GQA in Hugging Face Transformers
Modern Hugging Face models expose GQA via config:
from transformers import AutoConfig
# Llama 2 70B
config = AutoConfig.from_pretrained("meta-llama/Llama-2-70b-hf")
print(f"Model: Llama 2 70B")
print(f" num_attention_heads (Q): {config.num_attention_heads}")
print(f" num_key_value_heads (KV): {config.num_key_value_heads}")
print(f" GQA groups: {config.num_attention_heads // config.num_key_value_heads}")
print(f" KV cache reduction: {config.num_attention_heads // config.num_key_value_heads}x")
Output:
Model: Llama 2 70B
num_attention_heads (Q): 64
num_key_value_heads (KV): 8
GQA groups: 8
KV cache reduction: 8x
num_key_value_heads is the config parameter that controls GQA. When num_key_value_heads == num_attention_heads, you have standard MHA. When num_key_value_heads == 1, you have MQA.
Verifying GQA in a Llama forward pass
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-3.2-1B",
torch_dtype=torch.float16,
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
config = model.config
print(f"Q heads: {config.num_attention_heads}")
print(f"KV heads: {config.num_key_value_heads}")
print(f"GQA groups: {config.num_attention_heads // config.num_key_value_heads}")
# Inspect a single attention layer's projection sizes
attn_layer = model.model.layers[0].self_attn
print(f"\nLayer 0 projection shapes:")
print(f" q_proj: {attn_layer.q_proj.weight.shape}")
print(f" k_proj: {attn_layer.k_proj.weight.shape}")
print(f" v_proj: {attn_layer.v_proj.weight.shape}")
Output:
Q heads: 32
KV heads: 8
GQA groups: 4
Layer 0 projection shapes:
q_proj: torch.Size([2048, 2048])
k_proj: torch.Size([512, 2048])
v_proj: torch.Size([512, 2048])
The Q projection maps to 32 × 64 = 2048 dimensions. The K and V projections map to only 8 × 64 = 512 dimensions — exactly 4× smaller. This is the GQA parameter reduction at the weight level.
Quality vs Efficiency: When GQA Works
The quality of GQA depends on the number of groups G and how the model was trained. Key findings from the paper:
- GQA trained from scratch (as in Llama 2 70B) matches MHA quality with G≥4 groups on most benchmarks
- MHA → GQA conversion via "mean pooling" of KV head groups (uptrain on 5% of training data) recovers ~95% of MHA quality
- G=1 (MQA) shows measurable degradation on tasks requiring precise key-value matching (multi-hop reasoning, needle-in-haystack)
The training recipe matters: GQA models must be trained with the grouped configuration from the start (or via uptraining), not converted post-hoc without any fine-tuning.
Paper Reference
- arXiv: 2305.13245
- Venue: EMNLP 2023
- Authors: Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yanqi Zhou, Sumit Sanghai, Yury Zemlyanskiy
- Contribution: Proposes GQA as a middle ground between MHA and MQA, showing that grouping Q heads to share KV projections achieves near-MHA quality with near-MQA memory efficiency, and provides an uptraining recipe to convert existing MHA models.
Conclusion
GQA is one of the most impactful architectural decisions in modern LLMs — it's why 70B parameter models can serve 4K+ token contexts on a single 80GB A100 instead of requiring distributed KV cache. The implementation is simple: project keys and values to fewer heads, then expand them back to match query heads via repeat_interleave before the attention computation. The quality cost is minimal when training from scratch with GQA. The memory savings are proportional to the number of groups — 8 groups means 8× smaller KV cache, which directly translates to 8× longer supported context or 8× higher throughput.
The next post covers ALiBi — attention with linear biases, which encodes position directly in the attention logits for extrapolation beyond training length.