The spectrum: MHA, MQA, and the gap between them

Standard multi-head attention (MHA) runs h parallel heads, each with its own learned W_Q, W_K, W_V. For a model of width d the per-head dimension is d_head = d/h, so the cache must hold h distinct key vectors and h distinct value vectors per token, per layer. That is the memory cost that hurts.

Multi-query attention (MQA), introduced by Shazeer, takes the extreme opposite: all h query heads attend to a single shared key head and value head. The KV cache shrinks by a full factor of h — but collapsing every head onto one K/V pair throws away representational capacity, and in practice it can degrade quality and make training less stable. GQA, from Ainslie et al. (2023), fills the gap: instead of one shared K/V or h of them, use g of them, a tunable dial between the two extremes.

Advertisement

The GQA construction

Partition the h query heads into g equal groups of h/g heads each. Every group gets one shared key projection and one shared value projection; the query heads inside a group keep their own W_Q. So there are still h query heads (queries are cheap — they are not cached) but only g key heads and g value heads.

heads:      h query heads,  g key heads,  g value heads
grouping:   query heads {1 .. h/g} -> KV group 1
            query heads {h/g+1 .. 2h/g} -> KV group 2   ...
shared:     each KV group serves  h/g  query heads
endpoints:  g = h  ->  MHA  (every head its own KV)
            g = 1  ->  MQA  (all heads share one KV)

The attention math is unchanged within a head: softmax(QK^T / √d_head) V. The only change is which K and V a given query head reads — the shared ones for its group. Implementations either replicate (broadcast) each group’s K/V across its h/g query heads, or use a kernel that reads the shared K/V directly, avoiding the copy.

Advertisement

Why the KV cache, not the query, is the bottleneck

During autoregressive decoding the model generates one token at a time. To attend over the whole history without recomputing it, every layer stores the keys and values of all past tokens — the KV cache. Queries are not cached: the current token produces a fresh query each step and discards it. So the memory that accumulates with context length is entirely K and V.

That cache is read in full at every decode step. Modern decode is memory-bandwidth-bound: the arithmetic per token is tiny, but the hardware must stream the entire KV cache (and the weights) out of HBM to compute the next token. Time-per-token therefore tracks the bytes moved, not the FLOPs. Shrinking the KV cache does double duty: it frees memory so you can fit longer context or larger batches, and it cuts the bandwidth each decode step must pay — making decode faster. GQA attacks exactly this quantity.

The KV-cache size formula

For a single sequence, the cache holds, at every layer, a key and a value vector for each KV head and each token seen so far. In bytes:

KV_bytes = 2 · L · n_kv · d_head · S · B · p

  2       K and V
  L       number of layers
  n_kv    number of KV heads  = g   (groups)
  d_head  per-head dimension  = d / h
  S       sequence length (tokens cached)
  B       batch size
  p       bytes per element   (2 for fp16 / bf16)

The query-head count h does not appear — only n_kv = g does. MHA sets n_kv = h; GQA sets n_kv = g; MQA sets n_kv = 1. Dividing the GQA cache by the MHA cache, every factor cancels except the head counts:

KV(GQA) / KV(MHA)  =  g / h
reduction factor  =  h / g

Choosing g is thus a direct memory dial: h/g is precisely the factor by which GQA shrinks the KV cache versus multi-head.

A worked example: Llama-2 70B

Llama-2 70B has width d = 8192, h = 64 query heads, d_head = 128, and L = 80 layers, in bf16 (p = 2). The per-token, per-layer K+V footprint is 2 · n_kv · 128 · 2 bytes; multiply by 80 layers.

MHA  (n_kv = 64):  2 · 80 · 64 · 128 · 2  = 2,621,440 B  ≈ 2.50 MiB / token
GQA  (n_kv =  8):  2 · 80 ·  8 · 128 · 2  =   327,680 B  ≈ 0.31 MiB / token

reduction = h/g = 64/8 = 8×

at S = 4096 tokens (one sequence):
  MHA:  2.50 MiB × 4096 ≈ 10.0 GiB
  GQA:  0.31 MiB × 4096 ≈  1.25 GiB

The same 4096-token context costs 10 GiB of KV cache under MHA but only 1.25 GiB under GQA with g = 8 — an eightfold saving. That is the difference between fitting a handful of concurrent long-context requests on a GPU and fitting dozens, and Llama-2 70B indeed uses g = 8.

The quality-versus-memory trade

Why not always pick g = 1 and take the biggest saving? Because each KV group is a genuine bottleneck: the h/g query heads sharing it must all attend through the same key and value subspaces. Fewer groups means more heads crammed onto each shared K/V, less room to specialize, and measurably worse quality — the MQA failure mode.

The empirical finding is that the curve is sharply asymmetric. Going from MHA to a modest number of groups (commonly g = 8) captures nearly all of MQA’s memory and speed benefit while giving up almost nothing in accuracy — GQA-8 sits close to MHA on quality yet close to MQA on cost. Pushing further toward g = 1 keeps saving memory but starts to cost real quality. So g is chosen small enough to slash the cache, large enough to preserve per-head diversity; eight is the value that has become the de-facto standard.