From teacher to tiles — a from-scratch LLM distillation & serving engine: custom Triton/CUDA kernels, FSDP distillation, paged-KV continuous batching, speculative decoding, a Rust gateway, a JAX oracle, and interpretability tooling. This project addresses the challenge of distilling large language models into smaller, efficient student models and serving them effectively. It targets researchers and engineers needing a performant, end-to-end LLM stack, offering custom GPU kernels and optimized serving components for faster inference.

How It Works

Tessera employs a from-scratch approach to distill large teacher LLMs into smaller student models using FSDP/ZeRO-3 sharded training and knowledge distillation losses. The distilled student is then served via an inference engine featuring custom Triton and CUDA kernels for optimized operations like FlashAttention, fused RMSNorm, and int8 matmul. The engine further enhances performance with paged KV caching, continuous batching, and speculative decoding, all managed by a Rust gateway for efficient request handling.

Quick Start & Requirements

• Installation: Clone the repo, create a virtual environment, install dependencies (pip install -e ".[dev]" for CPU, pip install -e ".[dev,gpu]" for Triton/GPU).
• Prerequisites: Python, PyTorch (CPU or CUDA build), Triton (for GPU kernels), Rust (for gateway). NVIDIA GPU recommended for kernel performance.
• Resources: Can run on a laptop (CPU/MPS), but GPU is needed for kernel benchmarks.
• Links: examples/serve.py, examples/train_distill.py, examples/interp_demo.py, benchmarks/, docs/architecture.md.

Highlighted Details

• Custom Triton/CUDA kernels for FlashAttention (GQA, autotuned), fused RMSNorm, SwiGLU, and int8 weight-only matmul.
• From-scratch FSDP/ZeRO-3 implementation with sharded Adam optimizer and atomic, sharded checkpoints.
• Inference engine with paged KV cache (ref-counted allocator for prefix sharing), continuous batching scheduler (admission control, preemption), and speculative decoding.
• Rust gateway (tokio/axum) for HTTP handling and back-pressure, interfacing with the Python engine via PyO3.
• JAX/XLA reference implementation for parity checks.
• Interpretability tooling (activation hooks, logit lens, induction-head detection).

Maintenance & Community

The project appears to be a personal effort by Zengxiao He. No specific community links (Discord, Slack) or explicit maintenance team details are provided in the README.

Licensing & Compatibility

• License: Apache-2.0.
• Compatibility: Permissive license suitable for commercial use and integration into closed-source projects.

Limitations & Caveats

The project is actively under development. Missing features include a fused attention backward kernel, a fused paged-attention decode kernel, FP8 tensor-core GEMM for Hopper GPUs, and JAX pjit/shard_map training support. Some benchmarks are based on the PyTorch reference path on Apple M2 Pro, serving as a performance floor.