A full transformer block·10 min

The attention sublayer

0 / 1

# grade this stage (from your course directory)
npx theinference test --stage attention-sublayer --watch

Wire the first half of a block: normalize the input, run multi-head attention on that clean copy, and add the result back into the original stream. This is the residual river from the last step, made real.

text

x = x + MHA( LayerNorm(x, ln_1.weight, ln_1.bias) )

Shape-preserving sublayer: [T,768] in, [T,768] out

x[T,768] ─ln_1─▶ [T,768] ─MHA─▶ [T,768] ─(+x)─▶ [T,768]

Shape in = shape out = [T, 768]. That invariance is the whole point — it is exactly what lets you stack these sublayers (and whole blocks) 12 deep without ever reshaping anything.

What to build

Your task

Implement runAttnSublayer

Implement runAttnSublayer(x, layer) → x + MHA(LayerNorm(x, ln_1)). Load h.{layer}.ln_1.{weight,bias}; reuse runLayerNorm and runMHA. The residual is an element-wise add of two [T, 768] tensors.

tests/adapters.ts

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

Compose, don't re-implement

Three lines: runLayerNorm(x, ln_1.weight, ln_1.bias) → runMHA(...) → element-wise add to the original x. Load only h.{layer}.ln_1.{weight,bias} here — runMHA loads its own c_attn/c_proj weights. Don't add the result back into the normalized copy; add it into the x you were handed.

How it's graded

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

Pre-LN + residual stream Why residual / pre-LN