Today for AI
ALL INTELLIGENCE
vLLM Integrates TorchTPU for Native Google TPU InferenceHacker News AIT2Open Source

Hacker News AI · 10/21/2026, 00:40:24

Google Open-Sources TorchTPU Native Backend for PyTorch on TPUs

78AI Score
Executive Summary

Google has open-sourced TorchTPU, a native backend enabling PyTorch models to run directly on Tensor Processing Units (TPUs). Featuring custom ATen kernels, compilation caching, and support for distributed training and torch.compile(), it simplifies deploying models like Llama and Qwen on TPU hardware.

SOURCE COVERAGEFull Original Report

Contents9 sections
<img src="docs/images/torchtpu_logo.avif" alt="TorchTPU Logo" style="width: 60%;">

Description

This project provides a PyTorch backend for Google's Tensor Processing Units (TPUs), enabling users to run PyTorch models and operations on TPU devices. It includes custom ATen kernels for various PyTorch operations, a compilation cache for optimizing execution, and examples for distributed training and model inference.

Features

  • PyTorch ATen kernel implementation for TPU.
  • Compilation cache for optimized execution of XLA computations.
  • Support for distributed training (Data Parallelism, Tensor Parallelism).
  • Examples for various models (Llama, Qwen, ResNet, DLRM, minGPT, UNet).
  • Integration with torch.compile() for graph mode execution.
  • Utilities for debugging and benchmarking.

User Guide and Interactive Notebooks

A user guide is available at https://torchtpu.readthedocs.io/. This is a static version of the interactive notebooks. They are available in the docs folder.

The notebooks are Marimo (similar to Jupyter). They can be run with marimo edit file.py and then you can access them in your browser. You will need to be running this from a TPU VM, a v6e GCE instance is easiest.

Installation

The same instructions are available in the user guide linked above.

TorchTPU's Python version support policy matches PyTorch's Release Compatibility Matrix. Currently, Python 3.11+ is supported. See the Python Version Support Policy for more details.

First create a Python virtual environment:

BASH
python -m venv .venv
source .venv/bin/activate

or, with uv

BASH
uv venv .venv
source .venv/bin/activate

Install via pip

Install the latest public release (0.2.0 on PyPI) (for uv prepend uv to all pip commands):

BASH
pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu

Build from source

Wheels can be built via:

BASH
bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel

All OSS Bazel builds default to optimized mode (-c opt) automatically via .bazelrc. Unoptimized builds (-c fastbuild, -c dbg) can degrade performance and are not used for wheel builds. If you are developing or debugging C++ code and explicitly want an unoptimized wheel, pass --//:allow_unoptimized_wheel=True:

BASH
bazel build -c dbg --config=wheel_common //ci/wheel:torch_tpu_wheel \
    --//:allow_unoptimized_wheel=True

--config=wheel_common is required, and the build fails at build time without it. The config sets --//:wheel_build=True, which routes the shared XLA/MLIR backend through the torch_version-reset transition so that pywrap factors it into a single libxla_base.so; built without the flag, every per-version common ships its own copy of the backend and the resulting wheel aborts on import torch with duplicate static registrations.

wheel_common also turns on the Remote Build Execution (RBE) remote cache and remote executor. If you do not have RBE credentials, append --config=no_rbe to strip those flags back out:

BASH
bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel --config=no_rbe

The order matters here: --config=no_rbe has to come after --config=wheel_common, because the later config is the one that wins for the remote-cache and remote-executor flags they both set.

Install wheel via:

BASH
cd <path_to_repo>
python -m venv .venv; source .venv/bin/activate
# The extra-index-url provides the CPU version of torch rather than the default CUDA build.
pip install ../bazel-bin/ci/wheel/*.whl --extra-index-url https://download.pytorch.org/whl/cpu

Note that installing without --extra-index-url https://download.pytorch.org/whl/cpu will install the CUDA version of torch; pass --extra-index-url as shown above or install the CPU build separately.

Getting Started

Some pointers to get you started:

Dependencies

The project uses Bazel for dependency management. Key Python dependencies are:

TBD

Bazel setup

This project uses Bazel for building and dependency management.

To install Bazel, we actually install Bazelisk, which is a transparent wrapper that handles downloading and installing Bazel itself.

There are several ways to install Bazelisk. Here are some common ways:

  • If you have Go installed:
Text
    go install github.com/bazelbuild/bazelisk@latest
    ```

*   Download the binary directly:
Text
wget -O ~/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.27.0/bazelisk-linux-amd64
```

See the Bazelisk README for more install instructions

However you install it, add it's location to PATH.

To verify the install, run bazel info in the repo's root directory:

Text
bazel info

It should print information about the repo without errors.

（注：更多技术实现细节与完整 API 文档请查阅原项目 README）