Transformer block

Deep dive

How it works

In: token ids [B, T] · Out: logits [B, T, V] and, when decoding, the grown KV cache · State: weights (about 12d² per block) and the KV cache (2·L·Hkv·dh values per token) · Cost: 2N + 4·L·T·d forward FLOPs per token; training about 6N + 12·L·T·d

Notation: B batch, T sequence, d model width, H query heads, Hkv KV heads, dh = d/H head width (64 or 128), f MLP width, L layers, V vocabulary.

x [B,T,d] ─┬─ RMSNorm ─┬─ Wq ─ q [B,T,H,dh]   ─ RoPE ─┐
           │           ├─ Wk ─ k [B,T,Hkv,dh] ─ RoPE ─┼─ softmax(q·kᵀ/√dh + causal) · v ─ [B,T,H·dh] ─ Wo ─┐
           │           └─ Wv ─ v [B,T,Hkv,dh] ────────┘                                                  │
           └────────────────────────────────────────────────────────────────────────────────── (+) ◄──────┘
h [B,T,d] ─┬─ RMSNorm ─┬─ W1 [d,f] ─ SiLU ─┐
           │           └─ W3 [d,f] ───────(⊙)─ [B,T,f] ─ W2 [f,d] ─┐
           └──────────────────────────────────────────────── (+) ◄─┘

1. Embed: token ids [B, T] index a table [V, d]. Position enters through RoPE inside attention, not here.
2. Attention branch: RMSNorm, then project to q (H heads) and k, v (Hkv heads). Rotate q and k by position with RoPE. Per head, scores = q·kᵀ/√dh with a causal mask (key j ≤ query i), softmax in fp32, weighted sum of v. Concatenate heads and project with Wo [H·dh, d]. Add to the residual.
3. MLP branch: RMSNorm, then SwiGLU: W2 · (SiLU(x·W1) ⊙ x·W3). With f = 8d/3 it costs the same as a 4d GELU MLP. Llama 2 7B uses f = 11008 for d = 4096.
4. Head: after the last block, a final norm (needed with pre-norm because the residual stream grows with depth), then logits = h·W_U [d, V]. W_U may be tied to the embedding table.

GQA (Hkv < H) shares each KV head among H/Hkv query heads and shrinks the KV cache by that factor. MHA is Hkv = H; MQA is Hkv = 1.

Parameter count (no biases, SwiGLU, untied head):

GPT-2 small (d 768, L 12, 4d GELU MLP, tied 50257-token embedding, 1024 learned positions) gives 12·12·768² + 50257·768 + 1024·768 ≈ 124M.

FLOPs per token, forward (1 multiply-add = 2 FLOPs): every weight matrix costs 2 FLOPs per parameter, so 2N in total. Attention scores add 2·T·d for q·kᵀ and 2·T·d for the weighted sum of v, i.e. 4·T·d per layer counting the full T×T matrix. Backward costs 2× forward. Training is therefore 6N + 12·L·T·d per token, the formula PaLM uses for MFU. Causal kernels that skip masked tiles do about half of the attention term.

KV cache: for decoding, store post-RoPE k and v of every past token. Per layer, K and V each have shape [B, Hkv, T_max, dh]; a decode step computes q [B, H, 1, dh] and attends over the cache. Bytes per token = 2·L·Hkv·dh·bytes_per_value.

Code

One block forward pass for one sequence from scalar primitives, Llama layout (RMSNorm, GQA, RoPE, SwiGLU). The test checks causality: a prefix gives the same outputs as the full sequence.

// Pre-norm decoder block (Llama-style): RMSNorm -> GQA attention with RoPE -> residual,
// RMSNorm -> SwiGLU MLP -> residual. Row-major; x is [t, d]; weight W is [in, out], y = x·W.
pub struct Block {
    pub d: usize, pub n_head: usize, pub n_kv: usize, pub d_ff: usize,
    pub g_attn: Vec<f32>, pub g_mlp: Vec<f32>, // RMSNorm gains [d]
    pub wq: Vec<f32>,                          // [d, n_head*dh]
    pub wk: Vec<f32>, pub wv: Vec<f32>,        // [d, n_kv*dh]
    pub wo: Vec<f32>,                          // [n_head*dh, d]
    pub w1: Vec<f32>, pub w3: Vec<f32>,        // gate, up: [d, d_ff]
    pub w2: Vec<f32>,                          // down: [d_ff, d]
}

fn matmul(x: &[f32], w: &[f32], t: usize, k: usize, n: usize) -> Vec<f32> {
    let mut y = vec![0.0; t * n];
    for i in 0..t {
        for p in 0..k {
            let a = x[i * k + p];
            let (wr, yr) = (&w[p * n..(p + 1) * n], &mut y[i * n..(i + 1) * n]);
            for j in 0..n { yr[j] += a * wr[j]; } // i-p-j order: unit-stride inner loop
        }
    }
    y
}

