Skip to content

← Attention Variants step 3 of 14

Hard Research

Grouped-Query Attention

Grouped-query attention (GQA) lets several query heads share a key head and a value head. Each query head still computes its own attention weights: sharing keys and values does not make the query heads identical.

With four query heads and two key/value heads, query heads 0 and 1 use K/V head 0, while query heads 2 and 3 use K/V head 1. These are contiguous groups, not alternating assignments.

Your task

Implement grouped_query_attention(Q, K, V) for floating-point arrays:

  • Q: shape (seq_len, n_heads, d_k).
  • K and V: shape (seq_len, n_kv_heads, d_k).

All dimensions are positive, and n_heads is divisible by n_kv_heads. For query head h, choose its key/value head with integer division:

group_size = n_heads // n_kv_heads
g = h // group_size
scores = Q[:, h, :] @ K[:, g, :].T / sqrt(d_k)
weights = softmax(scores, axis=-1)
output_h = weights @ V[:, g, :]

Each score matrix has shape (seq_len, seq_len). Softmax runs across key positions in each row, separately for each query position. Use a numerically stable softmax.

Return shape (seq_len, n_heads, d_k), with output_h at output[:, h, :]. Keep the head axis separate. This exercise has no causal mask, dropout, batch dimension, input projections or final output projection.

How the head counts relate

If n_kv_heads == n_heads, every query head has its own K/V pair, as in multi-head attention (MHA). If n_kv_heads == 1, all query heads share one pair, as in multi-query attention (MQA). Intermediate counts form groups.

The diagram uses eight query heads. Change the number of K/V heads to see which queries share each pair. Your function must also work for other valid head counts.

Loading visualization…