Multi-head attention·12 min

One head from QKV

0 / 2

# grade this stage (from your course directory)
npx theinference test --stage one-head --watch

This is the moment the head you built by hand in Chapter 2 runs on actual GPT-2 weights. No new attention math — just slice head h's 64 columns out of each third of the QKV block and call runAttention.

Head h occupies a 64-wide slice of each third

q_h = qkv[:,    h·64 :    h·64 + 64]
k_h = qkv[:,  768 + h·64 :  768 + h·64 + 64]
v_h = qkv[:, 1536 + h·64 : 1536 + h·64 + 64]

Concretely, for head 0

q0 = qkv[:,    0 :   64]
k0 = qkv[:,  768 :  832]
v0 = qkv[:, 1536 : 1600]

runAttention(q0, k0, v0)  ->  [T, 64]

What to build

Your task

Implement runHead

Implement runHead(x, layer, head): project x to QKV (reuse runQKV), slice head head (64-wide, at the offsets above), run causal attention (reuse runAttention), and return [T, 64].

tests/adapters.ts

export function runHead(x: number[][], layer: number, head: number): number[][] {
  // x: [T, 768]  ->  [T, 64]
}

head_dim = 64, not 768

Each head sees only its 64-channel slice. Passing the full 768 to runAttention would compute one giant head, not GPT-2's 12 — and the shapes wouldn't even match the [T, 64] the grader expects.

Because d = 64 inside the head, the √d scaling in your scores is √64 = 8. Your Chapter-2 code already computes this from q[0].length, so it's automatic — just feed it the 64-wide slice.

How it's graded

Run on block 0, head 0; your [T, 64] output is compared with allclose, atol 1e-4.

Knowledge check

GPT-2 splits the 768-wide representation across 12 heads. What is each head's dimension, and how do the heads divide the channels?

QKV projection Why multiple heads