tpu-mini-sglang：基于JAX和TPU的教育级LLM推理库

tpu-mini-sglang is an educational LLM inference library inspired by mini-sglang, built with JAX for Google TPU. It retains the core architecture of SGLang while stripping production-level complexity, making it ideal for learning modern LLM service framework mechanisms. The project is open-sourced under the Apache 2.0 license, emphasizing knowledge sharing and educational accessibility.

SGLang is renowned for efficient structured generation and parallel scheduling, but its full codebase is too large and complex for learners. tpu-mini-sglang was created to fill this gap—it is positioned explicitly for educational use, with smaller code volume and clearer structure, allowing learners to focus on core LLM inference concepts without being overwhelmed by engineering details.

JAX is chosen as the core computing framework (fit for TPU with functional programming and auto-differentiation). The library maintains a complete modular design consistent with production frameworks:

• entrypoints/: Handles API requests
• kernels/: Core computing operations (e.g., attention mechanisms)
• layers/: Neural network layer implementations
• managers/: Resource coordination (memory/computation)
• mem_cache/: KV cache optimization
• model_executor/: Model execution engine
• models/: Supported model definitions
• sampling/: Sampling strategy implementations

• ModelConfig Class: Parses critical parameters from HuggingFace configs (num_heads, num_kv_heads, hidden_size, head_dim, intermediate_size, dtype, context_len, EOS/BOS token IDs)
• Flexible Dtype Support: Automatically selects optimal data types (e.g., bfloat16) balancing precision and performance
• Sharding Support: Basic model/data parallelism via sharding.py (key for large-scale LLM deployment)

Core Dependencies: FastAPI (≥0.110), Flax, JAX, Transformers (≥4.57.1), Tokenizers (≥0.21.1), SafeTensors Optional Backends: CPU (jax[cpu]), GPU (jax[cuda12]), TPU (jax[tpu]—primary target) Development Tools: Ruff (formatting/linting), MyPy (static typing), Codespell (spelling check), pre-commit hooks (automated checks)

Target Learners: Deep learning framework developers, TPU/JAX users, SGLang enthusiasts (finding full source too complex), education researchers Recommended Path: 1. Understand model config via model_config.py; 2. Explore attention mechanisms in kernels/;3. Track request flow from launch_server.py;4. Experiment on Google Colab TPU environment

Project	Scale	Target Platform	Main Use Case
SGLang	Large	Multi-platform	Production deployment
mini-sglang	Medium	GPU	Education/research
tpu-mini-sglang	Small	TPU	Education/TPU-specific
llm.c	Extra-small	CPU	Minimalist education

Unique value: TPU-optimized educational implementation, filling the gap in JAX/TPU ecosystem for educational LLM inference frameworks.

tpu-mini-sglang demonstrates 'small but beautiful' educational value—with ~760 lines of core code, it covers key LLM service components (model config, kernel computation, memory management, sampling, service interface). It is an ideal starting point for learners and a lightweight foundation for TPU-based LLM deployment. As JAX ecosystem matures and TPU accessibility improves, such projects will play an increasingly important role in AI education.