The Transformer as a Computation

I wrote GPT-2’s forward pass out by hand. It is about forty lines of tensor operations: an embedding lookup, twelve repetitions of a block built from layer norm, attention and a two-layer feed-forward network, a final layer norm, and a matrix product. Fed the weights of the real model, it returns logits that agree with the library’s to within \(1.2 \times 10^{-4}\) on logits as large as 323 in magnitude, and picks the same most probable token at every position of four different prompts.

That is the claim of this chapter in one sentence: the transformer is a short, fully specified program, not an opaque object. Once it is a program, every question that matters for inference becomes arithmetic that can be checked: how many parameters, how many floating-point operations per token, how much memory for each intermediate tensor, which parts depend on how long the context is, and which do not.

Chapter 3 ended at the boundary of the model. Text had become ids, ids had become embedding vectors, and at the other end a vector had become scores. This chapter fills in what lies between. We will do it the way the book does everything: write it, measure it, and say precisely what the measurement does and does not establish.

Where this sits

Chapter 1 timed the forward pass as one block and found it was 99.7 per cent of a request. Chapter 2 showed that the model is a bundle, including twelve causal-mask buffers that are not parameters. Chapter 3 showed that a third of GPT-2’s parameters are the embedding table, which is also the output projection. This chapter opens the remaining two thirds.

The Inference Lab piece added here has two parts. The first is a reference implementation: a plain, unfused, readable forward that later chapters use as an oracle. Every optimisation in the book, from caching to fused kernels to quantisation, must agree with it within a stated tolerance, and a reference that cannot be trusted is worse than none. The second is a cost model: functions that return parameter counts and matmul FLOPs from a configuration, so that predictions can be made before anything is run and then compared with measurements.

For the Pi series, the reference implementation is the piece that will matter to the GPU-computing and serving volumes that follow this one. Writing a kernel, or porting a model to another runtime, is only checkable if there is a slow, obviously correct version to compare it against. The same goes for any claim that two runtimes “run the same model”.

The whole computation

GPT-2 small has twelve layers, a model dimension \(d = 768\), twelve attention heads of dimension 64, a feed-forward width of \(4d = 3072\), a vocabulary of 50,257 and a maximum of 1,024 positions. For a batch of \(B\) sequences of \(T\) tokens, the forward pass is:

x = wte[ids] + wpe[0..T-1]                       # [B, T, d]
for layer in 1..12:
    x = x + attention(layer_norm(x))             # [B, T, d]
    x = x + mlp(layer_norm(x))                   # [B, T, d]
x = layer_norm(x)
logits = x @ wte.T                               # [B, T, V]

    flowchart TD
    IDS[token ids] --> EMB[token embedding + position embedding]
    EMB --> X0[residual stream x]
    X0 --> LN1[layer norm]
    LN1 --> ATT[causal self-attention]
    ATT --> ADD1[add]
    X0 --> ADD1
    ADD1 --> LN2[layer norm]
    LN2 --> MLP[feed-forward network]
    MLP --> ADD2[add]
    ADD1 --> ADD2
    ADD2 --> NEXT[next layer, ×12]
    NEXT --> LNF[final layer norm]
    LNF --> OUT[× embedding table transposed]
    OUT --> LOG[logits]
  

Two structural facts matter before any detail. The data passes through the blocks as a single tensor \(x\) of shape \([B, T, d]\), called the residual stream, and each sublayer reads from it and adds to it. And only attention moves information between positions: embedding, layer norm and the feed-forward network all act on each position independently. If you remember one thing from this chapter, it should be that. It is why inference has the shape it has. The positions are independent except where attention joins them, so a token’s computation can reuse earlier positions’ work only through what attention reads from them, which is the observation Chapter 10 turns into the KV cache.

The code lives in examples/ch04/gpt2_from_scratch.py. It uses nothing but matmul, softmax, tanh and arithmetic. One detail of the weights is worth knowing before reading it: GPT-2 stores its linear layers as matrices of shape [in, out], so a layer is computed as x @ W + b, the transpose of the convention PyTorch’s nn.Linear uses.

Embeddings and positions

The first line adds two lookups. The token embedding \(E_{\text{tok}} \in \mathbb{R}^{50257 \times 768}\) is the table from Chapter 3. The position embedding \(E_{\text{pos}} \in \mathbb{R}^{1024 \times 768}\) has one learned row per position, 786,432 parameters:

