GQA Math — Grouped-Query Attention

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.

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.

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.

Uptraining an MHA model into GQA

A key practical result from the GQA paper: you do not need to pretrain a GQA model from scratch. You can convert an existing multi-head checkpoint and then briefly continue training. The conversion is a mean-pool: for each group, average the h/g original key projections into one shared key projection, and likewise for values.

W_K^group = mean( W_K^head  for head in group )     # per group
W_V^group = mean( W_V^head  for head in group )
W_Q          unchanged   (still h heads)

Mean-pooling is a better initializer than picking one head or random init because it preserves the average behavior the model already learned. After conversion, uptrain on a small fraction of the original pretraining budget — roughly 5% — and the model recovers to near-MHA quality. This makes GQA cheap to adopt: take a trained MHA model, pool, uptrain briefly, and ship a version with a far smaller inference footprint.

Why Llama-2 and Llama-3 use it

Llama-2 adopted GQA (with g = 8) for its larger 34B and 70B models, while the small 7B and 13B kept plain MHA — the KV cache only becomes a dominant cost once models are large and contexts long. Llama-3 went further and uses GQA across the board, including the 8B, with 8 KV heads throughout.

The reasoning is squarely the decode economics above. These models target long context and high-throughput serving, where the KV cache dominates memory and decode bandwidth. An eightfold smaller cache means more concurrent sequences per accelerator, longer usable context within the same memory budget, and faster tokens because each step streams less data. Uptraining kept the quality cost negligible. GQA is therefore not an exotic optimization but the default attention variant in most current open-weight LLMs — the standard trade of a hair of quality for a large, permanent reduction in serving cost.

Shapes, pitfalls, and CPU-SLM notes

A few things to keep straight. Divisibility: g must divide h evenly — groups are equal-sized. Head dimension is unchanged: GQA reduces the number of K/V heads, not d_head; queries still have h heads of size d_head. Broadcast vs. kernel: a naive implementation materializes the shared K/V by repeating each group h/g times, which restores MHA-shaped tensors and gives up the memory-bandwidth win inside the kernel — fused GQA kernels instead read the shared K/V once, which is where the decode speedup actually comes from.

For small models running on CPU, the same logic holds and often matters more: RAM bandwidth is scarce, and a GQA cache that is h/g times smaller both fits in cache-friendly space and moves fewer bytes per token. GQA does not reduce prefill FLOPs meaningfully — it is a decode-memory optimization — but decode is precisely the bandwidth-bound phase that dominates interactive latency.