SGLang-TorchTPU

PyTorch-native LLM serving on Google TPUs.
The SGLang serving engine. The PyTorch ecosystem. Built for TPU.

SGLang-TorchTPU brings SGLang to Google TPUs through torch_tpu and Pallas kernels. Serve large language and vision-language models with SGLang's OpenAI-compatible API and scheduling engine, while developing in PyTorch with eager execution and TPU compilation.

Documentation

PyTorch-native TPU serving. Develop with familiar PyTorch workflows and deploy on TPU7x (Ironwood), with torch.compile and direct compilation. Tensor, expert, data, pipeline, and context parallelism cover a range of serving configurations.

Efficient inference at scale. Combine prefix caching, HiCache KV offload to host memory, disaggregated prefill/decode with TP resharding, and encoder disaggregation for vision-language models. Low-precision inference includes FP8, MXFP4, INT4 W4A16, and FP8 KV cache.

Speculative decoding. Explore EAGLE3, DFlash, DFlash2, DSPark, and MTP for supported model configurations. The speculative decoding developer guide below covers worker selection and model-specific requirements.

Pull the TPU image, which has the engine and TPU plugin installed:

docker pull lmsysorg/sglang:v0.1.0-tpu

On a TPU7x host, start a container with access to the TPU chips:

docker run -it --rm --privileged --net=host --shm-size=64g \
  -e SGLANG_PLATFORM=tpu_oot \
  lmsysorg/sglang:v0.1.0-tpu

The container starts in the repository root, /sgl-workspace/sglang.

On a TPU7x host (Linux x86_64, Python 3.11 to 3.14), install the engine and TPU plugin wheels from the v0.1.0 release. Install both together: sglang here is this repository's build, not the PyPI package.

pip install --extra-index-url https://download.pytorch.org/whl/cpu \
  https://github.com/sgl-project/sglang-torchtpu/releases/download/tpu-v0.1.0/sglang-0.5.22-py3-none-any.whl \
  https://github.com/sgl-project/sglang-torchtpu/releases/download/tpu-v0.1.0/sglang_tpu-0.1.0-py3-none-any.whl

The serving recipes below live in the repository, so clone it to run them.

On a TPU7x host, install the engine with the TPU dependency manifest, then the TPU plugin:

git clone https://github.com/sgl-project/sglang-torchtpu.git
cd sglang-torchtpu
cp python/pyproject_tpu.toml python/pyproject.toml
pip install --extra-index-url https://download.pytorch.org/whl/cpu -e python/
pip install --no-deps -e .

From the repository root, launch GPT-OSS-20B with the single-host serving recipe, configured for 4 chips and 8 TensorCores:

python -m benchmark.tpu.launch_server --config recipe/configs/gpt_oss_20b/server.py

Once the server is ready, send an OpenAI-compatible request from another terminal:

curl http://localhost:30000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model": "gpt-oss-20b", "messages": [{"role": "user", "content": "Hello"}]}'

Development and testing focus on TPU7x (Ironwood), from a single host to multi-host serving and disaggregated prefill/decode.

Selected model integrations include the following. Recipe availability and test coverage vary by model; experimental integrations are included.

See the full support guide for support tiers, model-specific requirements, and experimental capabilities.

Explore model recipes, TPU benchmarks, and the CI support matrix for configurations and test coverage.

A vision-language model can run its vision encoder in a separate encode server that sends image embeddings to a language-only server (#665). EPD is supported for the Qwen family and verified on Qwen3-VL-8B and Qwen3-VL-235B-A22B; other model families should work but are not verified.

• Latency. At the same request rate, TPOT drops 4-30% on 8B and 5-24% on 235B-A22B, and TTFT stays low as load rises: at 5 req/s on 235B-A22B, 392 ms vs 2091 ms with the encoder on the same host.
• HBM. The language server no longer compiles or runs the encoder, so after warmup it uses about 2 GB less HBM per TensorCore (8B and 235B-A22B alike).
• Independent topology. The encoder and the language model pick their own TP size and hosts: for example, a TP8/EP8 235B-A22B language server with a TP8 encoder on a second host. Scale the encoder by adding encode servers or TP; encoder data parallelism is not supported on TPU.

# Encode server (requires --enable-torch-compile on TPU)
sglang serve --model-path $MODEL --device tpu --tp-size 8 --enable-torch-compile --encoder-only
# Language server; without a registered encoder it returns 503 for image requests
sglang serve --model-path $MODEL --device tpu --tp-size 8 --enable-torch-compile --language-only \
  --encoder-urls http://$ENCODER_HOST:$ENCODER_PORT

SGLang-TorchTPU is the PyTorch-native TPU backend in the SGLang ecosystem, alongside the JAX-based SGL-JAX. It brings TPU execution to the SGLang serving interface, so applications can use familiar OpenAI-compatible clients.

TPU kernels, memory management, and backend integrations live in sglang_tpu/.

Contributions are welcome in model support, kernel optimizations, serving features, benchmarks, and documentation.

1. Read the contributor guide for development setup and the pull request workflow.
2. Explore the TPU backend and speculative decoding developer guide.
3. Add coverage for your changes and include accuracy or performance results where relevant.

CPU-only unit tests run without a TPU:

python test/run_suite.py --suite base-a-test-cpu

Kernel and end-to-end tests require TPU hardware. A maintainer applies the run-ci label and comments /tag-and-rerun-ci to run them on your pull request.

• Bugs and feature requests: Open an issue, including your TPU generation, topology, and torch_tpu version for bug reports.
• Security: See our security policy for reporting vulnerabilities.

For engine issues that also reproduce on GPUs, use the upstream SGLang issue tracker.

SGLang-TorchTPU is released under the Apache 2.0 license.

SGLang-TorchTPU is a joint effort of Google, Meta, RadixArk, and the SGLang community. Thanks to all contributors.

See NOTICE for attribution details.