Optimizing a Tiny Transformer on CPU

Results

v1 naive · v2 layout · v3 algebra · v4 token-fused GEMM · LP solver-selected (fuse-uv-only) · v6 algebra+vectorization stack. D=128. Full mechanism-per-version detail in the Appendix.
^*v6's 94.0µs beat the prior champion's 97.4µs in the same window, but in a slightly slower machine state than the one that produced the 87.3µs LP-robust record; state-normalized estimate is ~84µs (~2.5× over v1, ~2.9× including the 512-bit vectorization gain already folded in).

End-to-end latency per variant (lower is better); current best (v6) in red. *see state-normalization footnote above.

Decomposition law: speedup = (work ratio) × (rate ratio). The algebra's work cut (v3, extended by v6) is compiler-, ISA-, and memory-regime-invariant; every rate gain from layout, fusion, or vectorization was contingent on machine state — which is why v4 (fusion) is slower than v3 alone at this size, and why the LP had to pick a fusion subset rather than "fuse everything."

Mechanism — why it works

Six algebraic identities, applied offline to the weights rather than at runtime, cut total MACs from 3.96M to 2.11M (a 46.7% reduction) at D=128 — before any layout or scheduling change. v6 adds a seventh, cutting MACs again to 1.80M. This cut is exact, not heuristic: each identity is a theorem about the model, true for any D, any compiler, any machine.

• QK merge: q·k = x_n^t(W_q^T·W_k)x_n — only the product is observable; the K projection disappears, keys become the normalized activations themselves.
• OV merge: W_o·Σa_s(W_v·x_ns) = Σa_s(W_oW_v)x_ns — no nonlinearity between V and O, so they compose. 4 D×D projections/token → 1 fused one at half the size.
• Layernorm folding: γ, β feed only linear maps → absorb offline. Runtime LN is just (x−mean)·rstd.
• Softmax shift-invariance deletes the query-side bias cross-terms outright; softmax-sums-to-one turns the value bias into a plain bias — two "free" deletions that exist only in the math.
• Last-token pruning: only logits[last] are read → final layer's MLP/attention for tokens 0..8 is dead code no compiler can see. Decays as 1/L with depth — but is exactly incremental decoding: we accidentally derived the KV cache.
• OV-through-the-sum (v6): the OV-merged projection W_g is linear, so Σ_sa_s(W_g·x̂_s) = W_g(Σ_sa_s·x̂_s) — sum the raw, attention-weighted activations first, project once. One D×D GEMV per computed output row instead of one per token. With last-token pruning, layer 2's single pruned row drops from 20D² to 2D² MACs.
• Caveat filed: QK/OV merges invert into a pessimization for multi-head with H ≥ 3 (rank D/H products vs dense D²). Single-head only, or keep factored.

Worked example: the QK merge

Standard attention computes k = W_kx_n (a D×D matvec) for every token, then scores q·k. But k is never read anywhere else — only the scalar q·k is. Substituting q = W_qx_n:

q·k = (W_qx_n)·(W_kx_n) = x_n^T(W_q^TW_k)x_n

M = W_q^TW_k is a constant D×D matrix, computable once, offline. Runtime cost per token drops from two D×D matvecs (q and k) to one (Mx_n) plus a dot product — and the K projection has disappeared entirely; the "key" is just the normalized activation itself. This substitution and its OV-merge counterpart are most of the 46.7% MAC reduction above. v6's OV-through-the-sum identity is the same move applied one level up: defer a linear projection past a summation, and it runs once instead of once per token. Because v6 shares Wu/Wg and biases with v3, both paths must agree bit-for-bit — measured agreement: 4.6e-7.

Fusion speedup (v4 ÷ v3) by regime. Below the 1.0× line, fusion hurts — as it does under frontend contention with the unrolled build (red).

Regime dependence: fusion is worth 0.92× at D=128 (L2-resident, i.e. fusion is slightly counterproductive), 1.52× at D=512, 1.88× at D=2048 (DRAM-bound), and 0.33× under frontend contention with the unrolled build. The algebra's ~1.9× held in every regime — fusion did not.

Why the algebra survives the memory wall: in matvec, FLOPs and bytes are locked in proportion (2 FLOPs per 4 weight bytes), so every FLOP deleted was a byte deleted. Measured byte-ratio prediction at D=2048: 1.88×; measured speedup: 1.90×.

The floor: what's left after algebra

v6's 1.80M MACs = 3.60M FLOPs ÷ 99 GF/s peak vector throughput = a 36.4µs floor. Add the cost of reductions and short spans that can't reach peak rate, and the honest floor is ~42–45µs — meaning v6, at 94µs, is sitting at roughly 2× over its own floor, and the gap is no longer GEMVs. It's twenty layernorms, sixty-five score dots, and softmax — all running at serial-chain rates on a machine whose vector unit sits idle while they compute. That's the thesis this ladder ends on: every projection — the part everyone optimizes — is within sight of physics, and what's left is the connective tissue. The next 1.5× isn't in a better matvec; it's in restructuring the small stuff, or in the plainer conclusion that at 94µs on one noisy rented hardware thread, the cheapest remaining optimization is a machine that's actually yours.

