sgl-project/sglang-jax

JAX backend for SGL

What it solves

SGL-JAX addresses the challenge of deploying Large Language Models (LLMs) on Google TPUs with high throughput and low latency. It provides a specialized inference engine that maximizes hardware utilization for demanding production workloads.

How it works

The engine uses a distributed architecture consisting of an OpenAI-compatible HTTP server, a scheduler for continuous batching, and Tensor Parallel Workers that execute the model using JAX. It optimizes performance through several key mechanisms:

  • Radix Tree KV Cache: Manages memory efficiently and allows requests with common prefixes to share cached data, reducing redundant computation.
  • Continuous Batching: Dynamically groups incoming requests to keep the TPU fully utilized.
  • FlashAttention: Integrates high-performance kernels to speed up attention calculations for long sequences.
  • Tensor Parallelism: Splits large models across multiple TPU devices to handle models that exceed a single chip's memory.

Who it’s for

It is designed for developers and organizations deploying large-scale LLMs or multimodal models (such as text-to-video or vision-language models) specifically on Google TPU hardware.

Highlights

  • OpenAI-Compatible API: Works as a drop-in replacement for existing OpenAI-based tools and SDKs.
  • Broad Model Support: Optimized for Qwen (including MoE variants), Llama, Gemma 2, DeepSeek, and others.
  • Multimodal Capabilities: Supports text-to-video (Wan 2.1/2.2) and vision-language (Qwen2.5-VL) models.
  • JAX-Based: Built from the ground up for high-performance execution on TPUs.

Related

  • Project
    vllm-project/tpu-inference

    A hardware plugin for vLLM that provides a unified backend for high-performance LLM inference on Google TPUs, supporting both PyTorch and JAX.

  • Project
    AI-Hypercomputer/maxtext

    A high-performance, scalable LLM library and reference implementation written in JAX for training and post-training models on Google Cloud TPUs and GPUs.

  • Project
    ekzhang/jax-js

    A high-performance machine learning framework for the browser that brings JAX-style numerical computing, automatic differentiation, and GPU acceleration via WebGPU and Wasm.

  • Project
    erfanzar/EasyDeL

    A JAX and Flax NNX-based framework for the high-performance training, fine-tuning, and serving of LLMs and multimodal models on TPU/GPU clusters.

  • Project
    sgl-project/sglang

    A high-performance serving framework for LLMs and multimodal models that provides low-latency, high-throughput inference across single GPUs and large distributed clusters.