Multi-Query Attention

Attention series: 10 of 11

Attention: the series

1. QKV Projection

2. Attention Computation

3. Self Attention

4. Cross Attention

5. Self Attention vs Cross Attention

6. Self Attention (Shared KV)

7. Multi-Head Attention

8. Fused QKV (Multi-Head)

9. Single vs Multi-Head Attention

10. Multi-Query Attention

11. Grouped-Query Attention

Multi-Query Attention (Shazeer, 2019) is an extension of Shared KV attention to the full multi-head setting.

Open the interactive diagram ↗

In standard Multi-Head Attention, every head has its own Q, K, and V. During autoregressive decoding, K and V from every previous token must be stored in a KV cache. With H heads, that cache is H times larger than a single head's.

MQA keeps H independent query heads but computes K and V just once. Each head can still ask a different question. But they all look up the answers in the same shared directory.

In this example, Head 1 is the full head: it computes Q, K, and V. Heads 2 and 3 each compute their own Q but reuse Head 1's K and V. The KV cache shrinks by a factor of H, with surprisingly little quality loss.

---

Next:
11. Grouped-Query Attention