Multi-Query Attention

An attention variant that shares one key-value head across all query heads to minimize key-value cache memory.

Multi-query attention (MQA) is an attention variant that keeps multiple query heads but shares one key-value pair across them, shrinking the key-value cache for long-context inference.

At a glance

Released

November 2019

Authors

Noam Shazeer

Optimizes

  • Kv Cache
  • Memory Bandwidth

What It Is

Multi-query attention (MQA) is an attention variant derived from multi-head attention. Every query head still computes its own attention scores, but all heads read from the same key head and value head.

Why It Exists

MQA targets key-value cache size and memory bandwidth during autoregressive decoding. Sharing one key-value pair across all query heads minimizes the number of tensors stored per token.

How It Works

Linear projections still produce multiple query heads, but the key and value projections collapse to a single head each. Every query head attends against the shared key and value tensors, and the outputs are then concatenated and projected as usual.

Math Or Compute Schema

The formulas below contrast how multi-head attention pairs every query head with its own key-value heads versus multi-query attention, which routes all query heads through one shared key-value pair.
Multi-head attention (MHA)
Attention(Qi,Ki,Vi)=softmax ⁣(QiKi⊤dk)Vi\text{Attention}(Q_i, K_i, V_i) = \mathrm{softmax}\!\left(\frac{Q_i K_i^{\top}}{\sqrt{d_k}}\right) V_i
QQ
Query vectors for head i.
KK
Key vectors for head i.
VV
Value vectors for head i.
HH
Number of query heads.
dkd_k
Key dimension per head.
ii
Query head index.
Multi-query attention (MQA)
Attention(Qi,K,V)=softmax ⁣(QiK⊤dk)V\text{Attention}(Q_i, K, V) = \mathrm{softmax}\!\left(\frac{Q_i K^{\top}}{\sqrt{d_k}}\right) V
QQ
Query vectors for head i.
KK
Shared key vectors from one key-value head.
VV
Shared value vectors from one key-value head.
HH
Number of query heads.
dkd_k
Key dimension per head.
ii
Query head index.

Compared To Nearby Modules

Compared with multi-head attention, MQA stores one key-value pair instead of one pair per head. Compared with grouped-query attention, MQA compresses further by sharing a single key-value head across all query heads rather than several groups.
Comparison dimensionMulti-Query AttentionMulti-Head AttentionGrouped-Query Attention
Key-value head count1 key head and 1 value headH key heads and H value headsG key heads and G value heads
Query-head flexibilityH query heads share one key-value pairH independent query and key-value head pairsH distinct query heads grouped into G shared key-value pairs
Cache footprint per token2 tensors (single shared key-value cache)2H tensors (full multi-head attention cache)2G tensors (G keys plus G values)

Example Architectures

MQA appears in decoder-only language models that prioritize inference throughput and lean key-value caches, including some large-scale serving configurations that accept the representational tradeoff.

Limitations And Tradeoffs

Sharing one key-value head across all queries can reduce representational flexibility versus full multi-head attention or grouped-query attention. Quality-sensitive workloads may prefer grouped-query attention or multi-head attention.

Why It Still Matters

MQA established the head-sharing design space. It shows how far key-value compression can go before grouped-query attention offers a middle ground between cache size and query diversity.

Tags

References

  1. Shazeer, Noam. "Fast Transformer Decoding: One Write-Head is All You Need." arXiv, 2019, https://arxiv.org/abs/1911.02150.

Multi-Head Latent Attention

Previous Page

Sliding-Window Attention

Next Page

On this page

What It IsWhy It ExistsHow It WorksMath Or Compute SchemaCompared To Nearby ModulesExample ArchitecturesLimitations And TradeoffsWhy It Still MattersRelatedTagsReferences