GQA and the KV cache — Linear algebra for LLMs

Generating one token at a time

To generate token $t + 1$ , the model needs the new token’s query against the keys of all $t$ earlier tokens, and their values. Those earlier keys and values never change (the mask means old tokens never see new ones), so instead of recomputing them every step, the model stores them: the KV cache.

Write

KV cache: keep every past token’s

k

and

v

, per layer and head.

Its size

\displaystyle \text{KV bytes per token} = 2 \times L \times h_{kv} \times d_h \times \text{bytes}

2 for keys and values, times the layers $L$ , times the key/value heads $h_{kv}$ , times the head width, times bytes per number. Multiply by the number of tokens (and by the batch size) for the total.

Write

KV cache per token =

2 \cdot L \cdot h_{kv} \cdot d_h \cdot \text{bytes}

.

With one K/V per query head

\displaystyle L = 32,\; h_{kv} = 32,\; d_h = 128,\; \text{BF16}

\displaystyle 2 \cdot 32 \cdot 32 \cdot 128 \cdot 2 = 524{,}288 \text{ bytes} = 512 \text{ KiB}

A Llama 2 7B-shaped model uses standard multi-head attention: every query head has its own keys and values. Half a megabyte per token: 4,096 tokens take 2 GiB, for one sequence.

Grouped-query attention

\displaystyle h = 32 \text{ query heads}, \quad h_{kv} = 8 \text{ key/value heads}

\displaystyle 2 \cdot 32 \cdot 8 \cdot 128 \cdot 2 = 131{,}072 \text{ bytes} = 128 \text{ KiB}

GQA lets groups of query heads share one key/value head. Llama 3 8B has 32 query heads and 8 KV heads: each KV head serves 4 query heads. The cache shrinks 4×, and $W_K$ , $W_V$ shrink from $4096 \times 4096$ to $4096 \times 1024$ .

Write

GQA:

h_{kv} < h

. Llama 3 8B: 8 KV heads, 128 KiB per token in BF16.

MHA, GQA, MQA

Variant	KV heads	Cache vs MHA
multi-head (MHA)	$h$	1×
grouped-query (GQA)	$h / g$	$1/g$
multi-query (MQA)	1	$1/h$

Fewer KV heads means less memory and less data to read per generated token; the queries keep their full count.

✏️ Your turn

Size a cache

\displaystyle L = 32,\; h_{kv} = 8,\; d_h = 128,\; \text{FP8},\; 32{,}768 \text{ tokens}

How big is the cache for one sequence? (FP8 is 1 byte per number.)

Work it out on paper, then go to the next card to check.

Answer: 2 GiB

\displaystyle 2 \cdot 32 \cdot 8 \cdot 128 \cdot 1 = 65{,}536 \text{ bytes per token}

\displaystyle \times\, 32{,}768 = 2{,}147{,}483{,}648 \text{ bytes} = 2 \text{ GiB}

Why it matters

In a transformer

During generation, every step reads the whole cache (and all the weights) from memory to do a little math with it. Long contexts and big batches make the cache the largest thing in memory after the weights, which is why serving systems page it, compress it, and quantize it.

📌 Recap

GQA and the KV cache in one card

Generation caches past keys and values instead of recomputing them.
Per token: $2 \cdot L \cdot h_{kv} \cdot d_h \cdot \text{bytes}$ .
MHA: $h_{kv} = h$ . GQA: groups share K/V. MQA: one K/V head.
Llama 3 8B in BF16: 128 KiB per token.

Write

The four lines above.