Google and Meta’s TorchTPU Runs PyTorch on Google’s Chips

Google and Meta have opened a new route for AI developers who want their PyTorch models to run on Google’s Tensor Processing Units, or TPUs. Released last week, TorchTPU brings those chips into familiar PyTorch workflows, reducing the need to rebuild models in another framework. Announcement

The simplest examples change the target device in ordinary PyTorch code. Serving a model also requires a working combination of hardware, kernels and runtime settings. The project’s tests distinguish a successful model port from a configuration that has undergone repeated serving checks. Paper Support

PyTorch is the software developers use to define and train models. TorchTPU supplies a backend, the implementation that makes its operations execute on TPU hardware. The repository includes optimized implementations of PyTorch operations, a cache of compiled computations, distributed training and model examples. Its public release is version 0.2.0 and requires Python 3.11 or later. The project uses the Apache 2.0 license. The install command deliberately selects the CPU build of PyTorch; the separate TPU package supplies accelerator execution. Code

Meta’s director of PyTorch engineering, Jana van Greunen, described hardware diversity as the reason for the collaboration in Google’s announcement. She said the work supports major language-model, advertising and recommendation workloads and gives Meta a way to deploy on TPUs while providing feedback to the engineering teams. Google also says Transformers v5 can select TPUs through its device setting, with nightly tests of the library. Announcement

What the device change covers

TorchTPU’s examples range from creating a tensor on one device to training across multiple machines. The introductory code asks PyTorch for the current accelerator, creates an array of numbers there and adds it to itself. More advanced examples demonstrate how to split model calculations or divide stored parameters between devices. The examples page says these demonstrations establish functionality and have not been optimized for TPU performance. Examples

The paper reports a pure-PyTorch Muse Glimmer 30B example processing input at 25,700 tokens a second on one TPU v7 chip with 5,120-token sequences. That chip exposes two PyTorch devices. This measures prefill, which processes input, rather than the speed at which a chatbot writes its answer. Paper

A separate Muse implementation originally written for GPUs required replacing FlexAttention with standard scaled-dot-product attention, then applying compilation and distributed sharding. The authors say FlexAttention was not yet supported. Their examples show portability for specific PyTorch implementations, including an identified change where a required operation differs. Paper

Public interface support also has levels. TorchTPU lists ordinary accelerator functions for checking devices, reading memory usage and synchronizing execution as stable interfaces. TPU-specific controls for execution modes, streams, numerical precision and many custom operations remain experimental. Stable functions go through a formal process before incompatible changes or removal; experimental ones can change without that process. Code using experimental TPU controls may need revision when the runtime changes. APIs

What the serving stack has tested

For serving models to users, the project provides vllm-torchtpu, a plugin for the widely used vLLM engine. It preserves PyTorch model definitions while adding TPU-specific kernels, optimized operations for attention and lower-precision arithmetic. Its installation requirements are more specific than the core package: Python 3.12 on a Cloud TPU virtual machine. Serving

Hardware support differs too. The serving repository recommends Ironwood, or TPU v7x, which is checked on every proposed code change and in nightly tests. Trillium, or v6e, is supported without that automated coverage. The older v5e and v5p generations are marked experimental. These labels identify the hardware paths receiving different levels of routine testing. Serving

The model table separates checks on every code change from nightly runs, component tests and example recipes. Qwen3.5-35B has a regularly tested low-precision FP8 checkpoint; the quickstart’s Llama-3.2-1B is listed as an example. For Kimi K3, routine testing uses a version with fewer experts, while the full model gets on-demand tests on a larger slice. A shared model-family name can therefore conceal materially different test configurations. Support

The same table says structured output has component tests and examples, while some multimodal Gemma and Qwen configurations have integrated checks. It also cautions that an absent model-feature pairing means that combination has not run in continuous integration, rather than establishing that it cannot work. The relevant unit of support is the particular model, feature and configuration. Support

Core-model coverage is broader than those serving recipes. The authors’ paper reports 4,821 of 4,917 models passing checks for execution without crashes. Separately, 638 of 657 numerical-comparison checks passed against a GPU reference. Those are different tests: a program can execute while producing outputs that still need comparison. They do not represent thousands of independent production-service benchmarks. Paper

Compilation and the next serving milestones

Starting a server also involves preparing code for the chip. The serving repository warns that the first startup can take several minutes to compile the configured shapes. Its dedicated cache guide distinguishes saved Python tracing artifacts from the native compiled executables. Each has a separate storage setting, so retaining one cache does not automatically preserve the other. Serving Cache

The cache guide also describes a detector for unexpected compilation during serving. An input shape, data type or code path missed during warmup can create a cache miss and stall the engine. The project recommends enabling the detector in testing and disabling it during production benchmarking to avoid its monitoring overhead. This gives developers a concrete diagnostic for a service that runs but pauses on unfamiliar requests. Cache

Google has also released an early-preview TPU Agent Suite to help with onboarding, profiling and kernel optimization. Its current packages support Claude Code and Antigravity; Codex and Cursor are listed as coming later. The suite’s prerequisites still include cloud or Kubernetes tools, credentials and a configured compute target. An agent can help operate the infrastructure, while the underlying environment remains part of deployment. Agents Discussion

The roadmap, dated October 19, targets DeepSeek V4 support, optimization and AgentX benchmarks for December. The current support table already lists two DeepSeek V4 checkpoints in nightly testing. That reports progress on specific configurations while leaving room for further optimization and benchmark work. Roadmap Support

One November target is a pipeline that would ingest, benchmark and publish support information for newly released open models within 24 hours. Another calls for public nightly checks across model versions, numerical formats and serving features. Roadmap