- Book: LLM Observability Pocket Guide
- Also by me: Thinking in Go (2-book series) — Complete Guide to Go Programming + Hexagonal Architecture in Go
- My project: Hermes IDE | GitHub — an IDE for developers who ship with Claude Code and other AI coding tools
- Me: xgabriel.com | GitHub
On October 20, Google put TorchTPU on GitHub, a native PyTorch backend for its TPUs. For anyone already running PyTorch on TPUs, the line to read is in the TorchTPU FAQ: "PyTorch/XLA is deprecated and replaced by TorchTPU, which is now public."
The PyTorch/XLA README now opens with the same message: "This project is deprecated and replaced by TorchTPU." It also says the repository will be archived on November 15, 2026. That gives teams on torch_xla under four weeks before the project goes read-only.
What Google shipped
Google announced TorchTPU at PyTorch Conference 2026 and says it was built together with the PyTorch team at Meta (Google Cloud blog). In Google's words, TPUs now join GPUs and CPUs as native PyTorch targets, with eager mode for debugging and torch.compile for production speed.
The repo is Apache 2.0, and the latest public release is 0.2.0 on PyPI, supporting Python 3.11 and up (README). The feature set page lists TPU v5e, v5p, v6e and v7 as validated. TPU v8i and v8t are on the near-term roadmap, not supported yet.
The serving story ships with it. Google's post links a vLLM backend built with Inferact and an SGLang backend built with RadixArk, and says Hugging Face Transformers v5 deploys to TPUs with device_map="tpu" (Google Cloud blog). On training, Google reports up to 52% MFU on Qwen 3 8B with TorchTitan, scaling close to linearly across up to 8,192 Ironwood TPUs. Those are Google's own numbers, not an independent benchmark.
TechCrunch frames the release as Google joining Amazon in going after Nvidia's hold on AI developers, one day after AWS announced its own PyTorch backend for Trainium. I covered that one in the TorchNeuron post.
The port starts with one string
The FAQ describes the change for existing workloads as .to("cuda") → .to("tpu"). Install the release with the CPU build of torch, because the README warns the default install pulls the CUDA build (README):
pip install torch_tpu==0.2.0 \
--extra-index-url https://download.pytorch.org/whl/cpu
Then the code from the VM getting-started guide, trimmed:
import torch
device = torch.device("tpu") # was "cuda"
a = torch.ones((5, 5)).to(device)
b = torch.randn((5, 5), device="tpu")
result = (a + b).cpu() # back to the host to print
print(result.sum())
The guide runs this on a TPU VM, with a v6e instance as its example. It also notes that recent TPU hardware needs GKE instead of a VM (guide).
Eager by default, and where it stalls
The execution model page explains what happens under that one-line change. In the default strict eager mode, each op becomes a single-op graph, compiled at -O1 and cached, so the cache warms quickly. torch.compile hands TorchTPU the whole FX graph and compiles it at -O2. Two environment variables sit in between:
# Errors surface at the op that caused them
TORCH_TPU_LAUNCH_BLOCKING=1 python train.py
# Fuse deferred ops into bigger graphs
TORCH_TPU_DEFER_AND_FUSE=1 python train.py
The docs claim up to 50% speedup over default eager for TORCH_TPU_DEFER_AND_FUSE, with no model changes. The same page lists what stalls the device. Data-dependent Python branches like if my_tensor.all() copy a value to the CPU and drain the queue, and in compiled or fused mode they also break the graph. Ops with data-dependent output shapes, such as torch.masked_select, do the same.
Growing shapes don't force a recompile on every step. The dynamic shapes page walks through a decode loop with a growing KV cache. With dynamic=False, every step recompiles. With the default dynamic=None, the growing dimension is treated as symbolic, and in the docs' example only the first three steps pay compile time. Under the hood TorchTPU compiles buckets aligned to multiples of 128 (one by default, up to eight) and pads inside them:
step = torch.compile(
decode_step, backend="tpu", dynamic=None
)
# or declare the range up front
torch._dynamo.mark_dynamic(kv_cache, 2, min=16, max=128)
The limits are on the same page: at most eight buckets per trace, no dynamic=True, no torch.nonzero-style output shapes, no custom Pallas kernels in that path. These bounded-dynamic executables are not saved to the AOTAutograd cache, so they recompile on every process start.
If you're on PyTorch/XLA today
The FAQ draws the line between the two: PyTorch/XLA leaned on JAX runtime components and its own APIs, while TorchTPU is a native backend with no JAX runtime dependency. It also says TorchTPU is not replacing JAX.
Start by finding every place your code touches the old stack:
git grep -n -e "torch_xla" -e "xla_device" -e "xm\."
Then check each hit against the feature set. Mixed eager and compiled runs, forward and backward passes, Pallas kernels and SimpleFSDP are listed as validated. FSDPv2 and SymInts are experimental, and FSDPv2 is marked as low performance. Once November 15 passes, the PyTorch/XLA repo becomes read-only: you can still read and clone the code, but nobody can open new issues or pull requests.
The TPU agent suite plugs into Claude Code and Antigravity
The second launch is the TPU agent suite, which Google's post calls a preview. It is a package of skills and subagents for coding agents, with one entry point, a tpu-agent skill, that routes to specialists for onboarding, profiling, hardware checks, cited TPU facts and kernels. Claude Code and Antigravity are supported; Codex and Cursor are listed as "coming soon". The repo does not accept outside contributions for now.
Read what the installer writes before you run it. Per the install guide, a project-scope install for Claude Code adds skills under .claude/skills/, one prompt per agent under .claude/agents/, a .claude/settings.json that pre-authorizes the ~/.tpu-agent-suite path, and a managed block in your CLAUDE.md. For Antigravity it writes AGENTS.md instead. Pass --no-workspace-memory to skip that block. Install on a branch and read git status and git diff before you commit anything.
One trap if you run agents in worktrees: a fresh git worktree checks out tracked files only. Uncommitted project-scope files stay in the folder you installed them in. The guide's --scope user install avoids that.
I build Hermes IDE, so weigh this with that in mind. Both agents the suite supports today are built-in agents in Hermes, and each task can get its own hermes/<name> branch and worktree, so a Claude Code run and an Antigravity run on the same TPU port stay out of each other's files (README). I haven't run the TPU agent suite myself. Hermes is not affiliated with Google, Anthropic or any agent vendor.
Hermes 2.3.4's agent picker: Claude Code and Codex installed, Antigravity CLI and the other supported agents listed, plus Custom agent.
Would you move?
I haven't run TorchTPU, and the limits above are where a real port will find out. The deprecation means teams on PyTorch/XLA don't get to wait and see.
Are you running PyTorch on TPUs today, and is the November 15 archive date on your roadmap yet? Tell me in the comments.
If this was useful
If you do move serving onto TPUs with vLLM or SGLang, the next question is whether the move paid off: latency, tokens per request, cost per answer, and evals that tell you the outputs didn't change. The LLM Observability Pocket Guide covers how to pick the tracing and evals tools that answer that.
If you run Claude Code next to Antigravity or Codex, Hermes IDE is source-available (BSL 1.1) and runs on macOS, Windows and Linux. Source is on GitHub.
Top comments (0)