Method — how it was found

Result: the robust, footprint-aware model (LP-2) selected fuse-uv-only — a config nobody had tried — and it turned out to be the measured optimum of all 16 configs in both machine states (contended: 122.2µs, predicted 121.6 — 0.5% error; quiet: 87.3µs). Rank correlation model-vs-measured: ρ = 0.906. Regret dropped from LP-1's 3.11× to 1.00×. It got there in two iterations:

LP-1 (compute rate + L3 BW + DRAM BW rows, calibrated by measurement): instantly re-derived every dominant algebraic choice (merge, prune — exact linear coefficients, they're theorems). Picked all-fuse everywhere. Scorecard: right at D=2048 (1.45×), wrong at 128/512 (0.34×/0.65×) — its pick ranked 15th of 16 measured configs at D=128. Regret 3.11×.

Root cause of the miss (a debugging saga in itself): the fused kernel ran 61–65 GF/s standalone and 14 GF/s in context — same shapes, same minute. The asm: 8 FMA instructions standalone vs 1006 in the composed binary — unrolling blew the hot code past the µop cache, a resource with no row in the model. And the penalty was time-varying: identical binaries swung 98µs ↔ 382µs across windows while tiny calibration kernels never moved — an unseen SMT sibling / co-tenant contending for the shared frontend. (Also: uptime 7 min mid-session — the VM had been silently re-placed. Some of "the machine changed" was literally a different machine.)

LP-2 fixed it with three additions: (a) a code-footprint constraint row — the footprint turned out to be exactly additive in the decision variables (base 526 FMAs; marginals uv −24, w1 +248, w2 +256), so it's legitimately linear; (b) a two-scenario robust objective (quiet / contended frontend) instead of a point estimate; (c) steady-state memory accounting (LP-1 charged cold-pass DRAM to hot loops — a bug). uv's footprint marginal is negative — its gemm replaced a larger per-token loop — so fuse-uv-only banks the 64 GF/s fused rate without crossing the contention cliff.

The formalism verdict: LP/roofline is exact for the algebra (FLOPs/bytes are theorems) and for regime identification (duals = which constraint binds); cut/partition (hypergraph — weights-as-nets, VLSI-style) is the right structure for fusion and device placement (CPU/GPU split is a literal poly-time s-t mincut); measurement owns everything contextual, compositional, or non-stationary — which this machine demonstrated four separate times is not a small residue. Trust models for 10× decisions, the stopwatch for 2× decisions.

Footnote — the LP doesn't survive v6 unchanged: fuse-uv-only was a statement about a kernel shape that no longer exists. v6's u projection is half-width and computed only for rows that need it, so the footprint marginals LP-2 solved against are stale; fusion would need re-deriving from v6's kernel set. At the current ~4µs margins between candidates, that's stopwatch territory, not model territory — the LP earns its keep at 10× decisions, not at the last few percent.

Scaling — how this generalizes

At this model size, a GPU already loses to a pinned CPU core: 30–50 kernel launches (≈200µs) plus PCIe round-trip exceed the 87µs CPU compute. GPUs only win once per-invocation work exceeds roughly 1ms or the working set demands HBM bandwidth (detail in the Appendix). For a CPU-bound tiny model, the remaining levers are batching, quantization, and core replication:

• Splitting one pass across cores has a size threshold: ~14 barriers × 1–5µs against an 87µs pass caps 8 cores below 2×. Worth it only when phase work ≫ 50× barrier cost, i.e. D ≳ 1024.
• Replication (one stream per pinned core, weights shared read-only) scales linearly until a shared resource saturates. At D=2048: unfused = 2.15 GB/pass → ~230 passes/s ceiling at ~500 GB/s, saturated by ~34 cores. Fused: 0.34 GB → ~1,470/s (6.4×). Fused + batch-8: ~11,000 inf/s (48×).
• Currency change: with many cores, bytes-per-pass (shared) replaces seconds-per-pass (private) as the objective. Fusion's 1.5× single-core value becomes 6.4× socket value. Shadow prices migrate with core count.
• Ordering: batch first, then replicate. Batching attacks the shared constraint; cores only the private one; they compose multiplicatively.
• SMT becomes a design variable: our contended-state data (small-footprint build loses 1.3× under frontend sharing, big build 3.5×) implies two compact streams per core (~1.54×) beat one unrolled stream with an idle sibling. The robust LP's "adversarial scenario" becomes a deliberate operating point.
• Parked for different hardware: v6's reduction-splitting fix (14 independent chains) stays parked on this box — it's the first thing to flip on a two-pipe Zen 5, where more independent chains map directly onto more FMA ports rather than more vector width.

Design projection, not measurement — the test box had 1 vCPU.

Appendix — the long version

The arc

This campaign started as a question about whether the OS process model limits how much compute you can push through a CPU. It turned into six generations of implementation on a tiny transformer (SEQ=10, D=100→128, FF=4D, 2 layers, single head), each measured exhaustively, modeled with an integer linear program, projected to multicore, and — after the LP's own kernel set was outrun — extended once more by hand.

Discoveries & surprises, ranked

1. UB deletes programs, not data — twice, silently, at -O3. In v1, load_weights never wrote ln_f_g/b; IPA proved them zero and constant-folded the entire final projection to memset(logits, 0, 40). In v2, float acc[DP] called with out_dim=FF overflowed the stack; the resulting UB let the compiler delete the loop exit and everything downstream as unreachable, leaving an orphaned infinite loop in the asm. ASan/UBSan belongs in the toolchain, not just the test suite.
2. The measured global optimum was a config no human proposed — LP-2 found fuse-uv-only via a negative footprint coefficient the solver could see and we couldn't.
3. A single divisibility property was worth 7×. Strips of 112 at od=112/224/448 ran at 36 GF/s; od=128 (112+16 ragged) dropped to 5.3 GF/s. Mixed strip widths defeat register promotion — the accumulator array stays in memory (14 loads + 7 FMAs + 7 stores per iteration instead of 7 register FMAs). Strip width must divide every output dimension.
4. Standalone kernel benchmarks lied by 4.4× about in-context performance. The fused kernel ran 61–65 GF/s standalone and 14 GF/s in context — 8 FMA instructions standalone vs 1,006 in the composed binary, unrolling past the µop cache. Calibration ≠ context.
5. Identical binaries varied 4.3× across time on shared infrastructure while calibration microbenchmarks held steady (98µs ↔ 382µs across windows) — the missing constraint was someone else's process on the same core.
6. gcc quietly performed our hand optimization. gcc auto-vectorized v1's reductions itself, making v1 ≈ v2 under gcc even though they differ 40% under clang. The "v1 is latency-bound" diagnosis was a property of the compiler, not the source.
7. The L3 cache was already doing implicit temporal fusion. v3's per-token weight re-streams at D=2048 ran at ~47 GB/s effective — L3 speed, not DRAM — because back-to-back streams of a matrix smaller than L3 hit cache. Expected 10× from explicit fusion, measured 1.88×; the gap is the cache hierarchy already doing the optimization.
8. Row-form vs. column-form matvec is the whole codegen story. Dot products give a serial FMA chain plus a 7-instruction horizontal-reduction ladder per output; axpy form gives persistent register accumulators and zero reductions. Reductions across vector lanes are poison; reductions across loop iterations into vector accumulators are free.
9. Inliner roulette: static inline was honored in one translation unit and ignored in another. Without inlining, no constant propagation, no strip specialization → 6× slower. Use always_inline when correctness-of-performance depends on it.
10. v4 reached the actual memory roofline — 1.34 GB of weights per pass at ~14 GB/s sustained, the machine's single-core streaming rate. Past that point "more ops through the CPU" is the wrong question; only quantization, batching, or cores move the needle.
11. GPUs lose at this model size. 30–50 kernel launches ≈ 200µs vs 87µs on CPU, and PCIe round-trip ≈ the whole computation. GPUs win at ≳1ms of work per invocation or when the working set wants HBM bandwidth. Tiny latency-critical models belong on a pinned CPU core (or FPGA).
12. The math never broke. Through deleted programs, inliner roulette, VM swaps, and a solver confidently choosing a 3× pessimization, relative error held at ~1e-7 in every single run. Exact reasoning above, empiricism below, hard interface between.
13. Belt and suspenders beats borrowed permission (v6). Explicit 4-way partial sums in every reduction put the parallelism in the source instead of relying on -ffast-math's license to reassociate. Same result, but now it's guaranteed by the code, not by a compiler flag someone could remove.

Operating principles distilled

• Algebra first (portable, provable), layout second, schedule last — each layer invisible to the tools below it.
• Model proposes, stopwatch disposes; every accepted move changes the binding constraint, so iterate.
• Only compare numbers from the same window / same binary pairing. Anything else is an anecdote.
• On shared infrastructure, optimize the regret distribution, not the point estimate.
• When prediction and measurement disagree, the disagreement is the next bug — and it's usually where the 2× is hiding.

Artifacts

tiny_transformer.c tiny_transformer_v1_fixed.c tiny_transformer_v2.c tiny_transformer_v2_fixed.c tiny_transformer_v3.c v6.c bench.c scaling.c opt.c calib.c lp.py lp2.py

tiny_transformer.c is v1 (naive), tiny_transformer_v2.c/tiny_transformer_v3.c are the layout and algebra rewrites. The _fixed variants carry the UB fixes described in the Appendix's first discovery. v6.c is the current champion — v3's algebra plus OV-through-the-sum, explicit partial reductions, and 512-bit vectorization, cross-checked numerically against v3 since both share the same trained weights. bench.c is the 3-way harness with numeric cross-check, scaling.c is the D-sweep, opt.c is the macro-parametrized fusion build (v4 / LP-driven), calib.c measures machine constants, and lp.py / lp2.py are the two LP iterations described above.