This morning I wrote about TorchNeuron and noted that TorchTPU was only a conference session with a repo on the way. That changed a few hours later. Google and Meta open-sourced TorchTPU at PyTorch Conference today, and two serving engines built on it shipped the same day: vLLM-TorchTPU, co-developed with Inferact, and SGLang-TorchTPU, co-developed with RadixArk. The engine repos are Apache-2.0, and the TorchTPU repo's LICENSE file is Apache 2.0 as well.

The headline is the same one every new backend uses: change .to("cuda") to .to("tpu"). The PyTorch blog post is unusually honest about why that sentence is only half the story, and the other half is what anyone running clusters needs to read.

Why this isn't PyTorch/XLA again

Google has shipped PyTorch on TPU before. PyTorch/XLA (2020) was built on lazy tensors, needed code changes to mark compilable regions, compiled slowly, had no real eager mode, and lagged on new distributed primitives. The team's own write-up says all of that directly.

TorchTPU starts from a different set of rules. There's no JAX dependency in the runtime. It registers as a PrivateUse1 backend in the PyTorch dispatcher. torch.compile(model, backend="tpu") goes through Dynamo rather than lazy tensors. Collectives sit at the Process Group level, through a tpu_dist backend that inherits from c10d::Backend, so DTensor, FSDPv2, SimpleFSDP and FairScale see an ordinary process group. It's the same architectural move AWS made with TorchNeuron, and the fact that both arrived on the same day is the real story for anyone who runs mixed fleets.

The coverage numbers at launch:

  • About 450 ATen ops, including forward and backward for the Core ATen set.
  • 4,917 Hugging Face models tested; 4,821 (98%) run without crashing.
  • 638 of 657 architecture representatives (97%) match GPU numerics.
  • With bounded dynamism, 834 of 839 causal LLMs (99%) produce correct numerics.

Eager has three modes, and one of them is new

TorchTPU's eager path can dispatch an op immediately, defer it until something forces a result, or anything in between. That gives three modes:

  • Strict (default): one ATen op at a time, asynchronous to the CPU, like CUDA eager.
  • Debug (TORCH_TPU_LAUNCH_BLOCKING=1): synchronous, with bounds checks. It's the CUDA_LAUNCH_BLOCKING equivalent.
  • DeferAndFuse (TORCH_TPU_DEFER_AND_FUSE=1): ops are recorded as deferred StableHLO builders and compiled together at the next materialization point, such as .item() or .to("cpu"). The post says this can be 50% or more faster than Strict and lowers peak HBM.

DeferAndFuse is the interesting design choice. A "Repeated Ops Heuristic" spots repeated sequences, like a stack of transformer blocks, and materializes early so the TPU isn't left idle.

Compiled graphs land in a three-tier cache: in-process memory, host-local shared memory across processes (on by default for distributed jobs), and an optional persistent tier on a remote filesystem such as GCS. The key is three 64-bit fingerprints covering graph structure, tensor dimensions, and compiler options plus TorchTPU, libtpu and PjRT versions. In practice a libtpu bump invalidates everything, which is correct, but you should plan your warm-up around it.

The parts that are not like a GPU

This is the section I'd hand to a platform team.

One process per TPU device, and only one. A TPU can be opened by a single user-space process. TorchTPU follows PyTorch's one-process-per-device model rather than JAX's one-process-per-host model. The post says plainly that RL colocation, where trainer and rollout share a chip, doesn't work out of the box yet.

A v7 chip shows up as two devices. Each TPU v7 chip exposes its two TensorCores as two PyTorch devices. Their single-chip prefill example runs FSDP with a world size of 2 for that reason. Plan your rank math the way you would for Trainium's NeuronCores, not by counting chips.

Distributed init needs the topology up front. A TPU process sees only its own device, so it has to know the world size and slice type before init_process_group. That means two variables torchrun can't set by itself: TORCH_TPU_TOPOLOGY (for example 2,2,1) and TORCH_TPU_SLICEBUILDER_ADDRESSES, one port per device for building the ICI mesh. python -m torch.tpu.distributed prints them for you on Google Cloud.

Collectives are compiled, not called. There's no NCCL equivalent. All-reduce and friends are StableHLO ops. Because PyTorch code often isn't strictly SPMD (rank-conditional branches, uneven DTensor splits), ranks can compile different modules. With embedded collectives that can deadlock or corrupt memory. So by default each collective is compiled into its own "skinny" module. That's safe, but the post puts the cost at 30 to 40% end to end, because XLA can't overlap communication with compute. You can opt out with an @spmd_safe decorator, or an RFC-stage compile handshake (TORCH_TPU_SPMD_HANDSHAKE_STAGE) that checks FX graphs match across replicas. Mark something @spmd_safe that isn't, and you've built the brittle case the authors warn about.

