Build your own GPT-2
Multi-head attention·6 min

The full MHA

0 / 1

Chapter 2 built one attention head. A real block runs 12 heads in parallel over different slices of the representation, then mixes their outputs. This chapter wraps your single head into GPT-2's full attention sublayer.

One head, one job

One head can track one kind of relationship. GPT-2 runs 12 heads at once over different 64-d slices of each token — one might follow "the previous token," another "the subject this verb agrees with," another "matching brackets." Then one projection mixes their findings back together.

And it's cheap: 12 heads × 64 = 768, the same width one big head would use. Splitting buys parallel specialization at no extra parameters.

How multi-head attention works

For input x of shape [T, 768], using block layer's weights:

  1. QKV projection — one matmul produces queries, keys, and values together. The 2304 output columns are [ all-q (768) | all-k (768) | all-v (768) ].
  2. Split into 12 heads of dimension 64 (768 / 12). For head h: q_h = qkv[:, h·64 : (h+1)·64], k_h = qkv[:, 768 + h·64 : ...], v_h = qkv[:, 1536 + h·64 : ...].
  3. Per-head attention — run your runAttention(q_h, k_h, v_h) (with the causal mask) to get [T, 64] for each head.
  4. Concatenate the 12 head outputs back to [T, 768] (head h occupies columns h·64 : (h+1)·64).
  5. Output projection — one final matmul mixes the concatenated heads back into [T, 768].

Same Conv1D [in, out] layout as the MLP — x @ W, no transpose.

Fused QKV projection
qkv = x @ c_attn.weight + c_attn.bias        # [T, 2304]   (2304 = 3 × 768)
The 2304 columns: three 768-blocks, not 12 interleaved chunks
c_attn: x[T,768] @ W[768,2304] + b  ->  qkv [T,2304]

columns:  [ --- all 12 heads' Q (768) --- | --- all K (768) --- | --- all V (768) --- ]
          0                              768                   1536                  2304

head h:   Q at  h*64 : h*64+64
          K at  768 + h*64 : ...
          V at  1536 + h*64 : ...
It's q|k|v split, not interleaved per head

The 2304 columns are three contiguous 768-blocks — all Q, then all K, then all V — not 12 interleaved [q, k, v] chunks of 192.

Get the base offsets wrong (0 / 768 / 1536) and every head reads the wrong numbers. Head h's query, key, and value each live h·64 into their own 768-block.

Output projection over the concatenated heads
out = concat @ c_proj.weight + c_proj.bias    # [T, 768]
Down-projectQKV projection