Grouped-query attention · AI Engineering

Grouped-query attention (GQA) lets several query heads share one key head and one value head. The model then stores fewer keys and values in its cache. With one group, GQA is multi-query attention. With one group for each query head, GQA is standard multi-head attention. Large open models such as Llama 3 use GQA.

Task

Write three functions.

repeat_kv(kv, n_heads) repeats each key or value head so that every query head has a matching head.
grouped_query_attention(x, w_q, w_k, w_v, w_o, n_heads, n_kv_heads) returns the output and the attention weights of GQA self-attention.
mean_pool_kv_heads(w, n_heads, n_kv_heads) turns the key or value projection of a multi-head checkpoint into a GQA projection.

Let H be n_heads, G be n_kv_heads, r = H / G and d_{head} = D / H. Head h of a projected array holds the columns h \cdot d_{head} to (h + 1) \cdot d_{head} - 1, in the same order. Query head h uses key and value group \lfloor h / r \rfloor. So the first r query heads use group 0, the next r use group 1, and so on.

The steps of grouped_query_attention are:

Q = x W_q, K = x W_k, V = x W_v. There are no bias terms.
Split Q into H heads and split K and V into G heads each. Repeat K and V to H heads.
In each head, compute \mathrm{softmax}(Q_h K_h^T / \sqrt{d_{head}})\, V_h. The softmax is over the key axis. There is no mask.
Join the H head outputs in head order, so that the result has shape (B, T, D), and multiply by W_o.

mean_pool_kv_heads takes a projection with H heads. Group g of the result is the mean of the heads g r to g r + r - 1.

All three functions raise ValueError when H is not a multiple of G.

Input and output

Name	Type	Shape	Meaning
`kv`	float array	`(B, G, T, d_head)`	keys or values split into `G` heads
`x`	float array	`(B, T, D)`	a batch of `B` sequences of `T` tokens with `D` features
`w_q`	float array	`(D, D)`	the query projection
`w_k`, `w_v`	float array	`(D, G * d_head)`	the key and value projections
`w_o`	float array	`(D, D)`	the output projection
`n_heads`	int		the number of query heads `H`
`n_kv_heads`	int		the number of key and value groups `G`
`w`	float array	`(D, H * d_head)`	a key or value projection with `H` heads

repeat_kv returns a float array of shape (B, H, T, d_head).
grouped_query_attention returns a tuple (output, weights). output has shape (B, T, D). weights has shape (B, H, T, T), and weights[b, h, i, j] is the weight that query i gives to key j in head h.
mean_pool_kv_heads returns a float array of shape (D, G * d_head).

Constraints

1 <= B <= 4, 1 <= T <= 32, 1 <= G <= H <= 16 and D is a multiple of H.
Divide the scores by \sqrt{d_{head}}, not by \sqrt{D}.
Do not change the input arrays.
Use numpy only.

Examples

python

kv = np.arange(8.0).reshape(1, 2, 2, 2)
repeat_kv(kv, 4)[0]

text

[[[0. 1.]
  [2. 3.]]

 [[0. 1.]
  [2. 3.]]

 [[4. 5.]
  [6. 7.]]

 [[4. 5.]
  [6. 7.]]]

python

x = np.array([[[1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0]]])
w_k = np.array([[1.0, 0.0], [0.0, 1.0], [0.0, 0.0], [0.0, 0.0]])
w_v = np.array([[1.0, 0.0], [0.0, 2.0], [0.0, 0.0], [1.0, 0.0]])
output, weights = grouped_query_attention(x, np.eye(4), w_k, w_v, np.eye(4), 2, 1)

text

output
[[[1.20333628 1.19777581 0.79666372 1.60444837]
  [0.79666372 1.60444837 1.20333628 1.19777581]
  [1.         1.50348984 1.         1.33333333]]]
weights
[[[[0.40111209 0.19777581 0.40111209]
   [0.19777581 0.40111209 0.40111209]
   [0.24825508 0.24825508 0.50348984]]

  [[0.19777581 0.40111209 0.40111209]
   [0.40111209 0.19777581 0.40111209]
   [0.33333333 0.33333333 0.33333333]]]]

python

mean_pool_kv_heads(np.arange(16.0).reshape(2, 8), 4, 2)

text

[[ 1.  2.  5.  6.]
 [ 9. 10. 13. 14.]]