Tessera: a from-scratch distillation and serving stack in one repository

The gap Tessera is aimed at: distillation that stops at the checkpoint

Most distillation code ends when the student weights are written to disk. The serving path is somebody else's problem, and the kernels that make the student cheap to run are a third project. Tessera's stated goal is to keep that chain in one repository, from a 40M teacher down to a 6M student and then out through an inference engine, a quantization step and a Rust front end. The README is explicit that the point is touching the pieces that matter in practice, and that none of it is meant to be a toy.

The audience is narrow. This is for engineers who want to read a sharded optimizer, a paged KV allocator and a Triton attention kernel in the same tree, and who are willing to run the GPU test markers themselves. It is not a drop-in replacement for a production serving framework, and the README never claims to be one.

What actually moves between the teacher and the student

The distillation side is conventional in its losses and less conventional in its plumbing. The losses live in tessera/distill/losses.py and are temperature-scaled KL divergence, an optional hard cross-entropy term, and hidden-state matching. So the student can be trained against the teacher's softened output distribution, against ground-truth labels, or against intermediate activations, depending on which terms are enabled.

The plumbing is the part worth reading. FSDP/ZeRO-3 is written from scratch in tessera/distill/fsdp.py as flat-parameter sharding with a sharded Adam optimizer. The README states it is checked to be numerically identical to single-process training, both in one process and across two gloo ranks. Checkpoints are atomic and sharded, with resume-from-latest, in tessera/distill/checkpoint.py. If you have ever tried to work out why a sharded run diverges from a single-process baseline, that equivalence test is the design decision that matters here.

The kernel set and its torch fallback

Tessera ships two kernel families. Under tessera/kernels/triton there is a FlashAttention forward kernel with online softmax, causal masking, grouped-query attention and autotuned tile sizes, plus a fused RMSNorm, a fused SwiGLU GEMM, and an int8 weight-only matmul that dequantizes inside the K-loop. Under tessera/kernels/cuda there are raw CUDA C++ versions of RMSNorm and attention, with nvtx ranges and Nsight notes.

The fallback rule is the important operational detail. The Triton and CUDA kernels target NVIDIA GPUs. On anything else the model falls back to a torch reference implementation, and the kernels are checked against that reference whenever a GPU is present. That is why the project runs and is unit-tested on a laptop, on CPU or Apple MPS, and why the published benchmark table is labelled a floor rather than a result. The M2 Pro figures in the README (roughly 100k tok/s for a forward pass on tessera-tiny at B=2, T=128, and about 200 tok/s for engine decode at 6 requests x 48 tokens) measure the reference path, not the fused kernels. The README says as much and points you at pytest -m gpu and benchmarks/ for the numbers that matter.

What is missing is also stated. There is no fused attention backward kernel, no fused paged-attention decode kernel, and no FP8 tensor-core GEMM for Hopper. Those are marked in the code as not done.

Serving: block-paged KV, batch recomposition, speculation

The engine is assembled from three pieces. tessera/serve/paged_kv.py implements a block-paged KV cache with a ref-counted allocator, which is what makes prefix sharing possible: two requests with a common prefix can point at the same blocks, and the refcount decides when a block is actually free. tessera/serve/scheduler.py is a continuous-batching scheduler that recomposes the batch every step, with admission control and preemption under memory pressure. tessera/serve/speculative.py implements speculative decoding with the standard accept/reject sampling.

The tests describe the intended invariants better than the prose does. Incremental decode with the KV cache must match a full forward pass. Self-speculation must reproduce greedy decoding exactly. The engine must drain every request under tight memory and preemption without leaking KV blocks. The README reports 98 to 100 percent acceptance for self-draft greedy speculation on the M2 Pro path. That last test, draining without leaking blocks, is the one to look at if you have been burned by a scheduler that quietly drops requests when the cache fills.

The Rust gateway and the PyO3 boundary

tessera-rs is a tokio and axum gateway that handles HTTP and admission back-pressure, then calls into the Python engine over PyO3. The split is deliberate: the request-facing layer is Rust, the model execution stays Python. The README's example is a single POST to localhost:8080/generate with a JSON body containing a prompt and a params object with max_new_tokens.