fn rms_norm(x: &[f32], g: &[f32], d: usize) -> Vec<f32> {
    let mut y = vec![0.0; x.len()];
    for (xr, yr) in x.chunks(d).zip(y.chunks_mut(d)) {
        let ms = xr.iter().map(|v| v * v).sum::<f32>() / d as f32;
        let r = 1.0 / (ms + 1e-5).sqrt();
        for i in 0..d { yr[i] = xr[i] * r * g[i]; }
    }
    y
}

// Rotate pairs (2i, 2i+1) of every head by pos * 10000^(-2i/dh) (interleaved RoPE).
fn rope(x: &mut [f32], t: usize, heads: usize, dh: usize) {
    for pos in 0..t {
        for h in 0..heads {
            let v = &mut x[(pos * heads + h) * dh..][..dh];
            for i in 0..dh / 2 {
                let (s, c) = (pos as f32 * 10000f32.powf(-2.0 * i as f32 / dh as f32)).sin_cos();
                let (a, b) = (v[2 * i], v[2 * i + 1]);
                v[2 * i] = a * c - b * s;
                v[2 * i + 1] = a * s + b * c;
            }
        }
    }
}

fn silu(z: f32) -> f32 { z / (1.0 + (-z).exp()) }

impl Block {
    pub fn forward(&self, x: &[f32], t: usize) -> Vec<f32> {
        let (d, nh, nkv) = (self.d, self.n_head, self.n_kv);
        let dh = d / nh;
        let h = rms_norm(x, &self.g_attn, d);
        let mut q = matmul(&h, &self.wq, t, d, nh * dh);
        let mut k = matmul(&h, &self.wk, t, d, nkv * dh); // k, v after RoPE = what a KV cache stores
        let v = matmul(&h, &self.wv, t, d, nkv * dh);
        rope(&mut q, t, nh, dh);
        rope(&mut k, t, nkv, dh);
        let scale = 1.0 / (dh as f32).sqrt();
        let mut att = vec![0.0; t * nh * dh];
        let mut s = vec![0.0f32; t];
        for hd in 0..nh {
            let kvh = hd / (nh / nkv); // GQA: nh/nkv query heads share one KV head
            for i in 0..t {
                let qi = &q[(i * nh + hd) * dh..][..dh];
                let mut m = f32::NEG_INFINITY;
                for j in 0..=i { // causal mask: only keys 0..=i
                    let kj = &k[(j * nkv + kvh) * dh..][..dh];
                    s[j] = qi.iter().zip(kj).map(|(a, b)| a * b).sum::<f32>() * scale;
                    m = m.max(s[j]);
                }
                let mut z = 0.0;
                for j in 0..=i { s[j] = (s[j] - m).exp(); z += s[j]; } // max-subtracted softmax
                let o = &mut att[(i * nh + hd) * dh..][..dh];
                for j in 0..=i {
                    let (vj, p) = (&v[(j * nkv + kvh) * dh..][..dh], s[j] / z);
                    for e in 0..dh { o[e] += p * vj[e]; }
                }
            }
        }
        let a = matmul(&att, &self.wo, t, nh * dh, d);
        let r: Vec<f32> = x.iter().zip(&a).map(|(x, a)| x + a).collect(); // residual 1
        let h = rms_norm(&r, &self.g_mlp, d);
        let gate = matmul(&h, &self.w1, t, d, self.d_ff);
        let up = matmul(&h, &self.w3, t, d, self.d_ff);
        let act: Vec<f32> = gate.iter().zip(&up).map(|(g, u)| silu(*g) * u).collect();
        let m = matmul(&act, &self.w2, t, self.d_ff, d);
        r.iter().zip(&m).map(|(r, m)| r + m).collect() // residual 2
    }
}

fn main() {
    let (d, nh, nkv, ff, t) = (16, 4, 2, 48, 5);
    let mut seed = 1u64;
    let mut rnd = |n: usize, sd: f32| -> Vec<f32> {
        (0..n).map(|_| { seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
            ((seed >> 40) as f32 / (1u64 << 24) as f32 - 0.5) * 2.0 * sd }).collect()
    };
    let dh = d / nh;
    let b = Block { d, n_head: nh, n_kv: nkv, d_ff: ff, g_attn: vec![1.0; d], g_mlp: vec![1.0; d],
        wq: rnd(d * nh * dh, 0.3), wk: rnd(d * nkv * dh, 0.3), wv: rnd(d * nkv * dh, 0.3),
        wo: rnd(nh * dh * d, 0.3), w1: rnd(d * ff, 0.3), w3: rnd(d * ff, 0.3), w2: rnd(ff * d, 0.3) };
    let x = rnd(t * d, 1.0);
    let full = b.forward(&x, t);
    let prefix = b.forward(&x[..3 * d], 3);
    for i in 0..3 * d { assert!((full[i] - prefix[i]).abs() < 1e-5); } // rows 0..3 ignore tokens 3..5
    println!("ok {:?}", &full[..4]);
}

