AI ENGVisual Encyclopedia

MODULE 3: THE PRE-TRAIN RECIPE · SCENE 12

Sharing the KV cache

Grouped-query attention lets query heads share K/V heads — cheaper memory with almost no quality loss.

KV VRAM SAVINGS: 75.0%
Q0Q1Q2Q3Q4Q5Q6Q7Q8Q9Q10Q11Q12Q13Q14Q15SHARED KV HEAD 0SHARED KV HEAD 1

GQA groups 8 Query heads per KV head, slashing KV cache memory overhead by 87.5%.

TECHNICAL BREAKDOWNModule 3: Modern Architectural Specifications (The Pre-Train Recipe)

Grouped-Query Attention (GQA) & KV Cache Optimization

Standard Multi-Head Attention (MHA) assigns independent Key and Value heads to every Query head, causing massive KV cache VRAM consumption during inference. Grouped-Query Attention groups multiple Query heads to share single Key/Value heads.

MHA vs. MQA vs. GQA

MHA: 1 Q head per 1 KV head (1:1 ratio). MQA: All Q heads share 1 KV head (N:1 ratio, extreme bandwidth saving, minor quality loss). GQA: G Query heads share 1 KV head (e.g. 8:1 ratio, optimal tradeoff).

KV Cache VRAM Reduction

With 8:1 GQA (e.g., 64 Q heads, 8 KV heads in Llama 3 70B), the KV cache footprint during generation drops by 87.5% (8x saving), allowing 8x larger batch sizes or context lengths.

Pre-Training Integration

GQA must be built into pre-training from step 0 so projections learn grouped representations.

MATHEMATICAL FORMULATION · GQA PROJECTION HEAD RATIOS
Heads_KV = Heads_Q / Grouping_Factor_G KV_bytes = 2 × Layers × Head_Dim × Heads_KV × Context_N × Batch × Precision_bytes

Dividing KV head count by grouping factor G directly divides KV memory overhead by G.

REAL-WORLD PRODUCTION ENGINEERING
  • Llama 3 8B and 70B both use GQA with 8 KV heads, enabling fast multi-user serving on single GPU nodes.
← 11 · RMSNorm and SwiGLU
13 · The state you pay to keep →