Torch-TPU: Running SGLang and vLLM Natively on TPUs

A Native PyTorch Path to TPU Inference — Without JAX

At the PyTorch Conference in San Jose, Qi from Google and Colin and Angela from Meta jointly introduced a new approach to large-model inference: Torch-TPU, a native PyTorch TPU backend that enables mainstream inference engines SGLang and vLLM to run natively on Google TPUs.

This work takes a fundamentally different approach from previous TPU adaptation efforts. Historically, running large models on TPUs meant one of two options: building an entirely new serving engine, or porting models into the JAX ecosystem. Both paths are costly, and both create a gap between TPU and the mainstream GPU toolchain. Torch-TPU's goal is to run existing SGLang and vLLM infrastructure natively on TPUs — not to start from scratch.

Why Reuse Existing Engines Instead of Rewriting?

SGLang and vLLM are the two dominant open-source inference engines in the LLM space today. Around them, a mature ecosystem has formed — including schedulers, batching systems, server APIs, and compilation workflows. Engineering teams are already well-versed in this toolchain, with stable production deployments running on GPUs.

The team emphasized that they preserved the components GPU users are already familiar with: schedulers, batching systems, server APIs, and compilation workflows. From the user's perspective, the cognitive overhead of migrating to TPUs is minimized — you're still using the same SGLang or vLLM, just with TPU under the hood.

This "reuse over rewrite" strategy reflects a pragmatic view of engineering reality. Rewriting a serving engine from scratch is not only time-consuming — it also means constantly playing catch-up with the upstream community's ongoing work on scheduling, continuous batching, prefix caching, and more. Rather than chasing, it's better to stand on the upstream community's shoulders.

Continuous batching is one of the core scheduling techniques in modern inference engines, introduced in the Orca paper (2022) and widely adopted by engines like vLLM. The key idea: rather than waiting for an entire batch of requests to complete before starting the next, the engine dynamically replaces finished requests with new ones at each iteration step, dramatically improving GPU/TPU utilization. Prefix caching is another critical optimization — for multiple requests that share the same system prompt or context prefix, the engine can reuse already-computed KV Cache to avoid redundant computation, significantly reducing time-to-first-token (TTFT). Both techniques are deeply embedded in scheduler logic and are extremely costly to reimplement. By reusing the scheduling layer from SGLang/vLLM, Torch-TPU gets these finely tuned mechanisms working on TPUs for free.

The Key Dividend of Staying Upstream: Zero-Cost Model and Feature Portability

One of the most important aspects of the Torch-TPU approach is that the serving engine stays upstream. This means new model support and features added by the SGLang and vLLM communities can run on TPUs directly — without requiring a separate TPU-specific implementation.

This is enormously valuable in practice. The large-model ecosystem evolves at a rapid pace, with new model architectures, quantization schemes, and sampling strategies emerging almost weekly. If TPUs required dedicated adaptation for every new feature, they'd always be half a step behind. But when the engine stays upstream, TPU users can theoretically enjoy the latest community developments almost in sync with GPU users.

From a hardware ecosystem competition standpoint, this is also a critical step toward lowering the barrier to TPU adoption and attracting GPU users to migrate. CUDA's moat has long been built largely on the richness of its software toolchain. Getting mainstream inference engines to natively support TPUs is precisely an attempt to close that gap at the software layer.

Cross-Vendor Collaboration Between Google and Meta

The composition of contributors here is noteworthy: Google provides the TPU hardware and backend, while Meta is a key driver of PyTorch and the broader inference engine ecosystem. Their collaboration on this project reflects PyTorch's value as a cross-vendor neutral framework — one capable of supporting NVIDIA GPUs while also extending to heterogeneous hardware like TPUs through native backends.

It's worth noting that the information released so far is only a conference preview. Specific performance benchmarks, the range of supported models, and the technical implementation details of the Torch-TPU backend were not disclosed in the preview — those are reserved for the full talk at the PyTorch Conference.

PyTorch's cross-hardware extensibility stems from its backend abstraction mechanism. Through torch.distributed and a device abstraction layer, PyTorch allows third parties to register custom operators and compilation backends for new hardware. XLA (Accelerated Linear Algebra) was Google's primary path to connecting TPUs with PyTorch, but XLA's compilation model has significant impedance with PyTorch's dynamic graph style — making it difficult to debug and incurring high compilation overhead. The "native PyTorch TPU backend" referenced by Torch-TPU implies a deeper level of integration — driving TPUs directly through PyTorch's Dispatcher and TorchInductor/torch.compile stack, rather than going through the JAX/XLA programming model. This means PyTorch users can deploy models on TPUs without learning a new programming paradigm.

Conclusion

Torch-TPU represents a clear engineering philosophy: don't fight against a mature upstream ecosystem — instead, implement hardware adaptation as a low-level backend so that upper-layer applications can migrate transparently. For teams looking to deploy large-model inference on TPUs without abandoning the SGLang or vLLM toolchain, this is a direction worth following closely. Full technical details are set to be revealed at the conference.