The trade-off is visible in the architecture. Every generate call crosses the PyO3 boundary, so the gateway can shed load but cannot avoid the Python runtime underneath it. If your reason for wanting a Rust front end is to escape the Python interpreter entirely, this design does not give you that. It gives you a typed, back-pressured HTTP layer in front of an engine that is still Python. Whether that is worth the extra build step depends on whether admission control at the edge is something you actually need.

Two independent references, and what they are for

The jax_ref directory holds a JAX/XLA reimplementation of the forward pass, used as an independent parity check against PyTorch. The README states the two agree to about 2e-4. This is a different kind of test from the Triton-versus-torch comparisons: it is checking the model definition itself against a second compiler and a second set of numerics, not checking a kernel against the same math written twice.

The same instinct shows up in tessera/interp, which provides activation hooks, a logit lens and induction-head detection. Those are standard mechanistic interpretability tools. Their presence here is not incidental: if you are distilling a teacher into a student, having the logit lens and activation hooks in the same tree means you can inspect what the student learned without bolting on a separate library. The README does not describe a workflow that connects the two, so treat this as adjacent tooling rather than an integrated feature.

Getting it running, and where it will stop you

The quickstart assumes Python 3.10 or newer. After cloning and creating a virtual environment, the README installs torch from the PyTorch CPU wheel index (or a CUDA build), then pip install -e ".[dev]". The CPU test run is pytest -m "not gpu", and the README notes that kernel tests skip without a GPU. tessera info lists presets and parameter counts. Three example scripts cover the main paths: examples/serve.py for continuous batching plus speculative decoding, examples/train_distill.py with a --steps flag, and examples/interp_demo.py.

On Linux with an NVIDIA GPU the extra is pip install -e ".[dev,gpu]", which adds Triton, followed by pytest -m gpu to check the Triton kernels against the torch reference. The Rust side is a separate build: cd tessera-rs, cargo test, cargo run --release. The README badges list Rust 1.75 or newer.

The practical constraint is the split between the two environments. A laptop run exercises the torch reference and the Python tests. It does not exercise the kernels, the fused paths, or any of the numbers in the benchmark table. The README is honest that the M2 Pro figures are a floor, which is the right framing, but it also means a laptop quickstart tells you very little about serving performance.

Alternatives, licensing, and the maintenance question

The obvious comparison is a production inference server with a separate training stack. A framework like vLLM, for instance, concentrates on serving many models well and expects you to bring your own distillation code, or to use a separate library for it. Tessera inverts that: distillation is the centre, and the serving engine exists to run the student it produces. If your problem is serving a model you already have, at scale, with the widest possible kernel coverage, Tessera is the wrong tool, because its Status section lists a fused paged-attention decode kernel as unfinished and its benchmark numbers come from a fallback path.

If your problem is understanding or reproducing the whole chain, from the KD loss through the sharded optimizer to the paged cache, the inversion is the point. You get a tree where the FSDP implementation, the scheduler and the kernels are all small enough to read and all checked against references.

On licensing, the repository badge and the README both say Apache-2.0, but the metadata field for this repository reports NOASSERTION. Those two do not agree, and the discrepancy is worth resolving before you depend on the terms. Nothing in the README suggests a dual-licence or a contributor agreement, but the metadata is what automated tooling will read. This is a factual conflict in the material, not a legal question, and it is the kind of thing to confirm with the author rather than assume.

On maintenance, the material supports only a limited answer. There are no retrieved releases, so there is no version history to reason about upgrade cost from. The last push is dated 2026-06-05. The README points to docs/architecture.md, docs/kernels.md, docs/serving.md and docs/distillation.md as the deeper references, which is where the design rationale you would need for a fork lives. The unfinished items listed in Status (fused attention backward, fused paged-attention decode, Hopper FP8 GEMM, pjit and shard_map on the JAX side) are the natural upgrade surface: if you depend on any of them, you are depending on code that does not exist yet.

Editorial conclusion

Adopt Tessera if you want a readable, end-to-end distillation and serving pipeline where the FSDP sharding, the Triton kernels and the scheduler are all in one tree and all checkable against a torch or JAX reference. Do not adopt it if you need a fused paged-attention decode kernel or a Hopper FP8 GEMM today, since the Status section lists both as unfinished. Before committing, run pytest -m gpu on your target NVIDIA hardware and read docs/architecture.md, because the published throughput figures come from an Apple M2 Pro on the torch fallback path and say nothing about what the fused kernels do on a real GPU.

Official sources

1. Issues
2. README
3. zengxiao-he/tessera on GitHub