Grouped-Query Attention: the middle ground between MHA and MQA

The problem, from first principles

[03-multi-query-attention]'s single shared KV head can hurt quality on larger models, but [02-multi-head-attention]'s per-head KV cache is expensive. Grouped-Query Attention (GQA), used in Llama 2/3 and Mistral, is the middle ground: query heads are partitioned into n_kv_heads equal-sized groups, and every query head in a group attends to the SAME key/value head. n_kv_heads == n_heads recovers exactly MHA; n_kv_heads == 1 recovers exactly MQA.

From theory to code

Implement grouped_query_attention(X, W_Q, W_K, W_V, W_O, n_heads, n_kv_heads, mask=None):

Q_h = X @ W_Q[h]                            for h = 1..n_heads
K_g = X @ W_K[g],  V_g = X @ W_V[g]          for g = 1..n_kv_heads
head_h = softmax(Q_h K_g^T / sqrt(d_head)) @ V_g   where g = h // (n_heads / n_kv_heads)

Constraints

• n_heads must be divisible by n_kv_heads.
• Query head h attends to KV group floor(h / (n_heads / n_kv_heads)) — consecutive query heads are grouped in blocks of size n_heads // n_kv_heads.
• Reduces to MQA behavior when n_kv_heads == 1 and to MHA behavior when n_kv_heads == n_heads.
• Same scaling/masking rules as [01-scaled-dot-product-attention].

Hints