The Architecture Playground: What a Transformer Config Actually Buys You

A live derivation of the model's forward pass, specialised to the exact config above. Every shape, branch, and substituted value updates as you drag a slider or load a preset.

1. 1

   Input embedding

   $$x_0 = \mathrm{Embed}(t) \in \mathbb{R}^{S \times H}$$

   S = 8,192 tokens, H = 4096 → x₀ has shape 8,192 × 4096
2. 2

   Positional encoding · RoPE

   $$\mathrm{RoPE}(x, pos) = x \cdot e^{\,i\, pos\, \theta}, \quad \theta_j = 10000^{-2j/d_h}$$

   applied to Q and K inside attention — not to x₀. d_h = 128
3. repeated for each layer ℓ = 1 … 32
   3

   Residual block · pre-norm, RMSNorm

   × 32 layers

   3a · pre-norm residual

   $$\begin{aligned} h &= x + \mathrm{Attn}\big(\mathrm{RMSNorm}(x)\big) \\ x' &= h + \mathrm{FFN}\big(\mathrm{RMSNorm}(h)\big) \end{aligned}$$

   normalize inside each residual branch — stable gradients, the modern default

   3b · RMSNorm definition

   $$\mathrm{RMS}(x) = \frac{x}{\sqrt{\frac{1}{H}\sum_i x_i^2 + \epsilon}} \cdot \gamma$$

   drops mean-centering and bias — cheaper, ~same quality; H = 4096
4. 4

   Attention · Grouped-Query (GQA) · Flash

   × 32 layers

   4a · query projection

   $$Q = x\,W_Q \in \mathbb{R}^{S \times H}, \quad W_Q \in \mathbb{R}^{H \times H}$$

   32 heads × d_h 128 = H 4096

   4b · key / value projection · GQA

   $$K = x\,W_K, \quad V = x\,W_V \in \mathbb{R}^{S \times (n_{kv} \cdot d_h)}$$

   GQA: n_kv = 8 heads shared across 4 query groups → K/V width 1024

   4c · RoPE on Q, K

   $$\hat{Q}_i = \mathrm{RoPE}(Q_i, i), \quad \hat{K}_j = \mathrm{RoPE}(K_j, j)$$

   rotate Q and K by absolute position before the dot product

   4d · scaled dot-product

   $$\mathrm{Attn}(Q,K,V) = \mathrm{softmax}\!\left(\frac{Q K^\top}{\sqrt{d_h}}\right) V$$

   d_h = H / n_heads = 4096/32 = 128

   4e · output projection

   $$\mathrm{out} = \mathrm{Attn}(Q,K,V) \cdot W_O, \quad W_O \in \mathbb{R}^{H \times H}$$

   4f · kernel

   $$\textsf{FlashAttention (fused, tiled kernel)}$$

   same math — IO-aware tiling computes softmax in blocks, so the N×N score matrix is never materialised: O(N) memory instead of O(N²)
5. 5

   Feed-forward · dense SwiGLU

   × 32 layers

   $$\mathrm{FFN}(x) = \big(\mathrm{SiLU}(x\,W_{\text{gate}}) \odot x\,W_{\text{up}}\big) W_{\text{down}}$$

   SwiGLU: SiLU/Swish-gated, ⊙ is elementwise

   $$W_{\text{gate}}, W_{\text{up}} \in \mathbb{R}^{H \times F}, \quad W_{\text{down}} \in \mathbb{R}^{F \times H}$$

   3 matrices; H = 4096, F = 14336
   end layer loop
6. 6

   Output / LM head

   $$\mathrm{logits} = x_L\,W_{\text{lm}} \in \mathbb{R}^{S \times V}, \quad W_{\text{lm}} \in \mathbb{R}^{H \times V}$$

   V = 128256 vocab; head 4096 × 128256 = 525.3 M params (often tied to the embedding)

   $$p(t_{i+1} \mid t_{\le i}) = \mathrm{softmax}(\mathrm{logits}_i)$$

   next-token distribution over the vocabulary