Inside Transformer

The walk follows one GPT-2 block from token ids to the next token id, eighty-nine steps in eleven parts. Each part is one line of the forward pass below; the bench keeps every tensor at a fixed home, so the picture at any step is a pure function of the step, the scrubber, and the settings.

X  = wte[ids] + wpe[pos]           # (B, S, d)
Z  = ln_1(X)                       # normalize over d
F  = Z @ W_qkv + b_qkv             # (B, S, 3d)
Q, K, V = split(F)                 # 3 views, no copy
Qh, Kh, Vh = heads(Q, K, V)        # (B, H, S, d_h)
R  = mask(Qh @ Kh.T / sqrt(d_h))   # future keys → −∞
A  = softmax(R, dim=-1)            # rows sum to 1
M  = merge(A @ Vh)                 # (B, S, d) copy
X1 = X + (M @ W_o + b_o)           # residual add
H  = gelu(ln_2(X1) @ W_fc + b_fc)  # (B, S, 4d)
X2 = X1 + (H @ W_proj + b_proj)    # back to width d
logits = ln_f(X2) @ wte.T          # (B, S, V), tied
p  = softmax(logits[:, -1, :])     # last position
next = argmax(p)                   # one per sequence

All three LayerNorms use the population variance (correction=0 in PyTorch), ε = 1e-5, and one learned γ and β per feature, shared by every token. The causal mask writes −∞ into every score whose key comes after its query, so the softmax gives those keys weight 0. The output head reuses the input embedding table read as (d, V): a view, not a second matrix. Generation keeps only the last position of the logits and takes its greedy argmax.

This follows the original GPT-2 model code and the PyTorch LayerNorm definition.

The small example uses B × S × d = 2 × 4 × 8 with two heads of four features, so every drawn entry can be seen; Symbols labels the dimensions with letters, and Numbers shows their sizes. Both keep the same arithmetic. The GPT-2 scales in Display settings use the real widths (d = 768 or 1600, V = 50,257) with two sequences of 128 tokens, and draw a window: four tokens, eight features per home, two heads of four features. Every reduction, contraction, softmax, and argmax runs at full width before the window is chosen, so a running total jumps to the full sum at the last drawn cell. Values come from a deterministic generator, not from trained weights; the two synthetic prompts and every table row are fixed so the same numbers appear on every visit.

The arithmetic ledger counts each scalar add, subtract, multiply, or divide as one basic FLOP: a matmul of rows × k by k × columns costs rows · columns · (2k − 1), and bias and residual adds count separately. Exponentials, tanh calls, reciprocal square roots, and max comparisons are listed by name, never folded into the FLOP count. Replayed operations, such as the second and third LayerNorm, are charged once.

The memory ledger models one forward pass of retained teaching buffers: every tensor the walk keeps on the bench for inspection, at the selected storage precision, with FP32 statistics and int64 ids. Views, broadcasts, and reads allocate nothing; a contiguous copy charges its bytes once. From the value-mixing part on, each part's incoming subtotal continues the previous part's ledger, which starts from the softmax part's owners and omits the original embedding and the parameters of the earlier parts; the ledger's assumptions name that scope. Model parameter totals are shown separately from the one block that is evaluated.

A fused kernel can keep intermediates inside itself instead of allocating them, and changing the storage precision does not simulate rounding. These estimates exclude gradients, optimizer state, allocator overhead, kernel workspace, and the key/value cache of a decode loop, which is not part of this walk. Actual kernel instructions, memory traffic, and runtime depend on the implementation and the hardware.