Key-Value Head Sharing Changes Inference Memory

Distinguish query heads from K/V heads and calculate exact cache reductions for multi-head, grouped-query, and multi-query attention.

Multi-query attention and grouped-query attention change how query heads share key and value projections. The query-head count may stay fixed while the number of K/V heads falls, reducing the decode cache.

Keep the Head Counts Distinct

Let $h_q$ be query heads and $h_{kv}$ be key/value heads.

multi-head attention (MHA): usually $h_{kv}=h_q$ ;
grouped-query attention (GQA): $1<h_{kv}<h_q$ ;
multi-query attention (MQA): $h_{kv}=1$ .

Several query heads share each K/V head in GQA or MQA. Scaled dot-product attention remains, but the projection and cache contract changes.

Calculate the Cache

For batch $B$ , layers $L$ , cached length $T$ , and head width $d_h$ , keys and values contain

N_{KV}=2LBh_{kv}Td_h

scalars. With $B=1$ , $L=4$ , $T=16$ , $d_h=8$ , and eight query heads:

Design	$h_{kv}$	Cache scalars	float32 bytes
MHA	8	8192	32768
GQA	2	2048	8192
MQA	1	1024	4096

MQA uses one eighth of this MHA cache. This is a shape result, not a latency benchmark. Kernel design, memory movement, batching, and hardware determine measured speed.

Q1. Calculate a grouped-query cache

For $B=2$ , $L=3$ , $h_{kv}=2$ , $T=10$ , and $d_h=4$ , how many K/V cache scalars are stored?

Your answer

Compute it first, then check your number.

Hint

Evaluate

2LBh_{kv}Td_h

Solution

2\times3\times2\times2\times10\times4=960

scalars.

Preserve the Evidence Boundary

Shazeer (2019) reports multi-query decoding experiments for specified models and hardware. The cache formula generalizes under its declared tensor contract; the measured quality and speed results do not automatically generalize to every implementation.