I built the whole LLM stack from scratch in Rust

I'm a high school student in Germany. Over two and a half weeks this summer I built every layer of a small LLM by hand, in Rust, with no ML framework anywhere in the stack:

• Cadmus — a byte-level BPE tokenizer trainer (learning the merges, not just applying them). ~500 lines.
• Hephaistos — a Llama-style transformer trainer with hand-written forward and backward passes. No autograd, no tensor library, trained on CPU. ~2.8k lines.
• Talos — an inference engine: GGUF loader, KV cache, quantization, sampling, plus a Metal GPU backend with hand-written compute shaders. ~3.6k lines.
• Talos Forge — Talos compiled to wasm, running the trained model in the browser. That's the demo link above; the model generating text on that page went through zero external ML code on its way there.
• Mnemosyne — RAG on top of all of it: a contrastively fine-tuned embedding model, a vector index, and retrieval-augmented generation. Also in the browser (the "Mnemosyne" toggle in the demo). ~1.2k lines.

The rule I set myself: if I can't write it and test it against a ground truth, I don't get to use it. Allowed: rayon for threads, memmap2, an RNG, JSON. Not allowed: autograd, ggml, candle, HuggingFace tokenizers, or any kernel I didn't write.

I did this because I wanted to actually understand how this stuff works — not "I read the paper," but "I can make the gradients match to eight decimal places." What follows is what each layer taught me, and the numbers.

Part 1: Backprop by hand needs a truth machine (Hephaistos)

A "loss goes down" curve does not prove your gradients are right. Subtly wrong gradients often still learn — just worse — and you'll never know. So before trusting a single line of backward code, I built a gradient checker: every analytic gradient is compared against an f64 central-difference estimate. The forward and backward passes are generic over the float type, so the checker re-runs the exact same code in f64 — round-off can't hide a bug behind a precision difference.

