Skip to content

Google’s TorchTPU Runs PyTorch Natively on TPUs at 2.01 Cents per Million Tokens, but That Is a Hardware Cost Benchmark, Not Your API Bill

Google and Meta open-sourced TorchTPU, a native PyTorch backend for TPUs. The benchmark claims are strong, but the fine print matters more than the headline.

TorchTPU native PyTorch backend for Google TPUs

Google open-sourced TorchTPU on October 20, 2026, turning TPUs into a native PyTorch device instead of a side quest through JAX or lazy tensors. The project is a collaboration with Meta’s PyTorch team, and the pitch is simple: change model.to("cuda") to model.to("tpu") and keep your stack.

That pitch matters because TPUs have looked cheaper and faster on paper for years, while PyTorch developers kept hitting the same wall. PyTorch/XLA required source changes, compiled slowly and could not run in eager mode. TorchTPU is Google’s attempt to remove that wall without asking developers to learn JAX.

The headline number is 2.01 cents per million tokens for Qwen 3.5 on an Ironwood TPU, measured in the SemiAnalysis InferenceX benchmark. It is a Google-submitted benchmark result, not a public API price. The gap between those two things is the story.

What Google actually shipped

Google published four repositories at PyTorch Conference 2026: TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU and TorchTitan. The first is the device backend. The other three are serving and training stacks built on top of it.

The core runtime does three things. It runs standard PyTorch operations eagerly, compiles them through torch.compile(backend="tpu"), and sends distributed collectives through torch.distributed using a TPU process group called tpu_dist. Installation is one command on a TPU VM: pip install torch_tpu.

The coverage claim is broad. Google says TorchTPU supports roughly 450 Aten operations, including forward and backward passes for major model families, and was validated out of the box on the 5,000 most downloaded Hugging Face models. Transformers v5 loads a model on TPU with device_map="tpu". Diffusers supports eager, compiled and tensor-parallel paths.

The serving story matters most for inference. vLLM-TorchTPU is an out-of-tree plugin that keeps vLLM’s upstream scheduler, worker model, KV cache manager and OpenAI-compatible API server while adding TPU kernels. SGLang uses a similar design with a disaggregation backend. Both use Pallas kernels for attention, mixture-of-experts and quantized matmul.

Training is not an afterthought. Google says TorchTitan on TPU reached up to 52% Model Flops Utilization on Qwen 3 8B and scaled to 8,192 Ironwood TPUs. That is a vendor claim, but it covers multi-host training, not just one optimized decode path.

The old PyTorch/XLA problem

PyTorch/XLA was the first serious attempt to run PyTorch on TPUs. It used lazy tensors, which recorded operations into a graph and executed only when a result was needed. That made compilation possible, but debugging strange: a small shape change could trigger a long recompile, and developers often had to mark compilable regions by hand.

TorchTPU changes the contract. Operations run eagerly by default, one process per device, with normal tensor views and in-place mutations. For speed, TorchDynamo captures the graph and lowers it through OpenXLA and StableHLO. The runtime has no JAX dependency, though custom kernels can be written in JAX/Pallas and called as PyTorch operators.

The catch is that a device swap is not a performance port. Google says a model should run after changing the device, but peak performance still depends on shapes, sharding, batch size and kernels. A GPU-tuned model with custom CUDA kernels will run, but it will not automatically match a TPU-native implementation.

The 2.01 cent number, decoded

The Google Cloud announcement says the vLLM-TorchTPU stack reached the lowest cost per million tokens overall on InferenceX 8K/1K for Qwen 3.5 on Ironwood, at 2.01 cents. That is down from 3.6 cents in Google’s September submission. On AgentX, which tests multi-turn, long-context and sub-agent traffic, Google reports 0.83 cents per million tokens with four TPU v7 chips.

Google also reports that a Pallas megakernel for Kimi K3, a 2.8 trillion parameter model, runs the entire decode step as one kernel and serves more than 700 tokens per second to a single user on Ironwood. The vLLM team describes the same work as a custom kernel effort for a model that would otherwise be difficult to serve efficiently.

Those numbers are not a price list. Cost per million tokens in InferenceX is derived from the cost of the tested hardware and the measured throughput. It is a way to compare serving efficiency across accelerators, not the rate Google Cloud will charge your project. The result also depends on the model, the input and output lengths, the batch size, the quantization and the parallelism strategy.

The single-user Kimi number deserves the same care. More than 700 tokens per second to one user is an interactivity result. It does not tell you what a multi-tenant endpoint will deliver when hundreds of users share the same slice, or what happens when prompts are short, outputs are long and prefix caching does not help.

