SGL-JAX: High-Performance LLM Inference on JAX/TPU

SGL-JAX is a high-performance, JAX-based inference engine for Large Language Models (LLMs), specifically optimized for Google TPUs. It is engineered from the ground up to deliver exceptional throughput and low latency for the most demanding LLM serving workloads.

The engine incorporates state-of-the-art techniques to maximize hardware utilization and serving efficiency, making it ideal for deploying large-scale models in production on TPUs.

Key Features

• High-Throughput Continuous Batching: Implements a sophisticated scheduler that dynamically batches incoming requests, maximizing TPU utilization and overall throughput.
• Optimized KV Cache with Radix Tree: Utilizes a Radix Tree for KV cache management (conceptually similar to PagedAttention), enabling memory-efficient prefix sharing between requests and significantly reducing computation for prompts with common prefixes.
• FlashAttention Integration: Leverages a high-performance FlashAttention kernel for faster and more memory-efficient attention calculations, crucial for long sequences.
• Tensor Parallelism: Natively supports tensor parallelism to distribute large models across multiple TPU devices, enabling inference for models that exceed the memory of a single accelerator.
• OpenAI-Compatible API: Provides a drop-in replacement for the OpenAI API, allowing for seamless integration with a wide range of existing clients, SDKs, and tools (e.g., LangChain, LlamaIndex).
• Native Qwen Support: Includes first-class, optimized support for the Qwen model family, including recent Mixture-of-Experts (MoE) variants.

Architecture Overview

SGL-JAX operates on a distributed architecture designed for scalability and performance:

1. HTTP Server: The entry point for all requests, compatible with the OpenAI API standard.
2. TokenizerManager: Runs in the main process, handles text tokenization
3. Scheduler: The core of the engine. It receives requests, manages prompts, and schedules token generation in batches. It intelligently groups requests to form optimal batches for the model executor.
4. TP Worker (Tensor Parallel Worker): A set of distributed workers that host the model weights, distributed via tensor parallelism. They execute the forward pass for the model.
5. Model Runner(Included in TP Worker): Manages the actual JAX-based model execution, including the forward pass, attention computation, and KV cache operations.
6. DetokenizerManager: Runs in a subprocess, handles output token decoding

More details in architecture.

Getting Started

• Install SGL-JAX
• Cookbook (model recipes)
• Quick Start (Qwen3 on TPU v6e-4)
• Benchmark and Profiling
• Contribution Guide

Documentation

For more features and usage details, please read the Sphinx-style documents in the docs directory and the Mintlify-style model recipes in docs/cookbook.

Supported Models

SGL-JAX is designed for easy extension to new model architectures. It currently provides first-class support for:

• Qwen
• Qwen 2 / Qwen 2 MoE
• Qwen 3 / Qwen 3 MoE
• Llama
• Gemma 2
• DeepSeek V2 / V3
• GLM-4 MoE
• Grok-2
• Bailing MoE / Bailing MoE V2
• MiMo-7B
• MiMo-V2-Flash
• MiMo-V2.5-Pro

SGL-JAX also supports multimodal models with the same usage interface as LLMs. The architecture has been adapted to support flexible multimodal model architectures.

• Wan 2.1 T2V: Text-to-Video generation model.
• Wan 2.2 T2V: Text-to-Video generation model. Uses different DiT models at different noise stages for denoising.
• Qwen2.5-VL: Vision-language model series based on Qwen2.5.

For multimodal model usage, see the Usage Guide and Architecture Design.

Performance and Benchmarking

For detailed performance evaluation and to run the benchmarks yourself, please see the scripts located in the benchmark/ and python/sgl_jax/ directories (e.g., bench_serving.py).

Testing

The project includes a comprehensive test suite to ensure correctness and stability. To run the full suite of tests:

python test/srt/run_suite.py

Contributing

Contributions are welcome! If you would like to contribute, please feel free to open an issue to discuss your ideas or submit a pull request.

Before contributing, please read our Contribution Guide for setup instructions, coding standards, and contribution workflow.

You can also join our community on Slack to discuss ideas, get help, or collaborate with other contributors: 👉 Join the SGLang Slack workspace (https://slack.sglang.io/), then participate in discussions in the SGL-JAX Slack Channel.