The model is a genuine Llama block, not a GPT-2 toy: RMSNorm, RoPE, SwiGLU, AdamW with decoupled weight decay, dropout, all biases dropped the way Llama drops them. Everything runs on flat Vec<f32> buffers. A cache-blocked matmul with an 8-lane dot product lets the compiler emit SIMD, parallelized with rayon over rows and batch — while staying bit-identical to the naive reference loop (that's a test, too).

The part that makes it real rather than a sandbox: the trained checkpoint exports to GGUF (RoPE weights permuted to GGUF's interleaved layout, tokenizer embedded), so the result loads in llama.cpp and Ollama. If your file format, weight layout, and tokenizer conventions are all correct, someone else's engine will speak your model's words. That's a much stricter test than your own code agreeing with itself.

Reference check: a ~15.7M-param training run lands within ~1% total-corpus NLL of the same architecture trained in MLX on the same data.

Part 2: The last outsourced link (Cadmus)

Hephaistos originally borrowed one artifact: tokenizer merges learned by HuggingFace's tokenizers crate. Encoding and decoding were hand-written, but the training of the merge table wasn't. Cadmus closes that gap: GPT-2-style byte-level pre-tokenization (every byte maps to a printable char, so decode(encode(s)) == s for arbitrary UTF-8), count all adjacent symbol pairs over the corpus, fuse the most frequent, repeat until the vocab is full. The merge list's order is the encoder's priority table.

Small (~500 lines), but it means the capstone run — fresh tokenizer, fresh token stream, full training, GGUF out — touches no external ML code at any point. Bonus: on my corpus, Cadmus's merges compress ~5% better than the reference tokenizer I'd been borrowing.

Part 3: Inference is a different discipline (Talos)

Training taught me the math; inference taught me that the hard part is conventions. A GGUF reader that mmaps the file and hands out zero-copy tensor views. RoPE in the interleaved-pair convention matching the trainer's export. Grouped-query attention. A per-layer KV cache. Q8_0 and Q4_0 quantization with the dequant fused into the matvec. Temperature / top-k / top-p sampling. A teacher-forced perplexity harness, because generated samples "looking coherent" is not a metric.

The contract that kept it honest: Talos's logits must match Hephaistos's within 1e-4 for the same model and prompt. That's a test in CI, not a vibe. Every convention mismatch — RoPE layout, norm epsilon, quantization rounding — shows up as a parity failure before it can become a mystery.

Part 4: The GPU doesn't reward effort, it rewards architecture (Metal)

This is the part I'm proudest of, because each step is separately measured and honest about what it did not buy.

Step 1 — port the matvec to a Metal shader. One thread per output row, fused dequant for Q8_0/Q4_0 (weights stay quantized in GPU memory, decoded inside the shader byte-for-byte like the CPU path). Correct — and slow.

Step 2 — stop re-uploading the model. The naive version shipped each weight matrix to the GPU every token. Keeping tensors resident (upload once, keyed by name): 6.3× on a 4096×4096 matvec.

Step 3 — a kernel that actually uses the hardware. One simdgroup per output row, lanes striding the row so reads coalesce, simd_sum to reduce: 0.79 ms per matvec. Honest caveat: still 1.6× slower than my multithreaded CPU path, because every matvec was its own command buffer and the commit/wait overhead ate the win.

Step 4 — put the whole forward pass in one command buffer. Residual stream and KV cache live on the GPU; every op per token (rmsnorm, q/k/v, RoPE, attention, SwiGLU, output projection) is encoded into a single serial compute encoder; only the logits come back. The serial encoder gives memory coherence between kernels for free — race-freedom is structural, not hoped for. That amortizes the dispatch overhead, and the GPU finally wins end-to-end:

Q4_0 wins most because decode is bandwidth-bound — fewer bytes to move. The lesson that took me the longest: the GPU didn't get faster because my kernels got cleverer, it got faster when I stopped talking to it so often.

Correctness under all this: every kernel is tested against its CPU twin on random inputs with a relative tolerance (GPU and CPU sum in different orders — bit-equality is the wrong bar); the full GPU forward matches the CPU forward to ~1e-7 over a multi-token GQA sequence; and two GPU runs must be bit-identical, because a data race would show up as nondeterminism.

Part 5: The browser (wasm)

Talos compiles to wasm32, with a fused simd128 Q4_0 dequant-dot kernel (4.6× over the scalar wasm build on Q4_0). Talos Forge wraps it in a UI and runs the Hephaistos-trained model in a worker, client-side, no server. The wasm build is held to the same standard as everything else: same seed, same generations as the reference build.

Part 6: Teaching the stack to remember (Mnemosyne)

A language model's weights are its only memory, and at 15.7M params that memory is tiny — ask it "Who is Proserpina?" and it hallucinates freely. RAG fixes this by looking the answer up at generation time. I wanted the whole retrieval side from scratch too: the embedding model, its training objective, the index, the evaluation.

Embeddings for free, then embeddings for real. A passage embedding here is the mean over positions of the final-norm hidden state — the residual stream right before the lm_head — L2-normalized. The encoder stays causal, because Talos's KV-cache forward is causal by construction; a bidirectional encoder would have needed a whole second inference path for marginal gain. Stage 1 uses the pretrained LM unchanged and just pools its hidden states. Stage 2 fine-tunes those weights with a contrastive objective (InfoNCE, Contriever-style: positives are two random crops of the same passage, negatives are the other pairs in the batch, symmetric loss, τ = 0.05). Stage 3 makes the objective match the task: retrieval is a short query finding its long passage, so the final fine-tune trains two weight-sharing towers — a 16-token query sampled inside its 64-token passage — which needs nothing more than summing the two towers' gradients before the optimizer step. The gradient enters through a new backward entry point seeded directly at dL/d(hidden) — the lm_head never sees it — and the same f64 gradient checker that verified the language-model loss verifies the whole contrastive chain.

The numbers (retrieval over 34k overlapping 64-token chunks of the training corpus; queries are random 24-token crops; a hit = a retrieved chunk covering ≥50% of the crop):

Three things surprised me. First, the raw language model is already a passable embedding space — 0.495 recall@1 comes free with next-token pretraining. Second, how cheap the upgrades are: ~40 minutes of contrastive tuning on a CPU nearly doubles recall@1, and quantizing the encoder to Q4_0 with an f16 index costs almost nothing. Third, how much the shape of the training pairs matters: switching positives from "two crops of one window" to "short query inside its passage" — the thing retrieval actually does — was worth more (0.730 → 0.910) than the entire first fine-tune, and short 12-token queries went from broken (0.110) to usable (0.290). Honesty footnote: on this eval a plain token-overlap baseline scores 0.995 — verbatim crops are trivially findable by exact matching. That row is a ceiling by construction; embeddings have to win on paraphrase, and qualitatively they do ("Who is Proserpina?" now retrieves Ceres searching for her lost child; the symmetric encoder's top hit for the same query was an off-topic passage that merely sounded like a question).

The truth machine kept paying rent here, too: the gradcheck for the weight-sharing towers refused to pass, and the culprit was a real loss bug — my -ln(max(p, 1e-8)) clamp flattens the loss once the softmax saturates, so the analytic gradient disagreed with the numerical one exactly there. Computing the cross-entropy as logsumexp(sim) − sim[diag] is exact and smooth, no clamp needed. "Loss goes down" would never have caught it; the training even worked with the bug. The checker caught it in minutes.

In the browser, the encoder (8.5 MB) plus index (32 MB) load behind a toggle, retrieval over all 34k passages takes ~50 ms in wasm, and the packed prompt — passages plus your question — is displayed verbatim as the prompt echo, so you see exactly what the model reads. The generation is still 15.7M-grade prose; the point the demo makes is visible steering: same seed with the toggle off is myth soup, with it on the model continues the retrieved scene, characters and all.

War story from the index build: encoding a 4.8 MB corpus hung forever, because my BPE encode ran the merge loop over the whole input as one unit — O(n²), invisible on prompts, fatal on files. The fix (pre-tokenize per word, like the trainer does) is a few lines, and for vocabs trained with the same splitter the token stream is provably identical — the retrieval eval reproduced bit-exact afterwards, which is the kind of regression test this project keeps handing me for free.

What "from scratch" bought me

Concretely, things I now actually understand that I only thought I understood before:

• Why gradient checking needs f64 and central differences, and why "loss goes down" proves nothing.
• Why decode is matvec (not matmul) and therefore bandwidth-bound — and why quantization is a speed feature, not just a memory feature.
• That most "model doesn't work in engine X" bugs are convention bugs — RoPE layouts, tokenizer byte maps, norm epsilons — and that a 1e-4 logit-parity test catches all of them.
• That GPU programming is dominated by when you talk to the GPU, not what the kernels do.
• That a pretrained LM's hidden states are already a usable embedding space, and a contrastive objective is less "new model" than "reshaping geometry you already have."

What's next

FP16 compute and a proper attention kernel to push the Metal speedup on larger models; ulong index math so 7B-class models fit the kernels; and training something bigger than a myth-corpus toy.