TorchTPU Migration Guide: Native PyTorch on Google TPUs
Find the official TorchTPU, vLLM and SGLang repositories, compare TorchTPU with PyTorch/XLA, and plan a migration with correctness, compilation and serving checks.
BY RUIQI ZHOU · AIINFRAINTERVIEWS EDITORIAL · UPDATED NOVEMBER 9, 2026 · 5 MIN READ
PRACTICE THIS:GPU and accelerator questions ·Inference and serving questions ·Calculate KV-cache memory
TorchTPU is an open-source native PyTorch backend for Google TPUs. It supports eager execution and compilation through torch.compile, giving PyTorch developers a migration path that does not require rewriting every model in JAX. Google announced the release on October 21, 2026 and now recommends it over PyTorch/XLA for new workloads.
That is a substantial change in the programming interface. It does not establish that an arbitrary CUDA extension, quantization format or distributed layout will work unchanged. This guide identifies the public repositories and the evidence needed for a port. Sources were checked on November 10, 2026; no TPU execution is claimed here.
The repositories to start from
| Task | Official project | What to inspect first |
|---|---|---|
| Run ordinary PyTorch | google-pytorch/torch_tpu | Versioned docs, operators, examples |
| Serve with vLLM | vllm-project/vllm-torchtpu | Model matrix and matching engine build |
| Serve with SGLang | sgl-project/sglang-torchtpu | Release-specific engine/plugin pair |
| Distributed training | google-pytorch/torchtitan | TPU experiment and model recipes |
| Assisted porting and profiling | google-pytorch/tpu-agent-suite | Preview installation and compute profile |
These are separate projects. Installing the device backend does not install every production-serving integration. The Google Cloud entry page connects them and describes the supported workflow. The TorchTPU release announcement supplies the release context.
What changes from PyTorch/XLA?
PyTorch/XLA's familiar workflow records operations lazily and uses execution boundaries to materialize work. TorchTPU integrates with the PyTorch dispatcher, exposes the tpu device, and offers ordinary eager behavior. Internally, it can compile individual operations or grouped operations through XLA. Eager semantics and compiled machine execution can coexist.
In compiled mode, torch.compile(backend="tpu") captures a graph and lowers it toward StableHLO and XLA. StableHLO is an intermediate representation of tensor computations; XLA generates hardware-specific execution. The architecture documentation explains these boundaries. The Premium TPU software-stack reference compares them with JAX.
Step 1 of 4: PyTorch model
- PyTorch model
Identify standard operators, custom extensions and explicit CUDA assumptions.
- TorchTPU dispatch
Run a small fixed input and compare the output with the reference.
- Compiled region
Capture supported operations; compare eager and compiled results before timing.
- Distributed or serving workload
Add sharding, request scheduling and recovery one layer at a time.
Separate four migration risks
First, inspect device assumptions. Code may construct CUDA tensors internally, choose a CUDA-only attention implementation or depend on a custom extension. Merely moving the top-level model does not replace those paths. Prefer an existing tested TPU implementation before writing a kernel.
Second, compare numerical behavior. Fix model weights, input batches and preprocessing. Check outputs with tolerances appropriate to the selected precision; for training, check gradients and several optimizer updates. Exact text equality is often too strict for stochastic generation and too weak for diagnosing where a numerical difference began.
Third, examine shapes and host synchronization. Calls that read a device result into Python can make the host wait. Changing tensor dimensions can create new specializations or exercise bounded dynamic-shape support. Track compilation events, shape combinations and steady-state timing together; a faster warm step can still have an unacceptable cold start.
Finally, validate the distributed or engine integration. A backend supporting a collective does not prove the model's sharding is efficient. Check each rank's memory, the communication groups and checkpoint restoration. For serving, compare cancellation, stop conditions, prefix reuse and structured outputs on the exact engine release.
Find the first boundary that fails
Use the same saved weights and inputs for a reference run, TPU eager execution and TPU compiled execution. Begin with one layer or a small submodule whose outputs you can inspect, then widen the comparison. If eager TPU already disagrees, compiling the whole model adds another variable before the first discrepancy is understood.
| Observation | Next check |
|---|---|
| Import or device allocation fails | Check the installed backend, runtime and supported TPU generation. |
| Eager output differs | Match weights, inputs, masks, dtype and arithmetic precision; isolate the first differing layer. |
| Eager agrees; compiled differs | Reduce to the failing graph region and inspect mutation, aliasing and shape assumptions. |
| Both agree; the service is slow | Separate compilation, host synchronization, queueing and device execution in the trace. |
One subtle case is a tensor stored as FP32 whose matrix multiplication uses reduced input precision. The TorchTPU 0.2.0 precision guide documents a default medium mode and higher-precision modes that perform additional MXU passes. A shared dtype name therefore does not establish matched arithmetic across two backends. Record the selected matmul mode alongside the dtype, and compare quality before choosing the faster mode.
Use both absolute and relative error. Relative error alone becomes unstable around a reference value of zero; a loose absolute threshold can conceal substantial errors in small outputs. Check non-finite values and a task-level quality measure as well. For a ranker, a small numerical difference matters if it reverses important candidate orderings. For training, extend the check through gradients, updates and restoration from a saved checkpoint.
Keep diagnostic instrumentation out of the final timing loop. Copying every intermediate tensor to the host is useful for finding a mismatch, but it changes synchronization and traffic. Remove those copies after locating the issue, synchronize at the measurement boundary, and record compilation separately. A loop timing asynchronous submissions can report a speedup while the device is still processing its queue.
Compilation has a break-even point
Suppose a hypothetical compiled path adds 40 seconds of setup but reduces a synchronized step from 12 ms to 8 ms. The extra work pays back after:
Saving per step = 0.012 - 0.008 = 0.004 seconds
Break-even count = 40 / 0.004 = 10,000 steps
At 1,000 steps: eager 12 s; compiled 40 + 8 = 48 s
At 100,000 steps: eager 1,200 s; compiled 40 + 800 = 840 s
These invented timings illustrate the decision; they are not TorchTPU measurements. A persistent compilation cache can change the setup cost, while a new shape or incompatible runtime can change cache reuse. Measure the actual restart path rather than subtracting compilation from every result. The free serving benchmark reference explains how to define the measurement interval.
Read the matrix at the layer you use
The TorchTPU 0.2.0 feature table lists v5e, v5p, v6e and v7 in its validated hardware row, while TPU 8 is listed as near-term work. It marks FSDPv2 as experimental with performance limitations. This is more specific than saying that the ecosystem supports native distributed PyTorch.
The vLLM-TorchTPU repository separately recommends , lists v6e as supported without CI coverage, and calls v5e/v5p experimental. SGLang-TorchTPU's reviewed instructions focus on Ironwood and a matched engine/plugin release. These distinctions are reasons to select a known recipe, not reasons to assume all older hardware is equivalent.
Record a container digest or exact package set alongside the model revision, topology and test inputs. Link the corresponding release documentation in the run record. Follow the official installation instructions for that combination; mixing a current backend with an arbitrary engine wheel creates an additional compatibility experiment.
A practical acceptance record
For the first port, require a numerical comparison, a synchronized warm measurement, a cold-start measurement and a restart test. For a service, also record offered requests, completed requests, failures, first-token latency and token gaps. A throughput gain that violates the service target is not usable capacity.
Use TPU Agent Suite to organize some of this investigation if its preview supports your environment. For hardware fit, use the generation comparison. Continue to inference-serving questions when you can explain which result would make you keep the original deployment.
Turn it into offers. Work the real questions and concepts this maps to:
FAQ
Yes. Google announced its open-source release on October 21, 2026. The official repository is google-pytorch/torch_tpu; the documentation reviewed here is version 0.2.0.
Google recommends TorchTPU for new PyTorch workloads and migration from PyTorch/XLA. TorchTPU offers native eager semantics and torch.compile integration. Existing workloads still need numerical, operator and distributed-execution validation.
No universal CUDA-kernel conversion is established by the release. Standard PyTorch operations can dispatch through the backend, while custom CUDA extensions and device-specific assumptions require a separate implementation or supported replacement.
The native PyTorch path is vllm-project/vllm-torchtpu. The older tpu-inference project is now vLLM JAX. Choose a backend with a tested recipe for the model and features you actually need; their support matrices differ.
Discussion (4)
Put restart recovery in the migration acceptance criteria. A training step that matches numerically is only half a port if the restored optimizer state changes the next update.
I would compare a short loss trajectory as well as the first forward pass. A backward implementation can be wrong while the logits look perfectly reasonable.
Keep the same batches and initial parameters for that check. Different data order makes the discrepancy much harder to interpret.
An OpenAI-compatible endpoint also needs the same stop behavior and streaming contract. A backend migration can pass a tensor check and still break a client.