Problem 107 · Attention & Inference

Multi-Head and Grouped-Query Attention Shapes

MediumAI engineering

01Problem

Implement multi-head self-attention on lists: project `X` (T x d_model) with `Wq`, `Wk`, `Wv` (d_model x d_model for queries; d_model x (n_kv_heads * head_dim) for keys and values), split into heads, run causal attention per head, concatenate, and project with `Wo`. Support grouped-query attention, where `n_heads` query heads share `n_kv_heads` key/value heads (`n_heads % n_kv_heads == 0`). Also write `kv_cache_bytes(layers, n_kv_heads, head_dim, seq_len, batch, bytes_per_value)`.

02Key insight

Heads are a reshape, not extra parameters: d_model is split into n_heads slices of head_dim, each attending independently, and Wo mixes them back. GQA cuts the K and V projections (and the KV cache) by n_heads / n_kv_heads while keeping the query heads, which is why it is the default in modern open models: decode is memory-bandwidth bound, and the cache is what you stream every step. The cache formula (2 for K and V, times layers, kv heads, head_dim, tokens, batch, bytes) is the number you size GPUs with.

03Solution

solution.pypython47 lines

import math

Matrix = list[list[float]]


def matmul(A: Matrix, B: Matrix) -> Matrix:
    cols = list(zip(*B))
    return [[sum(a * b for a, b in zip(row, col)) for col in cols] for row in A]


def _causal_attention(Q: Matrix, K: Matrix, V: Matrix) -> Matrix:
    scale = 1.0 / math.sqrt(len(Q[0]))
    out: Matrix = []
    for i, q in enumerate(Q):
        scores = [sum(a * b for a, b in zip(q, K[j])) * scale for j in range(i + 1)]
        m = max(scores)
        w = [math.exp(s - m) for s in scores]
        z = sum(w)
        out.append([sum(w[j] * V[j][c] for j in range(i + 1)) / z for c in range(len(V[0]))])
    return out


def _head(M: Matrix, h: int, head_dim: int) -> Matrix:
    return [row[h * head_dim:(h + 1) * head_dim] for row in M]


def multi_head_attention(X: Matrix, Wq: Matrix, Wk: Matrix, Wv: Matrix, Wo: Matrix,
                         n_heads: int, n_kv_heads: int | None = None) -> Matrix:
    d_model = len(X[0])
    n_kv = n_kv_heads or n_heads
    if d_model % n_heads or n_heads % n_kv:
        raise ValueError("d_model % n_heads and n_heads % n_kv_heads must be 0")
    head_dim = d_model // n_heads
    if len(Wk[0]) != n_kv * head_dim or len(Wv[0]) != n_kv * head_dim:
        raise ValueError("Wk/Wv must project to n_kv_heads * head_dim")
    Q, K, V = matmul(X, Wq), matmul(X, Wk), matmul(X, Wv)
    group = n_heads // n_kv
    heads = [
        _causal_attention(_head(Q, h, head_dim), _head(K, h // group, head_dim), _head(V, h // group, head_dim))
        for h in range(n_heads)
    ]
    concat = [sum((heads[h][t] for h in range(n_heads)), []) for t in range(len(X))]
    return matmul(concat, Wo)


def kv_cache_bytes(layers: int, n_kv_heads: int, head_dim: int, seq_len: int, batch: int, bytes_per_value: int) -> int:
    return 2 * layers * n_kv_heads * head_dim * seq_len * batch * bytes_per_value

Try it first.

TimeO(T * d_model^2 + n_heads * T^2 * head_dim)

SpaceO(T * d_model + T * n_kv_heads * head_dim)

tests.py · 7 passing

tests.pypython23 lines

import random
rng = random.Random(0)
rand = lambda r, c: [[rng.uniform(-1, 1) for _ in range(c)] for _ in range(r)]
T, d = 5, 8
X = rand(T, d)
Wq, Wk, Wv, Wo = rand(d, d), rand(d, d), rand(d, d), rand(d, d)
out = multi_head_attention(X, Wq, Wk, Wv, Wo, n_heads=4)
assert len(out) == T and all(len(r) == d for r in out)
one = multi_head_attention(X, Wq, Wk, Wv, Wo, n_heads=1)
ref = matmul(_causal_attention(matmul(X, Wq), matmul(X, Wk), matmul(X, Wv)), Wo)
assert all(abs(a - b) < 1e-9 for r1, r2 in zip(one, ref) for a, b in zip(r1, r2))
assert multi_head_attention(X, Wq, Wk, Wv, Wo, 4, 4) == out
Wk2, Wv2 = rand(d, 4), rand(d, 4)
gqa = multi_head_attention(X, Wq, Wk2, Wv2, Wo, n_heads=4, n_kv_heads=2)
assert len(gqa) == T and len(gqa[0]) == d
for bad in [dict(n_heads=3), dict(n_heads=4, n_kv_heads=3)]:
    try:
        multi_head_attention(X, Wq, Wk, Wv, Wo, **bad)
        assert False
    except ValueError:
        pass
assert kv_cache_bytes(32, 8, 128, 8192, 1, 2) == 2 ** 30
assert kv_cache_bytes(32, 32, 128, 8192, 1, 2) == 4 * 2 ** 30

04Staff-level follow-up

Why not go all the way to one KV head (MQA)? Quality drops measurably; GQA with 8 groups recovers most of it. Then discuss how tensor parallelism shards heads across GPUs (n_kv_heads must divide or be replicated across the TP degree), and how MLA-style latent compression shrinks the cache further at the cost of extra compute.