AI inference glossary
Software

SGLang-JAX

Also known as SGLang-JAX, sgl-jax, SGL-torchtpu

In plain English

SGLang-JAX is the JAX-native version of the SGLang serving engine for TPUs, separate from the PyTorch SGLang runtime.

Technical definition

SGLang-JAX is a JAX-native serving engine that pairs SGLang-style scheduling and prefix caching with JAX model implementations and TPU-specific kernels, rather than translating the PyTorch SGLang runtime onto TPU.

Engineering details

Where vLLM reached TPU through TorchAX translation, SGLang took a different route and reimplemented its engine in JAX. That gave clean access to mature TPU primitives at the cost of maintaining a second codebase that diverges from upstream SGLang. Google and RadixArk have announced SGL-torchtpu as a PyTorch-native alternative built on TorchTPU, which would let SGLang on TPU share upstream code paths again.

Why it matters

Two engines with two different TPU strategies is a sign the ecosystem had not settled. TorchTPU is meant to converge them: both vLLM and SGLang get a native PyTorch device, so kernel work in Pallas can be shared and day-0 model support can land on TPU alongside NVIDIA and AMD.

How to read it in InferenceX

The TPU InferenceX preview results use TorchTPU vLLM. SGLang on TPU results will follow once TorchTPU SGLang is open sourced, at which point SemiAnalysis moves TPU benchmarking into the public InferenceX repo.