The residual stream

Each sublayer is wrapped as

The residual stream is also the one place where the magnitude of the numbers is worth looking at, because Chapter 12 will ask whether they survive being stored in fewer bits. I recorded, after each layer, the largest absolute value in the stream for a 17-token prompt:

From layer 3 onwards, one coordinate of the first token’s vector is more than a thousand times the typical value. The root-mean-square of the stream over all positions at layer 4 is 23.8; excluding position 0 it is 2.5. At the last layer the anomaly shrinks, and the two figures are 15.9 and 16.1. This is one prompt on one model, and I make no claim about how general it is. What it shows is that the numerical range of the stream is dominated by a few entries, and any scheme that represents the stream with a uniform scale will either clip them or waste its precision on them. We come back to this in Chapter 12.

Layer normalisation

Layer norm acts on each position’s vector separately, over the feature dimension:

Attention

Attention is where positions meet. From the normalised stream, one matrix multiplication produces queries, keys and values together:

The factor \(1/\sqrt{d_h}\) keeps the scores from growing with head size. If the entries of a query and a key each have unit variance, their dot product over \(d_h\) terms has variance \(d_h\), and without the rescaling the softmax would saturate towards a single position as \(d_h\) grew.

Look at the shape of \(S\). It is the only tensor in the model whose size grows with the square of the sequence length. For \(T = 1024\) and 12 heads in 32-bit floating point it is \(12 \times 1024^2 \times 4 = 50.3\) MB per layer, per sequence, and a naive implementation creates it, masks it, normalises it, and multiplies it by \(V\). Everything else in the layer grows linearly. This is the origin of every attention optimisation in Part III, and the subject of Chapter 8.

What the mask does

The mask is what makes the model usable for generation. Without it a position’s output would depend on tokens to its right, and the computation could not be run one token at a time. I tested that on the real model with a 10-token input, changing only the final token and comparing the logits at the nine earlier positions:

With the mask the earlier outputs are bit-for-bit unchanged, which is the property that allows a prefix to be computed once and reused. Without it, changing the last token rewrote three earlier predictions. The effect on the last position is large in both cases, as it should be: changing its own input is supposed to change its own output.

The exact zero is a stronger result than the tolerance used elsewhere in this book. It holds because masked entries become exactly \(-\infty\), exponentiate to exactly zero, and contribute exactly nothing to any sum. That is a structural guarantee and not a numerical coincidence, and it is the kind of thing a correct optimised kernel must preserve.

The feed-forward network

The second sublayer acts on each position separately:

Counting what is there

The shape information is enough to count parameters without looking at the model. The formula in parameter_count gives, for GPT-2 small:

This agrees exactly with the model’s actual parameter count, so the formula, and with it our understanding of the architecture, is complete. The feed-forward networks hold 45.5 per cent of all parameters, attention 22.8 per cent and the embedding table 31.0 per cent. A count that matches to the last parameter is a strong check. A missing bias vector or an extra layer norm would break it.

Counting what it costs

Parameters are memory. Time is closer to arithmetic. A matrix product of a \([T, m]\) matrix by an \([m, n]\) matrix costs \(2Tmn\) floating-point operations: one multiply and one add per term. So a linear layer costs 2 FLOPs per weight per token, and everything with weights follows that rule. Attention adds a term with no weights: the products \(QK^\top\) and \(AV\) each cost \(2 T d_h\) per head per query token, and across all heads and both products that is \(4Td\) per layer per query token when a token attends over \(T\) positions. For the output of one token at context \(T\):

Three consequences are worth drawing out.

The cost per token is nearly constant in context. It rises by 15 per cent while the context grows by a factor of 64. Most of the arithmetic is in weights that every token touches regardless of how long the prompt is. This matters for the next two chapters, because it means the cost of generating a token is dominated by something the context length barely changes: reading the weights.

The output projection is a third of the work at short context. That is the 77 MFLOP figure from Chapter 1, now identified as 31 per cent of a token’s matmul work. It is also work that a generating loop only needs at the final position, and a naive loop computes it at every position and discards all but one. A runtime that computes only the last row would save up to roughly 30 per cent of the matmul arithmetic of a prefill pass over a short prompt. I have not measured that saving here; it follows from the counts and is a prediction for a later chapter to test.

The 0.25 GFLOP per position estimated in Chapter 1 was right. The exact figure at a 16-token context is 0.248 GFLOP.

