Grouped-query attention

Problem

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.

1. repeat_kv(kv, n_heads) repeats each key or value head so that every query head has a matching head.
2. 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.
3. 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:

1. Q = x W_q, K = x W_k, V = x W_v. There are no bias terms.
2. Split Q into H heads and split K and V into G heads each. Repeat K and V to H heads.
3. 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.
4. 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

• 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

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

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

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

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

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

pythonCopyx = 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)

textCopyoutput
[[[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]]]]

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

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

Source paper GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, Ainslie et al., 2023.