Grouped-Query Attention (GQA)

The problem

While generating text, a model keeps the keys and values of all previous tokens in memory, and that memory is what limits how much text it can handle at once. Grouped-query attention shrinks it: there are still $H$ query heads, but only $G$ key/value heads, and each key/value head is shared by a group of $H/G$ query heads. With $G = H$ this is ordinary multi-head attention; with $G = 1$ it is multi-query attention, where all query heads share one key/value head.

Which query head uses which key/value head

Consecutive query heads form a group. Query head $h$ uses key/value head

where $\lfloor \cdot \rfloor$ rounds down, and $H/G$ is the size of a group.

Example. $H = 4$ query heads and $G = 2$ key/value heads: groups of $4 / 2 = 2$.

query head $h$	$h / 2$	key/value head
0	0	0
1	0.5	0
2	1	1
3	1.5	1

So query heads 0 and 1 share key/value head 0, and query heads 2 and 3 share head 1. The memory holds 2 key/value heads instead of 4: half the size.

The simplest way to compute it: copy each key/value head $H/G$ times, so that there is one per query head, in the order of the table, and then run ordinary causal attention.

What to implement

gqa_attention(q, k, v), with causal masking:

• q has shape $(B, H, T, d)$;
• k and v have shape $(B, G, T, d)$;
• return the output, shape $(B, H, T, d)$.

The attention and causal_mask functions are given.

Requirements

• Raise ValueError if $H$ is not divisible by $G$.
• Query heads $0, \dots, H/G - 1$ share key/value head 0, the next $H/G$ share head 1, and so on.
• No loops over heads.

Hint

Along the head axis, np.repeat(k, 2, axis=1) turns the heads [k0, k1] into [k0, k0, k1, k1]: the order of the table. np.tile would give [k0, k1, k0, k1], which pairs query head 1 with key/value head 1: the wrong grouping. (Real implementations avoid making the copies, but the result is the same.)