$ show sgl-project/sglang-jax
SGLang-JAX is a high-performance LLM inference engine optimized for TPU/GPU hardware using JAX. It implements continuous batching, radix tree KV cache management, FlashAttention integration, tensor parallelism, and disaggregated serving (prefill-decode separation). The system provides OpenAI-compatible API endpoints for text generation, chat, embeddings, and tool use, with support for Mixture-of-Experts models, LoRA adaptation, quantization, and speculative decoding. Architecture features modular components: tokenizer manager, scheduler, model executor, memory hierarchy (including HiCache host-device transfer), optimized JAX/Pallas kernels, and LoRA adapter management. Designed for production serving workloads on TPU infrastructure with first-party support for Qwen, Llama, Gemma, DeepSeek, and multimodal models.
The LLM identified 5 files worth highlighting for reuse:
Implements a radix tree (prefix tree) for efficient KV cache management in LLM serving, with optional HiCache extension for host-device memory tiering to support larger context windows. Demonstrates the PagedAttention pattern with a smart host-device hierarchy for memory efficiency.
Abstracts the transfer of KV cache and metadata between prefill and decode nodes in a disaggregated serving setup, supporting multiple backends (JAX sockets, Raiden). Enables separation of compute phases for scalable deployment.
Provides utilities like `device_array` to canonicalize JAX arrays for sharding, preventing unnecessary PJIT recompilation and improving performance in distributed JAX applications. Solves a common problem in JAX-based systems.
Implements batched LoRA adaptation using grouped matrix-vector multiplication (BGMV) to efficiently serve multiple LoRA adapters concurrently. Reduces overhead by combining operations across adapters, demonstrating an optimization for parameter-efficient fine-tuning serving.
Defines an abstract base class for speculative decoding workers, with concrete implementations for algorithms like Eagle and DFlash. Separates the speculative decoding logic from the specific algorithm, enabling easy extension and reuse.