Grouped-Query & Multi-Query Attention

The LLM Inference Memory Bottleneck

When serving Large Language Models (LLMs) to multiple users, the model stores past token Key and Value vectors in GPU RAM so it does not have to recompute them at every new token step.

This is called the KV Cache.

In standard Multi Head Attention (MHA) with 32 Query heads, 32 Key heads, and 32 Value heads:

KV Cache Memory = 2 * Layers * Batch_Size * Sequence_Length * Heads * Dimension

For a 70B model serving a batch of 32 users with 4k context:

KV Cache alone consumes over 20 Gigabytes of VRAM!

┌──────────────────────────┬──────────────────────────┬──────────────────────────┐
│ 1. MULTI HEAD (MHA)      │ 2. MULTI QUERY (MQA)     │ 3. GROUPED QUERY (GQA)   │
├──────────────────────────┼──────────────────────────┼──────────────────────────┤
│ 8 Query Heads            │ 8 Query Heads            │ 8 Query Heads            │
│ 8 Key / Value Heads      │ 1 Shared Key / Value Head│ 2 Grouped Key / Value    │
│ (1:1 Ratio)              │ (8:1 Ratio)              │ Heads (4:1 Ratio)        │
│ Highest VRAM Memory      │ Lowest VRAM Memory       │ Optimal Sweet Spot!      │
│ Maximum quality          │ Minor quality drop       │ Near MHA quality         │
└──────────────────────────┴──────────────────────────┴──────────────────────────┘

Visualizing the Attention Variants

  MULTI HEAD ATTENTION (MHA)        MULTI QUERY ATTENTION (MQA)      GROUPED QUERY ATTENTION (GQA)
  Q1 Q2 Q3 Q4 Q5 Q6 Q7 Q8          Q1 Q2 Q3 Q4 Q5 Q6 Q7 Q8          Q1 Q2 Q3 Q4  Q5 Q6 Q7 Q8
  │  │  │  │  │  │  │  │            \  \  \  │  /  /  /  /            \  \  │  /    \  \  │  /
  K1 K2 K3 K4 K5 K6 K7 K8                     K1                               K1            K2
  V1 V2 V3 V4 V5 V6 V7 V8                     V1                               V1            V2

1. Multi Query Attention (MQA - Shazeer, 2019)

MQA collapses all Key and Value heads down to a single shared Key head and Value head.

• Pros: Reduces KV cache memory by factor $h$ (for example, 32x memory reduction!). Enables serving 10x larger user batch sizes on a single GPU.
• Cons: Slight drop in model accuracy on complex reasoning tasks because all Query heads are forced to share the exact same key representations.

2. Grouped Query Attention (GQA - Ainslie et al., 2023)

GQA divides Query heads into $G$ groups. Each group shares a single Key and Value head.

For example, with 32 Query heads and 4 KV heads ($G = 4$):

• Every 8 Query heads share 1 Key head and 1 Value head.

GQA achieves 8x reduction in KV cache memory while matching 99 percent of standard Multi Head Attention quality.

GQA is the inference standard used in LLaMA 3, Mistral, and Claude 3.

Say this out loud

Standard Multi Head Attention maintains equal numbers of Query, Key, and Value heads, creating huge KV cache memory bloat during inference. Multi Query Attention collapses KV heads to a single shared head for maximum memory savings. Grouped Query Attention groups Query heads to share a smaller set of KV heads, reducing KV cache VRAM by 8x while preserving model quality in LLaMA 3.

Followups to expect

1. How does GQA impact training speed versus inference speed? GQA has minimal impact on training speed because training does not use a KV cache. GQA benefits inference by reducing VRAM bandwidth bottlenecking when loading KV tensors.
2. What is PagedAttention (vLLM)? Manages KV cache VRAM dynamically using virtual memory paging concepts, eliminating memory fragmentation and allowing larger batch serving alongside GQA.