E2.4 · LESSON 4 OF 6

MQA and GQA

Calculate KV-cache memory and explain how layers, KV heads, context length and concurrent sequences affect it.

THE IDEA

Which head count counts?

Many models use fewer KV heads than query heads. Grouped-query attention lets several query heads share one KV head; multi-query attention uses a single one. The queries are not stored at all, so the cache depends only on the number of KV heads. With 32 query heads, going from 32 to 8 KV heads makes the cache 4 times smaller; a single KV head makes it 32 times smaller.

KV heads (32 query heads):

KV heads	name	KiB per token
32	multi-head	512
16	grouped-query	256
8	grouped-query	128
4	grouped-query	64
2	grouped-query	32
1	multi-query	16

All 32 query heads still compute; only the keys and values that must be remembered are shared. The cache formula uses the number of KV heads, not query heads.

Colours show which KV head each of the 32 query heads reads.

Predict first. Which head count belongs in the cache formula?

Chapter context

Causal attention can reuse keys and values for earlier tokens rather than recomputing them at every generated step. This saves work but creates state that grows with the number of retained token positions.

TRY IT YOURSELF

Part 1. Do it hands-on

Compare multi-head, grouped and multi-query caches by hand, then write the grouping and comparison code.

KV heads (32 query heads):

KV heads	name	KiB per token
32	multi-head	512
16	grouped-query	256
8	grouped-query	128
4	grouped-query	64
2	grouped-query	32
1	multi-query	16

All 32 query heads still compute; only the keys and values that must be remembered are shared. The cache formula uses the number of KV heads, not query heads.

KiB per token with 32 KV heads times smaller with 8 KV heads KiB per token with 1 KV head

To do: The same model with 32 query heads and 32 KV heads (full multi-head attention). KiB per token (L = 32, head size 128, 2 bytes)?
To do: How many times smaller is the cache with 8 KV heads (grouped-query) than with 32?
To do: With a single shared KV head (multi-query), KiB per token?

Part 2. Do it in code

bytes_per_token(layers, kv_heads, head_dim, bytes_per_value) as before (2 · layers · kv_heads · head_dim · bytes_per_value).
valid_grouping(n_q_heads, n_kv_heads): True if n_kv_heads divides n_q_heads and both are at least 1.
cache_reduction(n_q_heads, n_kv_heads): n_q_heads / n_kv_heads (ValueError for an invalid grouping).
kv_head_for(q_head, n_q_heads, n_kv_heads): the KV head a query head uses (consecutive groups).
compare_attention_types(layers, n_q_heads, head_dim, bytes_per_value): a dict with the bytes per token for multi-head (kv = n_q), grouped (kv = n_q // 4) and multi-query (kv = 1).

def bytes_per_token(layers, kv_heads, head_dim, bytes_per_value):
    pass

def valid_grouping(n_q_heads, n_kv_heads):
    pass

def cache_reduction(n_q_heads, n_kv_heads):
    pass

def kv_head_for(q_head, n_q_heads, n_kv_heads):
    pass

def compare_attention_types(layers, n_q_heads, head_dim, bytes_per_value):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Which head count belongs in the cache formula?

AQuery heads, because every query is cached.BWhichever is larger.CKV heads, because only keys and values are cached and queries are not stored.

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

For a uniform full-attention decoder:

M_{KV}=2L H_{KV}d_h T B b.

L is layers, H_KV KV heads, d_h head dimension, T retained tokens per sequence, B equal-length unshared sequences and b bytes per value. The factor 2 counts keys and values. For unequal lengths, replace TB by the sum of retained lengths.

Worked example

For L=32, H_KV=8, d_h=128 and b=2, each retained token needs 2×32×8×128×2=131,072 bytes, or 128 KiB. At T=4096, one sequence needs 512 MiB. Four equally long independent sequences need 2 GiB, before allocator overhead and other memory.

Reference code

The following example uses the Python standard library and does not call a remote service.

layers,heads,dim,bytes_per_value=32,8,128,2
per_token=2*layers*heads*dim*bytes_per_value
assert per_token == 131072
assert per_token/1024 == 128
assert per_token*4096/1024**2 == 512
assert per_token*4096*4/1024**3 == 2

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Use KV heads rather than query heads. Distinguish decimal GB from binary GiB. Hybrid, sliding-window and quantized-cache architectures require their own detailed accounting.