How TorchTPU compares with the old path

Capability TorchTPU PyTorch/XLA What it changes
Run existing code after a device swap ✓ — TorchTPU targets model.to("tpu"). The XLA path often needed code changes to mark compilable regions.
Eager execution by default ✓ — Developers can inspect tensors and debug step by step before compiling.
torch.compile through TorchDynamo ✓ — The compiled path uses OpenXLA and StableHLO instead of lazy tensors.
Upstream vLLM and SGLang serving ✓ — The new backends keep upstream schedulers, workers and API servers, with TPU work underneath.

The comparison is not a clean win in every row. PyTorch/XLA had years of production use and a large body of workarounds. TorchTPU is younger, and its ecosystem is still filling gaps. The difference is architectural: TorchTPU puts TPU support inside the PyTorch programming model instead of asking PyTorch to behave like a JAX program.

Where the stack is strong, and where it is not

The strongest part is reuse. A team already serving with vLLM can install a plugin and keep the same CLI, OpenAI-compatible server and much of its operational tooling. A team using Transformers can move a model with one argument. That is a lower migration cost than rewriting a stack in JAX.

The weakest part is maturity. The Google Cloud documentation recommends TorchTPU for new workloads, but vLLM-TorchTPU recommends TPU v7x Ironwood and treats v6e, v5e and v5p as experimental. That matters if your quota is only available on an older generation.

Hugging Face offers a reality check. Its launch post says 96.5% of the tests that run on TPU pass, with 34 issues reported and 19 already fixed. Google also says the engines run in production with preview customers. Preview is the operative word.

What this means for teams in India

The pitch lands hardest where GPU capacity is expensive or hard to reserve. Indian inference startups often rent GPUs by the hour in dollars, and a lower cost per token is meaningful when utilization is high. But TPU capacity is zone and quota dependent, so check data residency, latency and egress before moving production traffic.

For most teams, the practical test is not the leaderboard. It is a two-week pilot with one model, one real traffic trace and one cost sheet. Compare Ironwood against the GPU you can actually get, including idle time, migration work and tuning. The 2.01 cent number is a useful negotiating benchmark, not a budget.

The llm-d 1.0 release made a similar promise around agentic inference, and the DeepInfer release showed how much performance depends on what a vendor chooses to open. TorchTPU is more open than either, but the benchmark still belongs to the team that built the stack.

Confirmed vs not confirmed

Claim Status Evidence
TorchTPU is open source and available on GitHub Confirmed Google Cloud blog, PyTorch blog and the public repository all point to the October 20 launch.
It was co-developed with Meta’s PyTorch team Confirmed Both Google and PyTorch describe the collaboration.
It supports eager mode, torch.compile and native distributed collectives Confirmed Google Cloud documentation and the PyTorch technical post describe all three.
It reached 2.01 cents per million tokens on InferenceX Confirmed as a benchmark result The number is published, but it is a cost-efficiency metric, not a Google Cloud list price.
Kimi K3 serves more than 700 tokens per second to one user Confirmed as a vendor claim Google and vLLM describe the Pallas megakernel result. No independent replication is cited.
It is production-ready for every PyTorch workload Not confirmed TPU v7x is recommended, older TPUs are experimental and Hugging Face reports a 96.5% test pass rate.
It is cheaper than GPUs for every workload Not confirmed Cost depends on model, utilization, region, quota and migration work. The benchmark covers specific configurations.

FAQ

Is TorchTPU free to use?

The software repositories are public and open source. Running them still requires Cloud TPU capacity, which Google Cloud bills by the chip or slice.

Do I have to rewrite my model for TPU?

Most standard PyTorch models should run after changing the device, especially through Transformers or vLLM. Peak performance may still require sharding, compiled mode, quantization and TPU-specific kernels.

Does TorchTPU replace PyTorch/XLA?

Google’s documentation recommends TorchTPU for new workloads and says teams should migrate existing PyTorch/XLA workloads. PyTorch/XLA is still part of the story for older code, but it is no longer the recommended path.

Which TPU should I use?

For vLLM-TorchTPU, Ironwood TPU v7x is the recommended and continuously tested generation. TPU v6e, v5e and v5p are labeled experimental. Check your project quota and region before planning a deployment.

Is 2.01 cents what I will pay per million tokens?

No. It is a benchmark-derived cost estimate for a specific model and hardware configuration. Your bill depends on the TPU price, utilization, storage, networking and the throughput your workload actually achieves.

Image: Hugging Face, from its TorchTPU launch post.

Share this article

Leave a Reply

Your email address will not be published. Required fields are marked *

Loading the next article…

Continue reading