Grouped Query Attention › The Spectrum of Key-Value Sharing
Language AI HandbookAll chapters Grouped Query Attention › The Spectrum of Key-Value SharingOne section · Sections in this chapter2 / 14 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Before diving into GQA, let's place it within the range of attention mechanisms. The point is that key-value sharing exists on a spectrum, with Multi-Head Attention (MHA) and Multi-Query Attention (MQA) as the two extremes.
Grouped Query Attention divides query heads into groups, with each group sharing a single set of key-value heads. This provides a tunable trade-off between MHA's full capacity and MQA's memory efficiency, controlled by the number of key-value groups.
In standard MHA, each of the query heads has its own dedicated key and value projections. This maximizes representational capacity but requires storing key-value pairs in the cache during inference. In MQA, all query heads share a single key-value pair, reducing cache size by a factor of but potentially limiting what information the model can attend to in parallel.
GQA sits between these extremes. With query heads and key-value groups (where ), each group of query heads shares one key-value pair. The memory reduction factor is , which can be tuned to balance efficiency and quality.
def attention_variants_comparison(
num_query_heads: int,
num_layers: int,
head_dim: int,
seq_len: int,
batch_size: int,
):
"""Compare KV cache sizes for different attention variants."""
bytes_per_element = 2 # fp16
results = []
# Test various numbers of KV heads
kv_head_options = [num_query_heads, 8, 4, 2, 1] # MHA, GQA variants, MQA
for num_kv_heads in kv_head_options:
# KV cache: 2 (K and V) × layers × kv_heads × head_dim × seq × batch × bytes
cache_size = (
2
* num_layers
* num_kv_heads
* head_dim
* seq_len
* batch_size
* bytes_per_element
)
# Determine variant name
if num_kv_heads == num_query_heads:
name = "MHA"
elif num_kv_heads == 1:
name = "MQA"
else:
name = f"GQA-{num_kv_heads}"
results.append(
{
"name": name,
"kv_heads": num_kv_heads,
"cache_gb": cache_size / 1e9,
"reduction": num_query_heads / num_kv_heads,
}
)
return results
# LLaMA 2 70B-like configuration
results = attention_variants_comparison(
num_query_heads=64, num_layers=80, head_dim=128, seq_len=4096, batch_size=32
)KV Cache Comparison (LLaMA 2 70B config, batch=32, seq=4096) ============================================================ Variant KV Heads Cache Size Reduction ------------------------------------------------------------ MHA 64 343.6 GB 1x GQA-8 8 42.9 GB 8x GQA-4 4 21.5 GB 16x GQA-2 2 10.7 GB 32x MQA 1 5.4 GB 64x GQA-8 provides significant memory savings while retaining more representational capacity than MQA
The numbers tell a clear story. GQA with 8 key-value heads reduces cache size by 8x compared to MHA, a substantial improvement that enables larger batches or longer sequences. At the same time, it retains 8 separate key-value representations instead of collapsing everything to one, preserving the model's ability to attend to multiple distinct aspects of the input.