Transformer mechanics and GPU performance, from first principles

A small PyTorch project built to understand both the transformer and the machine it runs on. The transformer is written from PyTorch primitives (nn.Linear and below, no off-the-shelf attention), and numbers are predicted on paper before they are measured.

Two parts:

• train/addition: a 102K-parameter decoder-only transformer trained to add two-digit integers. Character-level tokenization, causal self-attention, masked next-token loss, autoregressive generation. Reaches 99.8% exact match on 1,000 held-out sums.
• benchmark: correct timing of asynchronous GPU work, plus napkin-math calculators for parameter memory, KV cache, FLOPs, memory traffic, and roofline-limited training and decode throughput.

The working method is to write the prediction down first, then measure, then explain the gap. For the capstone I committed predictions before renting an A100 for an hour.

Matmul throughput matched the prediction: 19.1 TFLOP/s FP32 (98% of datasheet peak) and 136 TFLOP/s TF32. Training did not. It took 131s against a predicted 3-6s, and without a written prediction that number would have passed as normal.

The phase breakdown localized the problem. The forward pass took 8.7x longer than the backward pass, so the extra time had to be in code only the forward pass runs. That was one line: a causal mask allocated on the host and copied to the device at each step, forcing a CPU-GPU sync twice per step. Allocating the mask on the device cut median training time from 131.3s to 5.3s. The bug could not have been caught on the local MPS machine, because unified memory makes the copy nearly free.

The other takeaway is about trusting instruments. The per-phase timer distorted the runs it measured: its torch.mps.synchronize() calls changed how MPS scheduled work, and a run with a 3.6s wall-clock time was reported as 8.66s. The check that caught it was re-measuring without the instrument, with a single synchronize at the end of the whole run, and the numbers that depended on the timer were corrected. The full reasoning, including the misses, is in notes/CAPSTONE.md.

Run

uv sync
uv run python -m train.addition.model
uv run python -m benchmark.calculator
uv run python -m benchmark.bench

Python 3.14 is required. Training and GPU benchmarks require a CUDA or MPS device.

Notes

Note	Summary
Timing GPU work	Correctly timing asynchronous GPU work with warmup and synchronization, then checking measurements against physically possible throughput.
Roofline & napkin math	Deriving roofline bounds, transformer memory and FLOP estimates, and the latency-throughput tradeoffs of batched inference.
The addition transformer	Building the addition transformer: tokenization, embeddings, attention, residual blocks, masked loss, held-out evaluation.
Profiling & MFU	Profiling training with parameter and FLOP sanity checks, MFU, Amdahl's law, and compute-optimal scaling arguments.
Predictions vs. an A100	Committing predictions before renting an A100, measuring them against the hardware, and explaining the misses.