Multi-Query Attention: sharing one KV head across all query heads

The problem, from first principles

[02-multi-head-attention]'s per-token KV cache grows with n_heads — at long context lengths and large batch sizes this becomes the dominant memory cost during decoding, not the model weights. Multi-Query Attention (MQA) fixes this: every query head still has its own learned projection, but all query heads share a single key head and a single value head, shrinking the KV cache by a factor of n_heads with only a modest quality cost.

From theory to code

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

Q_h = X @ W_Q[h] for h = 1..n_heads;   K = X @ W_K;   V = X @ W_V
head_h = softmax(Q_h K^T / sqrt(d_head)) @ V     -- every head uses the SAME K, V
output = concat(head_1, ..., head_H) @ W_O

Constraints

• Q is split into n_heads heads of size d_head = d_model // n_heads, same as MHA.
• K and V are each a single (seq_len, d_head) matrix, shared across every query head — W_K, W_V project to d_head, not d_model.
• Attention scale and masking rules are identical to [01-scaled-dot-product-attention].
• Output projection W_O has shape (n_heads * d_head, d_model).

Hints