L5.5 · LESSON 5 OF 6

MQA and GQA

Trace a transformer block, distinguish normalization variants and count a simplified block's parameters.

THE IDEA

Sharing keys and values

When a model generates text it must remember the keys and values of every earlier token (the KV cache). With ordinary multi-head attention every query head has its own K and V, so the cache grows with the number of heads. Multi-query attention (MQA) uses one K/V head shared by all query heads. Grouped-query attention (GQA) is in between: query heads are split into groups and each group shares one K/V head.

KV heads:

Q0 → KV0Q1 → KV0Q2 → KV0Q3 → KV0Q4 → KV1Q5 → KV1Q6 → KV1Q7 → KV1

setup (2 layers, head size 64, 1000 tokens, 2 bytes)	KV heads cached	cache bytes
multi-head	8	4,096,000
grouped (4)	4	2,048,000
grouped (2)	2	1,024,000
multi-query	1	512,000

There are still 8 query heads and 8 attention outputs per position; only the keys and values that must be remembered for generation are shared.

Eight query heads. Choose how many KV heads they share and read the cache size.

Predict first. If four query heads use two KV heads, are there still four queries per position?

Chapter context

A transformer combines attention with per-position feed-forward computation and residual paths. Exact ordering and normalization differ between architectures; learn a stated variant instead of assuming every transformer has an identical block.

TRY IT YOURSELF

Part 1. Do it hands-on

Work out which KV head a query head reads and how big the cache is, then implement grouped attention.

KV heads:

Q0 → KV0Q1 → KV0Q2 → KV0Q3 → KV0Q4 → KV1Q5 → KV1Q6 → KV1Q7 → KV1

setup (2 layers, head size 64, 1000 tokens, 2 bytes)	KV heads cached	cache bytes
multi-head	8	4,096,000
grouped (4)	4	2,048,000
grouped (2)	2	1,024,000
multi-query	1	512,000

There are still 8 query heads and 8 attention outputs per position; only the keys and values that must be remembered for generation are shared.

KV head used by query head 5 (8 query heads, 2 KV) cache bytes, 2 KV heads cache bytes, 1 KV head

To do: 8 query heads share 2 key/value heads (grouped-query). In which KV head does query head 5 (counting from 0) read? Type 0 or 1
To do: KV cache bytes = 2 (K and V) × layers × kv_heads × head_dim × tokens × 2 bytes. With 2 layers, head_dim 64, 1000 tokens, how many bytes for 2 KV heads?
To do: And for multi-query attention (a single KV head)?

Part 2. Do it in code

kv_head_for(q_head, n_q, n_kv): which KV head a query head reads. n_q must be a multiple of n_kv (ValueError); query heads are grouped consecutively: group size n_q // n_kv.
kv_cache_bytes(layers, n_kv, head_dim, seq, bytes_per=2): 2 (K and V) × layers × n_kv × head_dim × seq × bytes_per.
grouped_attention(Q, K, V, n_kv): Q has shape (n_q, n, d); K and V have shape (n_kv, n, d). Each query head attends with the K and V of its own KV head (scaled dot-product, no mask). Return (n_q, n, d).
savings(n_q, n_kv): the fraction of the KV cache kept compared with one KV head per query head: n_kv / n_q.

import numpy as np

def kv_head_for(q_head, n_q, n_kv):
    pass

def kv_cache_bytes(layers, n_kv, head_dim, seq, bytes_per=2):
    pass

def grouped_attention(Q, K, V, n_kv):
    pass

def savings(n_q, n_kv):
    pass

Ctrl/⌘ + Enter runs it

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

CHECK YOUR UNDERSTANDING

If four query heads use two KV heads, are there still four queries per position?

AYesBNo

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

One possible pre-norm block is

h=x+\operatorname{Attention}(\operatorname{Norm}(x)),\quad y=h+\operatorname{MLP}(\operatorname{Norm}(h)).

For standard full-width Q/K/V/O projections and an ungated d→f→d MLP, matrix weights total $4d^2+2df$ , excluding bias and normalization. GQA changes K/V projection dimensions.

For one feature vector x of width d, LayerNorm uses $\mu=\frac1d\sum_i x_i$ and $v=\frac1d\sum_i(x_i-\mu)^2$ , then normalizes $(x_i-\mu)/\sqrt{v+\epsilon}$ before its learned affine transform. RMSNorm instead scales using $\sqrt{\frac1d\sum_i x_i^2+\epsilon}$ without the same mean subtraction. Neither definition normalizes across arbitrary axes unless those axes are part of the specified layer.

Worked example

Let d=4 and f=8. Four attention matrices contain 4×4²=64 weights. The two MLP matrices contain 2×4×8=64. The simplified block has 128 matrix weights. Embedding tables, learned normalization scales, biases and any positional parameters are outside this count.

For x=[3,4] without affine transforms or epsilon, LayerNorm produces [-1,1], while RMSNorm divides by sqrt(12.5). The outputs differ because RMSNorm does not center the vector.

Reference code

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

width,hidden=4,8
attention=4*width*width
mlp=2*width*hidden
assert attention == 64 and mlp == 64
assert attention+mlp == 128
query_heads,kv_heads=4,2
assert query_heads//kv_heads == 2
from math import sqrt,isclose
x=[3.,4.]
mean=sum(x)/len(x)
var=sum((v-mean)**2 for v in x)/len(x)
assert [(v-mean)/sqrt(var) for v in x] == [-1.,1.]
rms=sqrt(sum(v*v for v in x)/len(x))
assert isclose(sum((v/rms)**2 for v in x)/len(x),1)

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

Common mistakes

Parameter-count formulas depend on the architecture. RMSNorm and LayerNorm do not perform identical transformations. Smaller KV caches do not mean the model has no attention or no queries.