An executable, progressively optimized history of the LLM stack — in Zig.
strata is not "another GPT-2 in Zig." It is a layered reference: every stratum
of the modern language-model stack — from a single scalar's derivative up to
grouped-query attention with rotary embeddings and a KV cache — implemented,
tested against a numerical oracle, and explained. The simple versions are never
deleted when the fast versions land above them. You can read the whole thing top
to bottom, or drill from a generated token down to the arithmetic that produced
it.
generated token
→ sampler
→ logits
→ transformer block
→ attention
→ matmul / softmax
→ the scalar chain rule
This repository is the spine of that vision: a GPT-2 architecture reference
and a modern dense
transformer, both trainable and both with cache-based decoders, sharing one
tape-based autograd and one tensor engine, all in pure Zig with no runtime
dependencies. The wider roadmap (GPU/PTX, a tensor IR, MoE, linear attention,
distributed training) attaches to this spine — see ROADMAP.md.
Everything below compiles and is covered by tests that run in CI. Nothing here is a stub. 248 tests (246 passing, 2 skipped). The two PyTorch-oracle tests used to skip whenever their fixture was absent, which meant they had never run once; the fixtures are now committed, and an absent one is a failure rather than a skip.
The 2 skips are the two Intel AMX tests that need the silicon
(src/kernels/cpu/amx_gemm.zig), and they are the only skips in the suite.
The AMX GEMM is written against a Backend enum with a .hardware engine (tile
assembly) and a .software engine — a plain-Zig model of TDPBF16PS that
consumes the same packed buffers and walks the same tile indices. So the
packing, the VNNI interleave, the tile index arithmetic, the k-block walk, the
placement of the four accumulators into C and the shape guards are all
differentially verified against the readable reference on every host,
including ones with no AMX. Only "does the silicon implement TDPBF16PS as
documented" is hardware-gated — and that was verified separately on an Emerald
Rapids host by zig build amx (1.9e-6 against an f32 reference over 256
outputs) before this container was migrated to a machine without AMX. See
docs/EVIDENCE.md and RN-0029.
| Capability | Where | Verified by |
|---|---|---|
Row-major tensor engine, SIMD matmul family (A·B, A·Bᵀ, Aᵀ·B) |
src/core/tensor.zig |
unit tests vs. hand-computed & transposed references |
| Tape reverse-mode autograd with fused transformer ops | src/autograd/tape.zig |
— |
| Backward rules for matmul, LayerNorm, RMSNorm, GELU, SiLU, embedding, cross-entropy, and fused causal attention (MHA/GQA/MQA, optional RoPE) | src/autograd/tape.zig |
finite-difference gradient checks at f64 (gradcheck.zig) |
GPT-2 architecture reference (learned positions, LayerNorm, GELU, biases, tied head) — rungs 1–2 of the ladder in docs/EVIDENCE.md, not "exact GPT-2": no official weights or tokenizer parity yet |
src/models/gpt2.zig |
end-to-end f64 gradcheck; KV-cache ≡ full-forward parity |
| Modern dense transformer (RoPE + RMSNorm + SwiGLU + GQA, bias-free) | src/models/modern.zig |
end-to-end f64 gradcheck; KV-cache parity with pre-rotated keys |
| AdamW (decoupled decay, bias-corrected), SGD, grad clipping, cosine+warmup | src/train/optim.zig |
converges on a quadratic; schedule endpoints |
| Temperature / top-k / top-p sampler (allocation-free) | src/inference/sampler.zig |
empirical frequency test over 20k draws |
| Byte-level BPE: trainer, encoder, decoder, GPT-2 pre-tokenizer & byte map | src/tokenizers/bpe.zig |
lossless round-trip incl. unseen UTF-8; determinism |
| safetensors & GGUF reader/writer | src/formats/ |
exact round-trip; canonical ordering; alignment |
| Capability | Where | Verified by |
|---|---|---|
Persistent fork-join thread pool + dynamic parallelFor |
src/runtime/pool.zig |
exact index coverage; correctness under contention |
| Cache-tiled, register-blocked, multi-threaded GEMM | src/kernels/cpu/gemm.zig |
parity vs. reference on ragged sizes; 178% of ggml like-for-like at 512³ single-thread, 92% parallel efficiency on 4 cores at 2048³ (both machine-attributed — see RN-0030/0031) |
Declarative compile-time quantization (SQ4_32_F16/SQ8_32_F16, f16/f32 scales) |
src/quant/format.zig |
round-trip ≤ ½ quant step; fused dequant-dot ≡ dequant-then-dot; ~3.8×/7× smaller. These are strata's own symmetric block formats — GGML byte-compatibility is not established, so they are not called Q4_0/Q8_0 |
Flash-style attention (online softmax, fwd + bwd, no T² scores) |
src/kernels/cpu/online_softmax_attention.zig |
forward parity vs. materialized softmax; backward vs. finite differences; O(block) memory |
| Mixture of experts (top-k routing, renorm gates, load-balancing loss) | src/nn/moe.zig |
full backward vs. finite differences; aux-loss balancing |
| Muon optimizer (Newton-Schulz orthogonalization) | src/train/muon.zig |
semi-orthogonal singular-value band; matrix-loss descent |
| Compile-time model planner (params, memory, KV cache, static validation) | src/plan/spec.zig |
param counts cross-checked against live GPT-2 & modern models |
| Checkpointing with deterministic restart | src/formats/checkpoint.zig |
resume ≡ uninterrupted run to f32 precision |
| Capability | Where | The claim it proves |
|---|---|---|
| Gated Linear Attention (recurrent + chunkwise) | src/nn/gla.zig |
chunkwise output ≡ recurrence to 1e-9; stable where the naive form underflows; BPTT gradchecked |
| Low-rank KV attention (the algebra under DeepSeek MLA) | src/nn/low_rank_kv_attention.zig |
absorption identity q·(Wᵁᴷc) = (Wᵁᴷq)·c verified numerically; end-to-end gradchecked. The 16× cache figure is arithmetic on the shapes — there is no latent-cache decode path yet |
| Speculative decoding | src/inference/speculative.zig |
lossless: emitted distribution ≡ target p (400k-trial Monte Carlo); greedy output identical to target with a zero-shared-weight draft |
| Sliding window / sinks / QK-norm / logit soft-cap | src/nn/attention_features.zig |
window=∞ ≡ full attention; each feature and all four composed gradchecked |
| LoRA | src/nn/lora.zig |
merge equivalence exactly; zero-init identity to 1e-14; r(in+out) « in·out; gradchecked (A,B only) |
| RoPE relative-position + context scaling | src/nn/rope_scaling.zig |
⟨R_i q, R_j k⟩ depends only on i−j; position interpolation folds extended positions into range |
The gradient checks, parity gates, and equivalence theorems are load-bearing: if a backward rule, a cache index, an optimized kernel, or a frontier mechanism is wrong, an oracle catches it immediately.
zig build test --summary all # 248 tests (2 AMX skips without the hardware)
zig build selfcheck # 17 internal-consistency checks
zig build conformance-gpt2 # tokenizer vs OpenAI's own encoder.py
zig build conformance-llmc # forward + backward vs karpathy/llm.c
zig build conformance-gpt2-official # rung 4: official GPT-2 124M vs onnxruntime
zig build conformance-fla # memory ladder vs the papers' own reference kernels
zig build bench -Doptimize=ReleaseFast # GEMM optimization ladder, measured
zig build plan # compile-time param/memory for GPT-2 & modern models
zig build hybrid # KV memory: KDA/MLA schedule vs all-MLA, by context length
zig build bench-lineage -Doptimize=ReleaseFast # lineage kernels: scalar vs vectorized, measured
./ci/ptx_occupancy.sh # register pressure/occupancy of specialized kernels, per arch
zig build roofline -Doptimize=ReleaseFast # kernels as % of a MEASURED machine bound
./bench/external/compare.sh # like-for-like CPU GEMM vs ggml (llama.cpp)
zig build gemm-nsweep -Doptimize=ReleaseFast # locate GEMM cache-blocking cliffs by shape
zig build gemm-tilesweep -Doptimize=ReleaseFast # sweep the micro-kernel register tile (mr x nr)
zig build gemm-scaling -Doptimize=ReleaseFast # parallel efficiency vs shape, with Amdahl fit per thread count
zig build vecwidth -Doptimize=ReleaseFast # what 128/256/512-bit vectors are worth on YOUR machine
zig build amx -Doptimize=ReleaseFast # Intel AMX from Zig: tile correctness, bf16 peak, ridge shift
zig build lab04 -Doptimize=ReleaseFast # controlled architecture ablationzig build plan computes, entirely at compile time, that GPT-2-small is
124,439,808 parameters — the real number — plus its weight, optimizer, and
KV-cache memory, before allocating anything.
Each lab is a small, permanent, runnable artifact. They are not refactored into the library's abstractions; they stay readable forever.
labs/
00_scalar_autograd/ scalar reverse-mode autodiff from nothing; learns XOR
01_character_gpt/ trains a real modern transformer; loss ~ln(V) → memorized
02_byte_bpe/ trains a BPE tokenizer; shows merges & compression
03_gpt2_forward/ GPT-2 KV-cache parity & determinism gates
04_ablation/ controlled arch ablation: GPT-2 vs modern, GQA sweep
Run them:
zig build lab00 # scalar autograd + XOR
zig build lab01 -Doptimize=ReleaseFast # character GPT training run
zig build lab02 -Doptimize=ReleaseFast # BPE tokenizer
zig build lab03 -Doptimize=ReleaseFast # GPT-2 forward parity
zig build lab04 -Doptimize=ReleaseFast # architecture ablationLab 01 output (abridged) — a genuine converging training loop:
init loss = 3.1126 ln(vocab) = 3.1355 (should match)
step 0 loss 3.1584 ...
step 399 loss 0.3200 ...
generating (greedy) from prompt "to be":
to bethe stio nobler in the mind to suffer ...
The labs teach; the library performs. Same tensor layouts and kernels back both the differentiable training path and the no-grad cache decoder, so there are not two unrelated engines.
src/
core/ tensor.zig rng.zig half.zig
autograd/ tape.zig gradcheck.zig
runtime/ pool.zig # fork-join thread pool
kernels/cpu/ gemm.zig online_softmax_attention.zig # tiled GEMM, online-softmax attention
quant/ format.zig # comptime quantization formats
nn/ moe.zig gla.zig low_rank_kv_attention.zig # experts, linear/low-rank attention
attention_features.zig # window, sinks, QK-norm, soft-cap
lora.zig rope_scaling.zig # adaptation, context extension
tokenizers/ bpe.zig
formats/ safetensors.zig gguf.zig checkpoint.zig
models/ gpt2.zig modern.zig
plan/ spec.zig # comptime model planner
train/ optim.zig muon.zig
inference/ sampler.zig speculative.zig # lossless speculative decoding
Requires Zig 0.16.0.
zig build test # full numerical suite (248 tests)
zig build test --summary all # per-step summary
zig build conformance-gpt2 # external: OpenAI encoder.py fixtures
zig build conformance-llmc # external: llm.c logits + gradients
zig build conformance-gpt2-official # external: official GPT-2 124M weights
zig build bench -Doptimize=ReleaseFast # kernel baselinesThe benchmark harness is deliberately honest: it prints the build mode, warms up, takes the best of N runs, and reports GFLOP/s and effective bandwidth. It never claims "fastest" — it records the baseline that future optimizations are measured against.
These are the rules that make the project a reference rather than a demo:
- Every optimized path has an oracle.
f64scalar reference →f32fast path → (future) SIMD/GPU/quantized. Correctness is checked at every layer; the finite-difference gradient checks are the current top of that chain. - Numerical parity is a gate, not a hope. The KV-cache decoder must equal the full forward pass to ~1e-9 at f64. A failing gate blocks the change.
- The simple versions survive. Labs are frozen. Producing a fast version never means deleting the readable one.
- Pure Zig, honestly. No BLAS, no cuBLAS, no PyTorch in the runtime or the
preprocessing path. Python/PyTorch may exist only as test oracles under
validation/, never as a dependency to train or run a model. - Reproducibility. Everything is seeded; same seed + same code + same precision ⇒ bit-identical results (there is a test for this).
src/models/modern.zig is the same architecture with five changes applied, each of which
the book explains and the code isolates:
| From (GPT-2) | To (modern) |
|---|---|
| Learned absolute positions | Rotary position embeddings (RoPE) |
| LayerNorm (mean+var, bias) | RMSNorm (scale only) |
| GELU MLP (4×) | SwiGLU (gated) |
| Biased projections | Bias-free projections |
| Multi-head attention | Grouped-query attention (GQA) |
Because the fused attention op takes n_head and n_kv_head independently, the
same code path covers full multi-head, GQA, and MQA — and a test confirms the KV
cache shrinks to n_kv_head / n_head of the multi-head size.
A written, executable textbook lives in docs/book/. It is being
built out chapter by chapter alongside the code:
- 00 — Introduction: what this is and how to read it
- 01 — From a scalar to a gradient (the autograd)
- 02 — Tensors and the matmul family
- 03 — Tokenization: byte-level BPE
- 04 — GPT-2: the architecture, and what "exact" would require
- 05 — From GPT-2 to a modern dense transformer
- 06 — CPU kernels: tiling, SIMD, and the parity oracle
- 07 — Quantization as a compile-time format
- 08 — Sparsity and long context: MoE, windows, sinks
- 09 — Training systems: optimizers, schedules, checkpoints
- 10 — Frontier attention: latent, linear, gated
- 11 — Efficient inference: KV caches, speculative decoding
docs/PAPERS.md is the literature map: for every
mechanism in the library, which paper it comes from, what specifically was
taken, and which test gates it — traced backwards from Attention Is All You
Need through its precursors and forwards through the frontier work since. It
also lists, explicitly, what the literature has that this repository does not.
docs/research/ holds the research notes with measured
results (RN-0001 … RN-0008), and docs/decisions/ the
architecture decision records.
See ROADMAP.md for the full multi-stage plan and the honest
scope boundary between what is built and what is designed.
MIT — see LICENSE.