KV Head Repeat

LLaMA

Easy

Expand KV heads to match query 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

Llama 3 attention is grouped-query: many query heads share one key/value head. The herd paper (Table 3) fixes 8 KV heads against 32 / 64 / 128 query heads. Scaled-dot-product attention is written as a dense product (QK^\top) that expects the same number of heads on (Q), (K), and (V). The official inference code therefore expands each KV head into a contiguous block of query-head slots immediately before that product.

The expansion is a view/reshape, not a learned map. Ainslie et al. (2023) describe the grouping; repeat_kv in llama/model.py is the concrete tensor op. The Llama 3 paper never names this helper — it is inferred from GQA plus the published module.

How it works

repeat_kv takes

x \in \mathbb{R}^{B \times S \times n_{kv} \times d_h}

and an integer (n_{\mathrm{rep}} \ge 1), and returns

x' \in \mathbb{R}^{B \times S \times (n_{kv}\, n_{\mathrm{rep}}) \times d_h}.

If (n_{\mathrm{rep}}=1) (true multi-head, or a rank that already has one query head per local KV head) it returns (x) unchanged. Otherwise the official body is equivalent to torch.repeat_interleave on axis 2:

1. Insert a length-1 axis after the KV-head axis: shape ((B,S,n_{kv},1,d_h)).
2. expand that axis to (n_{\mathrm{rep}}) (no copy in the expand itself).
3. reshape to ((B,S,n_{kv} n_{\mathrm{rep}}, d_h)).

So KV head (g) is copied into query-head indices ([g, n_{\mathrm{rep}},, (g+1), n_{\mathrm{rep}})). For the 8B model, (n_{\mathrm{rep}}=32/8=4): heads 0–3 share KV 0, heads 4–7 share KV 1, and so on. 70B uses 8; 405B uses 16.

Attention.forward calls this twice, on cached keys and cached values, after the cache read and after RoPE (RoPE was applied to the unrepeated (K)). Queries are never repeated. The subsequent transposes treat axis 2 as n_local_heads, which now matches xq.

Memory: the expand/reshape may materialize a tensor (n_{\mathrm{rep}}) times larger for the matmul. The cache stays at (n_{kv}) heads — that is the inference win the paper cites. Repeating only for the current attention op is the point of isolating this step.

The paper does not specify expand-versus-einsum. Some libraries contract GQA as a reshape of (Q) into ((B, n_{kv}, n_{\mathrm{rep}}, S, d_h)) and broadcast (K) to ((B, n_{kv}, 1, S_{\mathrm{cache}}, d_h)) without cloning heads. Algebraically that is the same grouping. The official repo uses explicit repeat then a standard batched matmul.

Official code

llama/model.py — function repeat_kv, and the two call sites in Attention.forward (keys, then values). n_rep = n_local_heads // n_local_kv_heads is set in Attention.__init__. There is no separate CUDA kernel in this repository.

Watch-outs

• Repeat the head axis, not the sequence axis. Interleaving on (S) looks like a valid 5-D expand and silently destroys alignment.
• (n_{\mathrm{rep}}) must be (n_q / n_{kv}) (or the local equivalent) and must divide evenly. A leftover head means the reshape shape is wrong.
• Do not write the repeated tensor back into cache_k / cache_v. Those buffers are allocated with n_local_kv_heads.
• expand is not repeat until the reshape (or a later contiguous op) forces storage. In-place mutation of the expanded view would corrupt all copies of a KV head at once.

Sources

• Paper: The Llama 3 Herd of Models (Llama Team, 2024), §3.2, Table 3 (head counts only)
• GQA: Ainslie et al., “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints,” 2023
• Code: meta-llama/llama3 llama/model.py (repeat_kv)