Grouped-Query Attention: Sharing KV Heads

Let a transformer layer have

$$H_q$$

query heads and

$$H_{kv}$$

key/value heads. Assume $H_q$ is divisible by $H_{kv}$ and define the group size

$$r = \frac{H_q}{H_{kv}}.$$

For query head $h$, the shared KV head index is

$$g(h)=\left\lfloor \frac{h}{r} \right\rfloor.$$

The attention output for head $h$ is

$$o_h = \mathrm{softmax}\!\left( \frac{Q_hK_{g(h)}^\top}{\sqrt{d_{head}}} \right)V_{g(h)}.$$

For one decoded token, $Q_h$ is a query vector with shape $[d_{head}]$. The cached $K_{g(h)}$ and $V_{g(h)}$ tensors have shape $[T, d_{head}]$ for the selected shared KV head. GQA changes which cached KV channel a query head reads; it does not remove context positions from the softmax.

The three common cases are:

• Multi-head attention: $H_{kv}=H_q$, so $r=1$.
• Grouped-query attention: $1 < H_{kv} < H_q$, so a group of query heads shares each KV head.
• Multi-query attention: $H_{kv}=1$, so every query head shares the same KV head.

During decoding, the KV cache for a batch of $B$ sequences, $L$ layers, context length $T$, head dimension $d_{head}$, and element width $b$ bytes is approximately

$$\mathrm{Mem}_{KV} \approx B\cdot L\cdot T\cdot H_{kv}\cdot d_{head}\cdot 2\cdot b.$$

The factor of 2 stores both keys and values. Compared with multi-head attention, the KV-cache ratio is therefore

$$\frac{\mathrm{Mem}_{GQA}}{\mathrm{Mem}_{MHA}} = \frac{H_{kv}}{H_q},$$

and the reduction factor is

$$\frac{H_q}{H_{kv}}.$$

That ratio is the core invariant. If a model keeps $H_q=32$ query heads but uses $H_{kv}=8$ KV heads, the decoding KV cache is one quarter of full multi-head attention for the same batch, layers, context length, head dimension, and dtype.

This is a logical fixed-shape cache estimate. Real serving stacks may add paging metadata, fragmentation, cross-device duplication, or temporary expanded KV views, so measured allocator memory can differ from the clean structural ratio.