One repo, every layer of the LLM stack, written by hand

What it does

Tessera is a deliberately small but complete LLM pipeline: distill a teacher into a student, then serve that student. The teacher is 40M parameters, the student 6M, and every layer in between — Triton/CUDA kernels, from-scratch FSDP, a paged-KV inference engine with continuous batching, speculative decoding, quantization, and a Rust HTTP gateway — is hand-built rather than imported.

The interesting bit

The whole thing runs and is unit-tested on a laptop, CPU or Apple MPS included. The kernels fall back to a torch reference off-NVIDIA, and the repo leans on cross-checks everywhere: a JAX/XLA reimplementation serves as an independent parity oracle for the PyTorch forward pass, and sharded Adam is verified step-for-step against torch.optim.Adam.

Key highlights

• FlashAttention forward in Triton (online softmax, GQA, autotuned tiles), plus fused RMSNorm, fused SwiGLU GEMM, and an int8 matmul that dequantizes in the K-loop
• FSDP/ZeRO-3 written from scratch — flat-parameter sharding with a sharded optimizer, checked numerically identical to single-process training
• Serving stack: block-paged KV cache with ref-counted prefix sharing, continuous batching with admission control and preemption, and speculative decoding
• Post-training quantization in three flavors: int8 weight-only, AWQ, and FP8 (E4M3)
• Rust gateway (tokio + axum) handling HTTP and back-pressure, calling the Python engine over PyO3

Caveats

• The published benchmarks are from an M2 Pro on the torch reference path — the author explicitly calls them a floor, not what the fused kernels do on real GPUs
• Known gaps, marked in the code: no fused attention backward, no fused paged-attention decode kernel, no Hopper FP8 GEMM, no sharded JAX training

Verdict

A study guide that actually runs — ideal if you want to understand how these pieces fit together by reading (and running) working code. If you need a production inference server, this isn’t it, and it doesn’t pretend to be.