Where the time goes

FLOP counts say how much arithmetic there is. They do not say how long it takes. I timed each component of the hand-written forward pass on three sequence lengths. These are medians of five runs after a discarded warm-up, on one thread, in float32 and unfused eager PyTorch:

Three observations, each with a different evidential status.

At 16 tokens the machine delivered about 31 GFLOP/s, almost exactly the 34 GFLOP/s that Chapter 1’s back-of-envelope estimate predicted for 18 tokens. That agreement was not guaranteed, and is the first time in this book that an estimate and a measurement have been compared. Throughput more than doubled at 128 tokens, to 71 GFLOP/s, and did not improve at 512. Short inputs underuse the processor, and by 128 tokens the matmuls are large enough to be bounded by the machine. This is a measurement, with the cause hypothesised from the pattern.

Second, the time shares track the FLOP shares closely at 16 and 128 tokens and diverge at 512. Attention’s share of matmul FLOPs, projections included, is 23.0, 23.7 and 25.8 per cent at the three lengths. Its share of time is 25.6, 24.6 and 34.7 per cent. At 512 tokens attention takes a third more of the runtime than its arithmetic share predicts. The score tensor at that length is \(12 \times 512^2\) values per layer, and masking, softmax and the reshapes touch every element with very little arithmetic per element. My hypothesis is that these operations are limited by memory traffic and not arithmetic, and that the effect grows as \(T^2\). I did not measure memory traffic, so this remains a hypothesis. It is the one Chapters 11 and 13 test.

Third, layer norm accounts for 1.4 to 2.9 per cent of the time while contributing almost nothing to the FLOP count. It is cheap, but it is not free, and the lesson generalises. A FLOP count is a lower bound on cost and a poor predictor of time for operations that do little arithmetic per byte moved.

These shares describe the unfused implementation in this repository, on one thread of one CPU. A production kernel that fuses the attention steps will distribute time differently. The measurement is evidence about where a naive implementation spends its time, which is precisely what an optimiser needs to know first.

Is it the same computation?

A hand-written forward pass is only useful if it is right. The check is direct: run both on the same inputs and compare logits.

The largest logit magnitudes in these four runs were 130, 172, 323 and 149, so a difference of \(10^{-4}\) is about one part in a million or smaller, which is what 32-bit floating-point arithmetic gives when the same sums are computed in a different order. That is the expected level of disagreement between two correct implementations. It is not exactly zero, and shouldn’t be.

What this establishes is narrower than it looks. It shows that on four prompts of 5 to 18 tokens, on a CPU, in float32, with this snapshot, the program above computes the library’s function to rounding error. It does not show agreement at 1,024 tokens, in half precision, on a GPU, or on batches. It also gives me confidence only that I have implemented this model correctly, not that my understanding of transformers in general is right. The tests in the repository add a second kind of evidence: a randomly initialised two-layer library GPT-2 of the same family, small enough to build without a download, is loaded into the same forward function and agrees to \(10^{-4}\), and the causal-mask property holds exactly there too.

The chapter’s earlier example, tiny_transformer.py, is a different kind of object, and it keeps a useful role. It is a randomly initialised model whose block has the same structure but not the same details: it has no bias on the attention projection, and uses the exact GELU. It is a fast fixture for testing shape reports and the causal property, and it makes no claim to reproduce any real model. The hand-written GPT-2 is the oracle; the tiny model is the test double.

What is not canonical

There is no single “transformer block”. Models differ in where normalisation sits, whether they use layer norm or a cheaper RMS norm, how positions are represented, whether the key and value projections are shared between heads, how wide the feed-forward network is and whether it is gated, and which activation function it uses. The block in this chapter is GPT-2’s. An inference engine that claims to run “a transformer” has to say which computation it runs, which is why the config is part of the artefact manifest from Chapter 2. Part IV returns to the variants that change what a runtime can do, and the invariant that holds across all of them is the one this chapter practised: a computation is defined by its operations and shapes, not by a diagram.

What the numbers do not show

