Repeat each key/value head n_rep times along the head axis.
The expansion underlying grouped-query / multi-query attention: a model
projects fewer key/value heads (num_kv_heads) than query heads
(num_heads) — sharing each K/V head across a group of
n_rep = num_heads // num_kv_heads query heads — which shrinks the K/V
projection and, crucially, the K/V cache. Before the attention scores are
computed, each K/V head is repeated n_rep times so the head count matches
the queries. Head h of the result maps to source head h // n_rep, so
query head q attends to K/V head q // n_rep — the grouping GQA
expects. n_rep == 1 (standard multi-head attention) is a no-op.
Parameters
hidden_statesTensor(B, num_kv_heads, T, head_dim).n_repintnum_heads // num_kv_heads (>= 1).Returns
TensorShape (B, num_kv_heads * n_rep, T, head_dim).
Notes
Equivalent to unsqueeze → expand → reshape (an
interleave-by-group, NOT a tile): the num_kv_heads heads expand to
consecutive blocks, so the output head order is
[kv0, kv0, …, kv1, kv1, …] with n_rep copies per block. This matches
the contiguous query-head grouping used by GQA models, where query heads
[g·n_rep, (g+1)·n_rep) share key/value head g.
Examples
>>> import lucid
>>> import lucid.nn.functional as F
>>> kv = lucid.randn(1, 2, 4, 8) # (B, num_kv_heads=2, T=4, head_dim=8)
>>> out = F.repeat_kv(kv, 3) # 8 query heads / 2 kv heads
>>> out.shape
(1, 6, 4, 8)
>>> bool((out[:, 0] == kv[:, 0]).all().item()) # heads 0,1,2 ← kv head 0
True
>>> bool((out[:, 3] == kv[:, 1]).all().item()) # heads 3,4,5 ← kv head 1
True