Decoder Block 1: The CPU Reference

Lesson

A decoder block is not one opaque attention() call. It is a precise sequence of tensor transformations, and this stage makes you write every one of them, in float, at a size where you can print and inspect anything:

n1 = RMSNorm(x, rms1)
Q, K, V = projections(n1)
Q, K = RoPE(Q, K)
attn = causal GQA attention(Q, K, V)
r = x + WO(attn)
n2 = RMSNorm(r, rms2)
ffn = Wdown(SiLU(Wgate(n2)) * Wup(n2))
output = r + ffn

The model shape is fixed for the whole path and scaled for teaching: d_model 512, eight query heads and two KV heads of 64 values each, FFN hidden width 1408. The weights arrive as ONE packed buffer whose offsets are file-scope constants above your function. That is deliberate: packed storage is an interface contract, and stage 3 fuses the three QKV projections into one GEMM precisely because their matrices are adjacent in this buffer. Do not infer offsets; use the constants.

Four contracts in this block reject the obvious shortcut, and the runner checks each one at its own named checkpoint:

- Pre-norm order: RMSNorm comes BEFORE each sublayer, and the residual adds the sublayer's output to the unnormalized stream. Adding before normalizing is a different model.
- RoPE half-split pairing: Each 64-value Q and K head rotates pairs (f, f+32), reading the cosine and sine for token position t and frequency f from the provided tables. Rotating adjacent pairs is a different layout, and rotating V is wrong outright.
- GQA head mapping: Eight query heads share two KV heads: query head h reads KV head h / 4. Nothing is copied; it is an index calculation.
- Causal masking with stable softmax: Row i may only attend to positions j <= i, and the softmax subtracts the row maximum before exponentiating. A one-token sequence, where the single probability must be exactly 1, is among the hidden cases.

The visible case is T = 4 and the goal is not speed. The runner verifies all nine checkpoint buffers against a double-precision reference in dependency order, so the first failure names the first operator that went wrong rather than leaving you to debug a garbled final output.

Advanced project checkpoints: RMSNorm, Q/K/V projections, RoPE, causal GQA attention, output projection/residual, gate/up, SwiGLU, down projection/residual. Validate one checkpoint before implementing the next.

What you will learn

• State the pre-norm decoder-block formula and the order of its nine intermediate tensors.
• Read weights from a packed buffer by documented offset instead of inferring layout from shapes.
• Apply RoPE's half-split rotation to Q and K only, using per-position tables.
• Map eight query heads onto two KV heads with an index calculation rather than a copy.
• Implement causal masking and max-subtracted softmax as correctness requirements.

Starter code

    // The model shape and packed offsets are the file-scope constants above:
    // D=512, HQ=8, HKV=2, DH=64, F=1408, KV=128, RMS_EPS, OFF_RMS1..OFF_WDOWN.
    // Buffers to fill, all row-major: norm1 (T x 512), q (T x 512), k (T x 128),
    // v (T x 128), attnOut (T x 512), residual1 (T x 512), norm2 (T x 512),
    // ffnOut (T x 512), output (T x 512). q and k are checked AFTER RoPE.
    float* scores = static_cast<float*>(std::malloc(static_cast<size_t>(T) * sizeof(float)));
    float* hidden = static_cast<float*>(std::malloc(static_cast<size_t>(F) * sizeof(float)));

    // 1. norm1 = RMSNorm(x) * w[OFF_RMS1..]: divide each row by
    //    sqrt(mean(x*x) + RMS_EPS), then scale elementwise.
    // TODO

    // 2. q = norm1 * WQ^T, k = norm1 * WK^T, v = norm1 * WV^T. Row o of a
    //    row-major (out, 512) weight matrix starts at OFF_W? + o * 512.
    // TODO

    // 3. RoPE: for each Q head (8) and K head (2), rotate the pair
    //    (value f, value f + 32) by ropeCos/ropeSin[t * 32 + f]. Never rotate v.
    // TODO

    // 4. Causal GQA attention into attnOut: query head h uses KV head h / 4.
    //    Row i sees only j <= i: scores = dot(q_i, k_j) * 0.125, subtract the
    //    row max, exponentiate, normalize, then weight v. Use the scores scratch.
    // TODO

    // 5. residual1 = x + attnOut * WO^T.
    // TODO

    // 6. norm2 = RMSNorm(residual1) * w[OFF_RMS2..]; hidden[f] =
    //    SiLU(gate) * up with gate = norm2 * Wgate^T and up = norm2 * Wup^T;
    //    ffnOut = hidden * Wdown^T; output = residual1 + ffnOut.
    // TODO

    std::free(scores);
    std::free(hidden);