• GPT-2 small, float32, CPU, one thread. Larger models shift the balance between components, GPUs shift it again, and lower precision changes the numerics.
• The reference is unfused. Its time shares describe it, not the library’s optimised kernels. I did not time the library for comparison.
• The agreement is on four prompts. It is evidence, not proof. Long contexts and batches are not covered.
• The residual statistics are for one prompt. They are enough to show that large values exist, not how common or how stable they are.
• FLOPs count matrix multiplications only. Softmax, GELU, layer norm and residual additions are excluded, and the timing shows they are not negligible.
• No memory-traffic measurement. The hypothesis that attention’s excess time is memory-bound is unconfirmed here.
• No variants. Rotary embeddings, grouped-query attention, gated MLPs and RMSNorm are described but not implemented.

Inference Lab increment

The Lab now has a trusted reference and a cost model:

• forward(ids, weights, config, causal=True, trace=None), the hand-written GPT-2 pass, with an option to switch the mask off to demonstrate leakage and an option to collect per-layer residual statistics.
• parameter_count(config) and flops_per_token(config, context), which predict memory and arithmetic from a configuration.
• A discipline: every later implementation of the forward pass is compared with forward before its speed is reported. A faster path that disagrees with the reference is a different model.

The request record from Chapters 1 to 3 can now carry something it could not before: not just which artefact ran, but a predicted cost for the request, to be compared with the measured one. That comparison is how Chapter 11 will decide whether a bottleneck is arithmetic or memory.

What you should now be able to do

• Write down the forward pass of a decoder-only transformer with the shape of every tensor, and say which operations mix positions.
• Compute a model’s parameter count from its configuration and reconcile it with the real count.
• Estimate FLOPs per token as a function of context length, and say which terms depend on context.
• Explain why the attention score tensor is the only quadratic term and compute its size.
• Test that a causal mask holds, and say why the result should be exactly zero, not merely small.
• Compare a hand-written forward pass with a library one, say what tolerance is expected and what the agreement does not show.
• Explain why FLOP share and time share can differ.

Key terms

• Residual stream: the \([B, T, d]\) tensor that every sublayer reads and adds to.
• Pre-norm: placing layer normalisation on the input of each sublayer.
• Layer normalisation: per-position normalisation over the feature dimension, with learned scale and shift.
• Query, key, value: the three projections of the normalised stream used by attention.
• Head: one of several parallel attention computations over a slice of the model dimension.
• Causal mask: the \(-\infty\) entries that stop a position attending to later ones.
• Feed-forward network (MLP): the per-position two-layer network with a wider hidden layer.
• FLOP: one floating-point operation; a multiply-accumulate counts as two.
• Oracle: a slow, trusted implementation used to check faster ones.
• Massive activation: an entry of the residual stream far larger than the typical value.

Reader exercises

1. Change n_layer in the tiny model and in parameter_count and show that the formula and the model agree for five different configurations. What do you have to add if you give the attention projection no bias?
2. Remove the \(1/\sqrt{d_h}\) scaling in attention and compare the entropy of the attention weights with and without it on the real model. What do you observe?
3. Replace gelu_new with the exact GELU and measure the change in logits. Is the difference visible in the argmax at any position?
4. Extend the timing command to 1,024 tokens. Does attention’s share of time keep growing as \(T^2\)? What would a memory-bound explanation predict that a compute-bound one would not?
5. Compute the logits only for the last position and report the saving in FLOPs and in measured time at 16, 128 and 512 tokens. Does it match the roughly 30 per cent prediction for short inputs?
6. Run the residual-stream experiment on five other prompts. Is the largest value always at position 0 and dimension 447?
7. Implement RMS normalisation in place of layer norm and describe what in the artefact must change for the model still to be valid.

References

• Attention Is All You Need — defines scaled dot-product attention, multi-head attention, and the original post-norm block.
• Layer Normalization — defines the normalisation used here.
• openai/gpt-2 — the first-party model repository.

Evidence in this chapter

Implementation status

• Manuscript: content/books/inference/04-chapter.md
• Runnable examples: examples/ch04/tiny_transformer.py, examples/ch04/gpt2_from_scratch.py
• Tests: tests/test_ch04_transformer.py, tests/test_ch04_gpt2_from_scratch.py (the real-model commands are not exercised by tests; the hand-written forward pass is checked against a small library GPT-2 built locally)
• Evidence: research/ch04/sources.md, research/ch04/results.json, research/evidence-register.md
• Verification command: python -m unittest discover -s tests -v

The chapter is a draft with draft = false. It makes no claim about GPU behaviour, fused kernels or other architectures. Chapter 5 puts the forward pass inside a loop and asks what generation does with it, one token at a time.