Default matmul precision is BF16. PyTorch on CUDA defaults to full FP32 for float32 matmuls. TorchTPU defaults to medium, which means BF16 inputs with FP32 accumulation. high and highest emulate FP32 with 3 and 6 MXU passes, so they cost about 3x and 6x, and they print a warning. There's a scoped torch.tpu.precision context manager for sensitive layers. If you're comparing loss curves between GPU and TPU runs, check this first.

Shapes are still mostly static. torch.compile on TPU forces dynamic=False for now. Bounded dynamism, where you compile once against a maximum shape and XLA pads and slices for you, is in active development and not on by default yet.

The numbers they did publish

Out of the box, with no custom kernels, on a self-contained pure-PyTorch Muse Glimmer 30B in the repo:

  • Prefill on one TPU v7 chip: about 25.7k tokens/s at 59% MFU at 5,120 tokens, and about 22.4k tokens/s at 51% MFU at 10,240.
  • Training on a 4-chip host with FSDP: about 5.6k tokens/s per TPU at 38% MFU (5,120), and 4.6k at 32% (10,240).
  • The TorchTitan GPU implementation, with FlexAttention swapped for F.scaled_dot_product_attention because FlexAttention isn't supported yet, reaches 5.3k at 36% and 4.2k at 29%.

At scale, Qwen3 MoE 235B in TorchTitan on 64 v7 chips ran at 1,828 tokens/s against 1,937 for JAX on MaxText, about 94%. Going from one 1,024-chip slice to four slices (4,096 chips) with HSDP took per-TensorCore throughput from 2,236 to 2,095 tokens/s and MFU from 28.8% to 26.9%, with step time going from 29.3 to 31.3 seconds. That run used Tokamax Splash Attention and other Pallas kernels, so it's not the out-of-the-box path.

The serving side

The vLLM-TorchTPU post describes an out-of-tree platform plugin, found through the vllm.platform_plugins entry point. The scheduler, workers, KV cache manager and API server are unmodified upstream vLLM. Models compile with torch.compile(backend="tpu"), and the TPU work lives in Pallas kernels registered as custom ops: Ragged Paged Attention v3, dense and sparse MLA, Gated DeltaNet and Kimi Delta Attention, fused MoE with expert parallelism, and FP8/INT8/MXFP4/NVFP4 matmuls.

Practical details:

  • pip install vllm-torchtpu needs Python 3.12 on a Cloud TPU VM. PyPI shows 0.29.0, uploaded Oct 19.
  • TPU v7x (Ironwood) is the recommended target and runs in CI on every PR. v6e, v5e and v5p are experimental.
  • First startup compiles graphs for your configured shapes and can take minutes. Set VLLM_XLA_CACHE_PATH so restarts reuse them.
  • All six vLLM parallelism dimensions work (TP, DP, PP, EP, prefill and decode context parallel). For large MoE they suggest DP=8, TP=1 with EP on one v7x host.

KV movement goes through TPU Sync, an open-source C++ library that reads device buffers directly, outside Python. vLLM uses it through its KV connector for zero-copy offload to host DRAM and for prefill/decode disaggregation. Prefill and decode can use different layouts, such as PCP=4, TP=2 for prefill and DP=8, TP=1 for decode, with the KV resharded in transit. SGLang uses it for disaggregation and as the HiCache I/O backend, and Google's post adds Managed Lustre as a further tier.

Two vendor claims belong in their own box. Google says its SemiAnalysis InferenceX submission for Qwen 3.5 on Ironwood reached 2.01 cents per million tokens. It also says a Pallas megakernel that runs Kimi K3's whole decode step as one kernel serves more than 700 tokens/s to a single user. Meta says SGLang-TorchTPU is in production there. I haven't reproduced any of this; check the InferenceX entries yourself before you plan around them. Google also says vLLM-JAX and SGLang-JAX will keep being maintained.

What I'd check first

  1. Run your model in Strict eager mode, then DeferAndFuse, and compare outputs before touching torch.compile.
  2. Set torch.set_float32_matmul_precision explicitly on both GPU and TPU so you're comparing the same math.
  3. Point the tier-3 compile cache, or VLLM_XLA_CACHE_PATH, at shared storage and pin libtpu, because a version change invalidates every entry.
  4. Size your ranks from TensorCores, not chips, and build meshes with torch.tpu.topology_aware_mesh() so TP lands on the fastest ICI ring.
  5. Measure with skinny collectives first. Only move to @spmd_safe on code you can show is identical on every rank.

Two native, non-CUDA PyTorch backends shipped in one day. The work of leaving CUDA is now mostly compile caches, shape discipline and topology, and that's work a platform team can do once for everyone.

Sources

Anpoo Sivanadi, Staff Software Engineer