AI Engineer Exhaustive
MediumPlusLLaMA

Grouped Query Attention

LLaMA

Medium

Multi-query attention with shared KV heads

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

Grouped-query attention (GQA; Ainslie et al., 2023) is the only structural change Llama 3 highlights versus Llama 2’s attention layout. Section 3.2: the herd uses GQA “with 8 key-value heads to improve inference speed and to reduce the size of key-value caches during decoding.” Table 3 keeps those 8 KV heads at every size while query heads scale with width (32 / 64 / 128). The paper cites GQA and states the head counts; the projections, RoPE, cache, repeat, and softmax live in Attention in llama/model.py.

GQA sits between multi-head attention (one (K,V) per query head) and multi-query attention (a single shared (K,V)). Query heads are partitioned into (n_{kv}) groups; each group shares one key head and one value head. For Llama 3 that sharing factor is (n_q/8).

How it works

Let (d) be the model width and (n_q) the query-head count. Head width is

d_h = d / n_q

(128 for all three Table 3 sizes). With (n_{kv}=8),

n_{\mathrm{rep}} = n_q / n_{kv} \in \{4,8,16\}

for 8B / 70B / 405B. After the input RMSNorm, a residual (x \in \mathbb{R}^{B \times S \times d}) is projected without bias:

Q = x W_Q \in \mathbb{R}^{B \times S \times n_q d_h},\quad K = x W_K,\; V = x W_V \in \mathbb{R}^{B \times S \times n_{kv} d_h}.

Reshape to ((B,S,n_q,d_h)) and ((B,S,n_{kv},d_h)). RoPE rotates (Q) and (K) (not (V)). In the official inference module the new (K,V) are stored in a cache of shape ((B_{\max}, S_{\max}, n_{kv}^{\mathrm{local}}, d_h)); reads return all keys/values up to start_pos + S.

To run ordinary batched attention, each KV head is repeated (n_{\mathrm{rep}}) times along the head axis (see the KV-repeat note) so (K) and (V) become ((B, S_{\mathrm{cache}}, n_q, d_h)). Then, with heads as the batch-like axis,

\mathrm{Attn}(Q,K,V)=\mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_h}}+M\right)V.

(M) is the additive causal (and cache-offset) mask. Softmax is in float32, then cast back. Heads are merged to (B \times S \times d) and multiplied by (W_O).

The official code also shards heads across model-parallel ranks (n_local_heads = n_q / MP, same for KV). n_rep is computed from the local counts, which equals the global ratio when both divide the world size. That sharding is an implementation detail of Fairscale, not something the paper specifies.

Ainslie et al. define GQA as grouping consecutive query heads onto one KV head — exactly what repeat_kv implements by expanding each KV head into n_rep adjacent query-head slots. The Llama 3 paper does not restate that grouping convention; it is inferred from Ainslie plus the official expand/reshape.

What the paper adds beyond Ainslie: a fixed (n_{kv}=8) at 8B, 70B, and 405B, and the observation that GQA also shrinks the all-gather of (K,V) in their context-parallel attention (later in §3). Those systems points do not change the per-layer math.

Official code

llama/model.py — Attention: wq / wk / wv / wo sizes, n_kv_heads defaulting to n_heads when unset, n_rep, RoPE via apply_rotary_emb, cache write/read, repeat_kv, scaled matmul, softmax, wo. Table 3’s “8 KV heads” is a checkpoint hyperparameter (n_kv_heads in params.json), not a hardcoded constant in the class.

Watch-outs

Sources