vLLM & SGLang Now Run Natively on TPU What PyTorch Devs Need to Know

The Short Version

For years, if you wanted serious inference throughput on TPU, you had two options: JAX-based backends, or a lot of patience. That changed this week. Google announced native PyTorch backends for vLLM (built with Inferact) and SGLang (built with RadixArk), both running on a new layer called TorchTPU.

The pitch is simple: bring your existing PyTorch model definitions, your existing torch.compile workflow, your existing torch.distributed collectives, and get TPU acceleration without rewriting your serving stack in JAX.

If you've been following the broader ecosystem shifts — like the React Foundation moving under the Linux Foundation — you'll recognize the pattern: open governance plus upstream alignment is how infrastructure projects actually survive.

So what's actually new here, and where are the sharp edges?

Under the Hood: How TorchTPU Works

TorchTPU makes the TPU a native PyTorch device. No wrapper tensor types. Tensors live on a tpu device, code runs eagerly or compiles via torch.compile through XLA, and collectives route through torch.distributed.

import torch
import torch_tpu  # registers the 'tpu' device backend

# Standard PyTorch — no wrapper types, no special tensor class
device = torch.device("tpu")
x = torch.randn(4096, 4096, device=device)
y = torch.randn(4096, 4096, device=device)

# Runs eagerly or compiles via XLA
z = torch.compile(lambda a, b: a @ b)(x, y)
print(z.device)  # tpu

The architectural choice that matters most: each device runs its own PyTorch process. Both engines keep their framework's native worker model, scheduler, KV cache manager, and API server — with TPU implementations underneath. Custom Pallas kernels are called directly from PyTorch as custom operators inside the compiled graph. This platform-plugin design means TPU optimizations can ship on their own cadence without forking the upstream framework.

The KV Cache Problem

Agentic workloads break naive serving. Agents resend expanding multi-turn conversations, so recomputing the prefix is pure waste. Both engines keep reusable prefixes in HBM, and with host offload enabled, spill older prefixes to host memory as a second tier.

Two features push KV cache off the computing device entirely:

• Disaggregated serving — prefill and decode run on separate TPU hosts, so KV cache travels from prefill host to decode host.
• Tiered KV offload — overflow goes to host DRAM and secondary storage via Google Cloud Managed Lustre, then comes back on prefix reuse.

Either way, the transfer must overlap computation or it lands directly on latency. Both engines solve this with TPU Sync, an open-source C++ library that reads device buffers behind PyTorch tensors directly — no staging copy on the host — and runs transfers on its own threads outside Python while prefill and decode keep churning. vLLM-TorchTPU exposes it via vLLM's KV connector interface; SGLang-TorchTPU uses it as a disaggregation transfer backend and as the I/O backend for HiCache.

Benchmarks and the Honest Caveats

Google benchmarked publicly through the SemiAnalysis InferenceX project, and the headline number is real: 2.01 cents per million tokens on Qwen 3.5 running on Ironwood TPU (TPU 7x) — the lowest cost per million tokens overall in their submission. They also published SOTA results to AgentX, the long-context and multi-turn track, where Ironwood leads across most of the Pareto frontier.

A Pallas megakernel for the 2.8T-parameter Kimi K3 runs the entire decode step as one kernel, serving 700+ tokens/sec to a single user on Ironwood.

Comparison at a Glance

Where This Gets Complicated

• TorchTPU is new. The wrapper-free tensor story is clean, but the Pallas kernel ecosystem is young. If your model has a weird op, you may be writing the kernel yourself.
• Host offload is not free. Tiered KV offload to DRAM/Lustre only helps if transfer overlaps compute. If your workload is latency-sensitive and short-context, the extra tier can hurt more than it helps.
• JAX isn't going anywhere. vLLM-JAX and SGLang-JAX remain maintained and supported. If your team already runs JAX, migrating to PyTorch-native is a lateral move, not an upgrade.
• Benchmark context matters. The 2.01¢ figure is on a specific model, specific hardware, and a specific workload shape. Your mileage will vary — especially on prefill-heavy vs. decode-heavy distributions.

If you're also tracking the feature-flag side of the Vercel platform, the Vercel Flags GA deep dive is a useful parallel read on how platform-level abstractions mature.

Should You Care?

If you're serving open models at scale and you've been paying the JAX tax — either in tooling friction or in team skill gaps — this is a genuinely meaningful shift. The PyTorch-native path means your existing model definitions, your existing debugging workflow, and your existing torch.compile intuition all carry over.

If you're running smaller workloads or you're already deep in JAX, the calculus is less obvious. The 2.01¢/M-token number is impressive, but it's a frontier benchmark, not a guarantee.

What to do next:

1. Check the upstream GitHub projects for vLLM-TorchTPU and SGLang-TorchTPU — read the issue tracker before you commit.
2. If you're contributing to vLLM, the community TPU capacity pool is worth applying to — free compute for active contributors is a rare deal.
3. Benchmark your own workload shape before you migrate. Prefill-heavy and decode-heavy distributions behave very differently under KV offload.

Source: Google Developers Blog announcement