ML//transformer//attention//GQA

GQA (Grouped Query Attention) is a variant of multi-head attention that shares the keys \(K\) and values \(V\) across groups of heads, where the standard form gives each head its own, and it is used to save memory at inference. Standard MHA with 96 heads has 96 queries, 96 keys and 96 values; GQA with 8 groups keeps 96 queries but only 8 keys and 8 values.

GQA (Grouped Query Attention) is a variant of multi-head attention that shares the keys KKK and values VVV across groups of heads, where the standard form gives each head its own, and it is used to save memory at inference. Standard MHA with 96 heads has 96 queries, 96 keys and 96 values; GQA with 8 groups keeps 96 queries but only 8 keys and 8 values.

It drastically reduces the size of the KV cache, because the cache grows per key-value pair rather than per head.

Llama 3 and Gemini use GQA: the quality loss is minimal and the memory savings at inference are major.

MQA (Multi-Query Attention) is the extreme case, where all heads share a single key-value pair. GQA is the practical middle ground.