TorchTPU
Also known as TorchTPU backend, native PyTorch TPU, torch device tpu
In plain English
TorchTPU makes a TPU look like an ordinary PyTorch device, so vLLM and SGLang can run their existing PyTorch code on Google chips.
Technical definition
TorchTPU is Google’s PyTorch-native TPU backend, built on the PrivateUse1 extension point, that exposes a real torch.Tensor on device tpu and lowers compiled graphs through StableHLO and XLA.
Engineering details
The PyTorch dispatcher routes ATen operations to the TPU backend directly instead of translating them into JAX. Developers can run eagerly for bring-up and debugging, or call torch.compile, where TorchDynamo and AOTAutograd produce an FX graph, TorchTPU lowers it to StableHLO, and XLA emits the TPU executable. Inductor and Triton are not used. Native scope covers tensors, dispatch, eager execution, compile entrypoints, and distributed APIs including DDP, FSDP2, DTensor, and both SPMD and MPMD. The kernel layer stays TPU-specific: TorchTPU can call Pallas and JAX-backed custom kernels, so kernels written for the TorchAX stack migrate rather than restart.
Why it matters
vLLM and SGLang can reuse upstream model code, schedulers, continuous batching, and feature logic instead of rebuilding them across a PyTorch-to-JAX boundary. That lowers the cost of bringing up new models and moves TPU toward day-zero support alongside NVIDIA. TorchTPU replaces TorchAX, which will be deprecated.
How to read it in InferenceX
The TPUv7 results in the InferenceX Official Preview come from the native TorchTPU vLLM stack serving Qwen3.5 397B in FP8. TorchTPU is in private beta as of early September 2026 and is expected to be open sourced around mid-October at the PyTorch Conference, at which point SemiAnalysis moves TPU benchmarking from its fork to the public InferenceX repo.