Google brings PyTorch-native vLLM and SGLang inference engines to TPU

Why Google built PyTorch backends

vLLM and SGLang already run on TPU through a JAX-based backend. Google says developers kept asking for a way to customize and debug their serving stack in PyTorch. The company responded with TorchTPU, which runs PyTorch natively on TPU, and then built new vLLM and SGLang engines on top of it.

According to Google, the goal is for the PyTorch community to bring models to TPU and make them fast with tools developers already know. The engines ship with production-ready serving features and custom Pallas kernels.

TorchTPU makes the TPU a native PyTorch device. Tensors are standard PyTorch tensors on a tpu device with no wrapper types. Code runs eagerly or compiles through torch.compile using XLA, and collectives go through torch.distributed. Both engines keep each framework's worker model, scheduler, KV cache manager and API server, and add TPU implementations underneath. Most models use the standard upstream PyTorch model definitions, with TPU-specific work limited to kernels and layer implementations.

Supported features and KV cache transfer

Google lists support across several workload types:

• Agentic and multi-turn: automatic prefix caching, KV cache offloading to host memory and Managed Lustre, tool calling and structured output.
• Low latency: speculative decoding with MTP, EAGLE-3, DFlash, DSpark and n-gram drafters.
• Long context: context parallelism for prefill and decode, FP8 KV cache, host KV offloading.
• Disaggregated serving: prefill and decode on separate hosts.
• Multimodal: vision-language models with image and video input, plus embedding models.
• Very large models: several parallelism modes across hosts, and FP8, FP4 (MXFP4, NVFP4), INT8 and INT4 checkpoints.

Agents repeatedly resend growing conversations, so finding a cached prefix instead of recomputing it matters. Both engines keep reusable prefixes in HBM and, with host offload enabled, move older prefixes to host memory as a second tier.

Disaggregated serving and tiered offload both move the KV cache away from the device that computed it. Both engines use TPU Sync for this, an open-source C++ library. According to Google, it reads device buffers behind PyTorch tensors directly, so no extra staging copy is needed on the host, and it runs transfers on its own threads outside Python while prefill and decode continue.

Benchmarks and production use

Google benchmarks publicly through the SemiAnalysis InferenceX project. Building on an earlier submission of Qwen 3.5 on Ironwood TPU (TPU 7x), the company says further optimization brought the lowest cost per million tokens overall, at 2.01 cents in high-concurrency configurations. It also published new results to AgentX, the project's track for agentic long-context and multi-turn workloads, and says Ironwood TPU leads across most of the Pareto frontier.

With Inferact, Google integrated a Pallas megakernel for the 2.8T-parameter Kimi K3 model that runs the entire decode step as one kernel. Google says it serves more than 700 tokens per second to a single user on Ironwood TPU. Inferact also led a first preliminary AgentX submission for Kimi K3 on vLLM-TorchTPU.

SGLang-TorchTPU, co-developed with RadixArk, is running in production at Meta. Zongwei Zhou, PyTorch-TPU lead at Meta, said it gives Meta "a remarkably mature, PyTorch-native serving stack that integrates effortlessly into our existing AI infrastructure." Simon Mo, CEO and co-founder of Inferact, said developers can "immediately run production workloads on TPUs" with vLLM.

What it means for developers

For teams that work in PyTorch, the launch offers a way to serve open models on TPU without moving to JAX. Google says most popular open model families are supported, including Qwen, DeepSeek, GLM, Kimi, GPT-OSS and Gemma4, and that it is working toward Day 0 support for new open models.

vLLM is also launching a community contribution TPU capacity pool that provides free TPU compute to developers who actively contribute to both engines. Google adds that more optimizations are under way and invites developers to open a GitHub issue to request collaboration.