How grouped-query attention shrinks the KV cache

Grouped-query attention kv cache reduction is the single most effective architectural lever for cutting inference memory without retraining. By letting multiple query heads share a single key-value head, GQA slashes the cache footprint proportionally to the grouping ratio while preserving most of multi-head attention’s quality. This post breaks down the mechanics, quantifies the savings, and maps where the tradeoffs bite.

What the KV cache actually stores

Every transformer decoder layer maintains a cache of past keys and values so each new token can attend to the full history without recomputation. For a model with L layers, H heads, head dimension d, and sequence length T, the cache holds two tensors per layer:

K_cache: [L, T, H, d]   # keys
V_cache: [L, T, H, d]   # values

Total elements per token per layer: 2 * H * d. At FP16 (2 bytes), that’s 4 * H * d bytes per token per layer. A 7B model with 32 layers, 32 heads, and 128-dim heads burns roughly 1 MB per 1K tokens just for the cache. At 32K context, that’s 32 MB per request — before activations, logits, or batch overhead.

Multi-head attention baseline

Standard multi-head attention (MHA) computes independent projections for every head:

# MHA per layer
Q = x @ W_Q  # [B, T, H, d]
K = x @ W_K  # [B, T, H, d]
V = x @ W_V  # [B, T, H, d]

# Each head gets its own K, V
attn = softmax(Q @ K.transpose(-2, -1) / sqrt(d)) @ V

Every head learns distinct key-value representations. This maximizes representational capacity but forces the cache to store H independent key-value streams. The memory scales linearly with head count.

How grouped-query attention changes the math

GQA partitions the H query heads into G groups. All queries in a group share one key head and one value head. The projection matrices shrink accordingly:

# GQA per layer
Q = x @ W_Q           # [B, T, H, d]  -- still H query heads
K = x @ W_K           # [B, T, G, d]  -- only G key heads
V = x @ W_V           # [B, T, G, d]  -- only G value heads

# Broadcast K, V to match Q's head dimension for attention
K_expanded = K.repeat_interleave(H // G, dim=2)  # [B, T, H, d]
V_expanded = V.repeat_interleave(H // G, dim=2)

attn = softmax(Q @ K_expanded.transpose(-2, -1) / sqrt(d)) @ V_expanded

The cache now stores G key-value streams instead of H. Memory per token per layer drops from 2 * H * d to 2 * G * d elements. The grouping ratio r = H / G directly determines the compression factor.

Llama-2-7B uses H=32, G=8 (4:1). Llama-3-8B uses H=32, G=8 (4:1). Mixtral uses H=32, G=8 (4:1). The pattern is consistent: 4x cache reduction with minimal quality loss.

Concrete memory savings

Let’s put numbers to a realistic serving scenario. Assume a 7B-class model: 32 layers, 4096 hidden dim, 32 heads, 128 head dim, FP16 cache.

At batch size 32 with 4K average context, MHA consumes ~4 GB just for KV cache. GQA 4:1 drops that to ~1 GB. That difference determines whether you fit on a 24 GB GPU or need model parallelism.

The cache write bandwidth also shrinks by the same factor. Each decoding step writes 2 * G * d * L elements instead of 2 * H * d * L. On memory-bound GPUs (which is most of them during decode), this directly improves tokens/second.

Tradeoffs: quality vs memory

GQA is not free. Sharing keys and values constrains the model’s ability to attend to different positions with different semantic lenses per head. The question is how much capacity you actually lose.

Empirically, 4:1 grouping (GQA) matches MHA quality on most benchmarks. 8:1 starts showing measurable degradation on tasks requiring fine-grained retrieval or long-context reasoning. 32:1 (MQA) degrades noticeably on code, multilingual, and long-context evals.

Why does 4:1 work so well? Attention heads in trained MHA models exhibit high redundancy. Many heads learn similar key-value projections — some specialize in syntax, others in semantics, but the key-value space compresses well. Queries remain diverse because they drive the routing decision: “which positions matter for this query?” Keys and values just need to represent “what is at this position?” — a question with fewer distinct answers.

You can verify this yourself. Take a checkpoint, compute the pairwise cosine similarity between key projection matrices across heads. You’ll find clusters of near-identical projections. GQA essentially forces that clustering architecturally.

When to use GQA vs MHA vs MQA

Use MHA when: training from scratch with unlimited compute, targeting maximum quality on complex reasoning, or building a base model for downstream distillation. The cache cost is a deployment problem, not a training problem.

Use GQA (4:1) as the default for any new model you intend to serve. It’s the Pareto frontier: 75% cache reduction, negligible quality loss, supported by every major inference engine (vLLM, TGI, TensorRT-LLM, llama.cpp). If you’re adapting a checkpoint via continued pretraining or full fine-tune, you can convert MHA to GQA by averaging key/value projections within groups — see the convert_mha_to_gqa script in the Hugging Face transformers examples.

Use MQA (1:1) only when: deploying sub-1B models on edge devices, or when KV cache is the absolute bottleneck and you’ll accept quality hits. MQA’s degradation compounds with context length; avoid it for 16K+ context windows.

Avoid GQA when: you need head-level interpretability (e.g., attention visualization for debugging), or you’re doing something exotic like per-head routing that assumes independent key-value spaces.

Implementation notes for inference engines

If you’re writing a kernel or integrating a model, the cache layout matters. Most engines store GQA cache as [L, T, G, d] and expand on-the-fly during attention. The expansion is a cheap repeat_interleave — no matmul, no extra memory allocation if you fuse it into the attention kernel.

# Fused GQA attention sketch (simplified)
def gqa_attention(Q, K, V, scale):
    # Q: [B, H, T, d], K: [B, G, T, d], V: [B, G, T, d]
    H, G = Q.shape[1], K.shape[1]
    assert H % G == 0
    repeats = H // G
    
    # Expand K, V to match Q heads without materializing full tensors
    # Kernel fuses this: each query head reads from its group's K/V
    K_exp = K.repeat_interleave(repeats, dim=1)  # [B, H, T, d]
    V_exp = V.repeat_interleave(repeats, dim=1)
    
    attn = (Q @ K_exp.transpose(-2, -1)) * scale
    attn = softmax(attn, dim=-1)
    out = attn @ V_exp
    return out

Production kernels (FlashAttention-2, xFormers) avoid the explicit expand by having each query head index directly into its group’s K/V slice. This saves shared memory and register pressure.

The decisive takeaway

Grouped-query attention kv cache reduction is a deployment-first architectural decision that costs almost nothing at training time. Adopt 4:1 GQA as your default for any model you plan to serve. It cuts cache memory and write bandwidth by 75% with benchmark-neutral quality impact. Only reach for MHA if you’re building a foundation model for others to compress, and only reach for MQA if you’re squeezing onto a Raspberry Pi. The math is settled; the engineering is standard; the decision is made.