A Transformer from Raw Tensors: What 112 Tests Taught Me

TL;DR: I built a decoder-only transformer from raw tensors — attention with RoPE, RMSNorm, SwiGLU, an incremental KV-cache, and greedy / temperature / top-k / top-p sampling — wrapped in 112 tests that run on every push. This post is not about the architecture; the Foundations series covers that. It’s about verification: the three properties that are easiest to get silently wrong — cache equivalence, the top-p boundary, causality — and the tests that pin each one. Plus the one bug that did ship: it lived in the only file the test suite didn’t cover.

What I Built

Phase 1 of rlvr-from-scratch is a complete decoder-only transformer written from raw tensors — no nn.Transformer, no Hugging Face, no borrowed attention. Scaled dot-product and multi-head attention with an additive causal mask, positional encoding (sinusoidal, learned, and RoPE), RMSNorm and LayerNorm, GELU and SwiGLU feed-forward variants, an incremental KV-cache, and a generate() loop with greedy, temperature, top-k, and top-p decoding. Every component ships with shape tests, numerical tests, and gradient tests — 112 in total, green in CI on every push.

Everything is public at rlvr-from-scratch, and the theory lives in the Foundations series: Part 1: Attention · Part 2: Positional Encoding · Part 3: Building a Transformer.

The point of “from raw tensors” was never reinvention. It was ownership: when every shape, mask, and gradient is code I wrote, every failure is mine to explain — and every property I claim is one I can prove with a test. What follows is the part the architecture posts don’t tell you: the three places where a transformer can be wrong while looking right, and what it takes to be sure it isn’t.

Three Properties That Fail Silently — and the Test That Pins Each

1. A KV-cache can lie to you

The KV-cache is an optimization with a dangerous failure mode. It changes the compute path — feed one token per step against cached keys and values instead of re-processing the whole sequence — but it must not change the math. If it does, nothing crashes. Generation keeps producing plausible text, just from the wrong distribution. Eyeballing samples will never catch it.

The signature test of Phase 1 makes the equivalence explicit: prime the cache on a prefix, decode the rest one token at a time, and require the per-step logits to match a single full forward over the whole sequence:

assert torch.allclose(incremental, full, atol=1e-5)

The same discipline shaped generate() itself. Additive positional encodings (sinusoidal, learned) have no valid incremental path — position is baked into the embedding at encode time — so caching is only trusted where it’s mathematically sound:

# Additive positional encodings have no incremental path -> recompute.
cache_ok = use_cache and self.pos_emb is None  # RoPE / none only

Rather than serve subtly wrong logits, the model falls back to full recompute. Slower and correct beats fast and quietly wrong.

2. The top-p off-by-one

Nucleus sampling has a boundary condition that buggy implementations get wrong: the token that crosses the threshold p belongs inside the nucleus. The definition is “the smallest prefix whose cumulative probability reaches p” — drop the crossing token and you’ve kept the largest prefix that doesn’t reach p, which is a different (and occasionally empty) set.

The implementation makes the fix visible — after masking cumprobs > top_p, the mask is shifted right by one position, so the boundary token survives and every row keeps at least one candidate:

remove = cumprobs > top_p
remove[..., 1:] = remove[..., :-1].clone()  # keep the token that crosses p
remove[..., 0] = False                      # at least one survivor per row

The test pins it with numbers small enough to check by hand: probabilities [0.50, 0.30, 0.15, 0.05] and p = 0.75 give cumulative sums [0.50, 0.80, 0.95, 1.00], so the correct nucleus is exactly tokens {0, 1}. Two hundred seeded draws must never leave that set — and token 1, the crossing token, must actually be reachable. A companion test checks the renormalized distribution is exactly [0.625, 0.375, 0, 0]: out-of-nucleus probability is zero, not merely small.

3. Causality, checked twice

”The mask blocks the future” is the kind of claim that deserves more than a shape check — a mask can have the right shape and still be applied to the wrong axis, or dropped entirely on some code path. Phase 1 checks causality by two independent routes.

The behavioral route: change the last token of a sequence and require every earlier position’s logits to stay identical at atol=1e-6 — while the last position itself must change. If the causal mask is dropped or mis-shaped, this fails loudly.

The gradient route: an end-to-end finite-difference check, run in float64 through the whole stack, comparing autograd’s analytic gradients against centred numerical estimates on the LM head — plus a test asserting gradients actually reach every stage (embedding, attention projections, FFN, final norm) and are finite and non-zero. Hand-rolled components mean hand-rolled backward risks; this catches the class of bug where the forward pass is right and the learning is quietly broken.

Sampling Trade-offs, Condensed

Decoding strategy is a set of trade-offs, not a single dial. The full reasoning lives in the code’s docstrings and tests; here is the compressed decision table:

One design decision worth naming: filters apply on logits (masking to -inf) before temperature and softmax, so masked tokens get probability exactly zero and the surviving distribution renormalizes for free. And the filters compose — top-k then top-p — with a test covering the combination.

Performance: What the KV-Cache Actually Buys

Measured on the reference config (d_model=256, 4 layers, 4 heads, vocab 1000, RoPE, float32, CPU, seed 0), median of 5 runs:

Without the cache, every step re-processes the whole sequence — O(T²) per generation — so doubling the tokens costs 42% of throughput. With the cache each step processes one token against cached K/V — O(T) — and the same doubling costs only 14%. The speedup widens from 2.6× at 128 tokens to 3.8× at 256, and keeps growing with context.

Now the confession. The one bug that shipped in Phase 1 wasn’t in the model — it was in the benchmark. The first committed version claimed a median over N_RUNS but actually timed each configuration once: rates = [time_once(...)] instead of a loop. The model code had 112 tests; the benchmark had none, and that’s exactly where the bug lived. The fix came with a lesson attached — the benchmark now opens with its own correctness gate before timing anything:

assert torch.equal(ids_nc, ids_c), "cache != no-cache in greedy — bug, not benchmark"

Benchmarks are code too. They need the same review, and ideally the same tests, as the thing they measure.

What Opens Next

Phase 1 is closed and public: the model, the tests, the benchmark, and the three Foundations articles that explain the theory. What comes next is deliberately not decided here. The active plan sequences Transformer Internals → Research Engineering Core → an Eval / Verifiable System with failure analysis — and the next front gets chosen at the monthly checkpoint, not promised in a closing paragraph. Pretraining stays parked, not abandoned.

If you want the full picture — code, phases, and the articles as they ship — it lives on the project page.

@article{sousa2026,
  title={A Transformer from Raw Tensors: What 112 Tests Taught Me},
  author={Sousa, Vitor},
  year={2026},
  url={https://www.vitorsousa.com/blog//}
}