Stretch: Grouped-Query Attention (GQA)

The problem, from first principles

[04-mha-split-heads]'s Multi-Head Attention gives every query head its OWN dedicated key and value head. During autoregressive text generation, a real serving system has to cache every previously-generated token's key and value vectors, for EVERY head (the "KV cache"), so they don't need to be recomputed at every single new generation step. With num_heads full-sized key/value heads, this cache grows LINEARLY with the number of heads, and for large modern models with many heads and long contexts, the KV cache can become the dominant memory cost of actually serving the model, often larger than the model's own weights for long enough sequences.

Grouped-Query Attention (Ainslie et al., 2023, used in Llama 2/3 and many other modern models) reduces this cost directly: keep the FULL number of query heads (queries are never cached across generation steps, so they don't contribute to this cost at all), but use FEWER key/value heads, with several query heads SHARING the same key/value head. This shrinks the KV cache by exactly the ratio num_query_heads / num_kv_heads, while empirically losing only a small amount of model quality compared to full Multi-Head Attention, a favorable, now widely-adopted tradeoff. (The extreme case, num_kv_heads = 1, a SINGLE shared key/value head for every query head, is called Multi-Query Attention, an even more aggressive predecessor to GQA.)

From theory to code

Implement repeat_kv_heads(x, num_repeats), repeating each key/value head CONSECUTIVELY num_repeats times along the head axis, so (batch, num_kv_heads, seq_len, d_k) becomes (batch, num_kv_heads * num_repeats, seq_len, d_k). Implement grouped_query_attention(query, key, value, num_query_heads, num_kv_heads, mask): split query into num_query_heads heads and key/value into num_kv_heads heads (both via [04-mha-split-heads]'s split_heads, already provided), repeat the key/value heads via repeat_kv_heads so their count matches num_query_heads, then call [01-scaled-dot-product-attention]'s scaled_dot_product_attention once on the now-matching-head-count inputs.

Constraints

• num_query_heads must be evenly divisible by num_kv_heads (assume this holds).
• repeat_kv_heads must repeat each kv head CONSECUTIVELY (kv head 0 repeated num_repeats times in a row, THEN kv head 1 repeated num_repeats times, and so on), not interleaved.
• Query head i must end up paired with kv head i // num_repeats (where num_repeats = num_query_heads // num_kv_heads), the standard "contiguous group" convention.
• When num_kv_heads == num_query_heads, grouped_query_attention must reduce EXACTLY to ordinary Multi-Head Attention (repeat_kv_heads with num_repeats=1 is the identity operation).

Hints