Grouped-Query Attention
An attention variant that reduces key-value cache memory by sharing key-value heads across query groups.
Grouped-query attention (GQA) is an attention variant that lets several query heads share fewer key-value heads, cutting key-value cache size without collapsing all heads into one shared pair.
At a glance
Released
May 2023
Authors
Joshua Ainslie, James Lee-Thorp, Seth R. Robertson, et al.
Optimizes
- Kv Cache
- Memory Bandwidth
- Long Context Inference
What It Is
Grouped-query attention (GQA) is an attention variant derived from multi-head attention. It keeps multiple query heads but groups them so each group reads from the same key-value head pair.Why It Exists
GQA targets key-value cache size, memory bandwidth during autoregressive decoding, and the cost of serving long contexts when head count would otherwise multiply cache size.How It Works
Queries are partitioned into groups. Each group shares one key head and one value head. Attention scores are computed per query head against the shared key-value pair for that group, and the outputs are then projected as usual.QueriesqqqqValuesKeysVVkkV to kV to kk to qk to qk to qk to q
Math Or Compute Schema
With multiple query heads and several key-value groups, each group serves a subset of the query heads. The formulas below contrast how multi-head attention pairs every query head with its own key-value heads versus grouped-query attention, which routes query heads through shared key-value pairs per group.Compared To Nearby Modules
Compared with multi-head attention, GQA reduces the number of key-value tensors that must be stored. Compared with multi-query attention, GQA keeps more distinct query heads while still sharing key-value heads within each group.| Comparison dimension | Grouped-Query Attention | Multi-Head Attention | Multi-Query Attention |
|---|---|---|---|
| Key-value head count | G key heads and G value heads | H key heads and H value heads | 1 key head and 1 value head |
| Query-head flexibility | H distinct query heads grouped into G shared key-value pairs | H independent query and key-value head pairs | H query heads share one key-value pair |
| Cache footprint per token | 2G tensors (G keys plus G values) | 2H tensors (full multi-head attention cache) | 2 tensors (single shared key-value cache) |