Complexity and constants

• Weights: 2 bytes/param in bf16, 13.5 GB for a 7B model. Training with Adam and mixed precision holds about 16 bytes/param (see mixed precision).
• Activations stored for backward, per layer, GPT-style block with a 4d GELU MLP: s·b·h·(34 + 5·a·s/h) bytes in 16-bit (s sequence, b batch, h width, a heads; Korthikanti et al.). The 5as/h term is the T×T score matrix; FlashAttention removes it.
• Attention FLOPs exceed matmul FLOPs per token once 4·T·d > 24·d², i.e. T > 6d. For d = 4096 that is about 24k tokens.
• Decode at batch 1 reads all weights per token: 13.5 GB / 3.35 TB/s (H100 SXM HBM) ≈ 4 ms per token lower bound for 7B. Batching amortizes the weight reads; the KV cache reads do not amortize.

Efficiency

• Cost: forward 2N FLOPs per token plus 4·T·d per layer for attention; training about 6N. Batch-1 decode is HBM-bound (about 4 ms per token for 7B on an H100 SXM, above); training is compute-bound.
• Levers:
  • FlashAttention — drops the 5·a·s/h score term of stored activations (Korthikanti et al.) — long sequences
  • GQA — KV cache ÷ (H/Hkv) (Ainslie et al.) — serving
  • fused QKV and W1/W3 GEMMs, bf16 with fp32 accumulate — fewer launches — all GPUs
  • parallel block — about 15% faster training (PaLM) — large models
• Naive → standard → state of the art:
  • post-norm MHA, full score matrix · O(T²) memory
  • pre-norm, RoPE, SwiGLU, GQA · 2N FLOPs per token
  • FlashAttention plus MoE or MLA · FLOPs per token do not grow with experts

Implementation notes

• Precompute RoPE cos/sin as a [T_max, dh/2] table. Two layouts exist: interleaved pairs (Meta's reference code, above) and split halves (rotate_half, Hugging Face). Weights converted between them need the q/k projection rows permuted.
• Fuse Wq/Wk/Wv into one [d, (H + 2Hkv)·dh] GEMM, and W1/W3 into one [d, 2f] GEMM.
• Softmax, norms and the loss in fp32; GEMMs in bf16 with fp32 accumulate.
• Init: normal(0, 0.02) and scale the output projections (Wo, W2) by 1/√(2L) so the residual stream variance stays bounded (GPT-2). See weight initialization.
• Pick d, f, V and dh as multiples of 64 or 128 so GEMMs tile cleanly on tensor cores (Llama 3: V = 128256 = 1002·128).
• KV cache: allocate in fixed blocks per sequence (paged KV cache), not one contiguous T_max buffer per request.

Pitfalls

• Missing final norm in a pre-norm model: logits blow up with depth.
• Softmax over bf16 scores without subtracting the row max: overflow to inf, NaN loss.
• Materializing full logits [B, T, V] in fp32: 8192 tokens × 128k vocab is 4 GB per sequence. Compute the loss in chunks.
• Applying RoPE to v, or caching k before RoPE and rotating again at every step.
• Causal mask built with the wrong diagonal (j < i instead of j ≤ i): the model cannot see the current token.
• Padding tokens not masked in batched prefill: real tokens attend to pad keys.

Variants and alternatives

In production

GPT-2/3, Llama 1-3, Mistral, Qwen, DeepSeek; reference implementations in Meta's llama repo, Hugging Face transformers, llama.cpp, Karpathy's nanoGPT. The same block family serves vision (ViT) and speech recognition (Whisper).

Implementations

Sources

• Vaswani et al., Attention Is All You Need, NeurIPS 2017. paper
• Radford et al., Language Models are Unsupervised Multitask Learners, 2019. paper
• Xiong et al., On Layer Normalization in the Transformer Architecture, ICML 2020. paper
• Zhang and Sennrich, Root Mean Square Layer Normalization, NeurIPS 2019. paper
• Shazeer, GLU Variants Improve Transformer, 2020. paper
• Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding, 2021. paper
• Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, EMNLP 2023. paper
• Chowdhery et al., PaLM: Scaling Language Modeling with Pathways, 2022. paper
• Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models, 2022. paper
• Touvron et al., Llama 2: Open Foundation and Fine-Tuned Chat Models, 2023. paper
• Dosovitskiy et al., An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale, ICLR 2021. paper
• Radford et al., Robust Speech Recognition via Large-Scale Weak Supervision (Whisper), 2022. paper