PyTorch on Cloud TPU

PyTorch on Cloud TPU lets you train, fine-tune, and serve open-source models on Cloud TPU hardware using standard PyTorch semantics, distributed APIs, and the ecosystem tools you already use, including Hugging Face, vLLM, SGLang, llm-d, TorchTitan, TorchRec, and Ray.

At the foundation of the stack is TorchTPU, an open-source, PyTorch-native backend co-developed with the PyTorch team at Meta. TorchTPU provides libraries, frameworks, and tools for every stage of the machine learning lifecycle, from core runtime execution and custom kernels to production inference, large-scale training, and cluster orchestration.

TorchTPU connected to Hugging Face, TorchTitan, vLLM, SGLang, Ray, and llm-d, with torch.profiler, DTensor, TorchRec, and Pallas.

Get started

TorchTPU: PyTorch-native backend for Cloud TPU

TorchTPU provides the device backend, graph compiler integration, and custom kernel APIs for PyTorch on Cloud TPU. Built with the PyTorch team, TorchTPU lets you debug models interactively in eager mode, compile model code with torch.compile, and scale across slices with torch.distributed (tpu_dist).

To install TorchTPU and run a tensor operation on a Cloud TPU:

pip install torch_tpu

import torch

device = torch.device("tpu")

x = torch.randn((4, 4), device=device)
y = torch.ones((4, 4), device=device)
result = torch.matmul(x, y).cpu()

print(f"Executed on device: {x.device}")
print(f"Result shape on CPU: {result.shape}, sum: {result.sum():.4f}")

Motivation

While Cloud TPUs have historically been paired with JAX and TensorFlow, running PyTorch on TPUs previously required lazy-tensor graph recording (PyTorch/XLA), which often meant adding code to mark compilable regions. TorchTPU runs operations eagerly by default using PyTorch's standard one-process-per-device execution model, so existing workloads and libraries run without structural rewrites.

Design

TorchTPU brings PyTorch-native TPU support to PyTorch around three design goals:

  • Standard PyTorch semantics: Runs one process per device and preserves PyTorch tensor views and in-place mutations without a JAX runtime dependency.
  • Eager execution by default: Standard PyTorch models run immediately in eager mode for interactive debugging and step-by-step inspection.
  • Compiler and kernel optimizations: In addition to operator-by-operator eager execution, TorchTPU provides operator fusion in eager mode, torch.compile integration with dynamic shape bounds to reduce recompilations, and custom kernel APIs.

For graph compilation, torch.compile captures the PyTorch computation graph with TorchDynamo and compiles it into TPU machine code through OpenXLA and StableHLO. Collective operations run at the torch.distributed process group level (tpu_dist).

Key strengths

  • Operator, model, and Hugging Face support: TorchTPU supports standard PyTorch (ATen) operators in forward and backward passes and is continuously tested against popular Hugging Face models. You can deploy models in Transformers v5 by passing device_map="tpu", run Diffusers pipelines with eager mode, torch.compile, and tensor parallelism, and profile workloads with torch.profiler in TensorBoard and XProf.
  • Single-chip to multislice scaling: The tpu_dist backend supports DistributedDataParallel (DDP), FullyShardedDataParallel (FSDP2), and DTensor. In addition, torch.tpu.distributed.topology_aware_mesh maps parallelism dimensions onto the physical TPU interconnects.
  • Configurable eager execution modes: Provides three eager execution modes:
    • Strict (default): Dispatches operators individually and asynchronously, with standard imperative PyTorch semantics.
    • Debug: Runs synchronously with tensor bounds checking.
    • DeferAndFuse: Fuses sequences of eager operators to reduce High Bandwidth Memory (HBM) traffic.
  • Full-graph compilation and caching: torch.compile(backend="tpu") compiles graphs into TPU executables and caches them across processes, hosts, and job restarts.
  • Custom TPU kernels: You can call custom Pallas kernels, including kernels that communicate directly between chips, as standard PyTorch functions, use prebuilt kernels from the Tokamax library, or author portable tile-based kernels with Helion.

For more information, see the TorchTPU documentation, ask questions in the PyTorch on TPU discussion forum, or report issues on GitHub.

Production inference and serving

For production deployment, PyTorch on Cloud TPU integrates with open-source inference engines and Kubernetes-native distributed serving infrastructure.

vLLM-TorchTPU and SGLang-TorchTPU: Inference engines

vLLM-TorchTPU and SGLang-TorchTPU are TorchTPU backends for the vLLM and SGLang inference engines. Both backends run open models such as Qwen, DeepSeek, GLM, Kimi, GPT-OSS, and Gemma 4. For more information, read the blog post Announcing New vLLM and SGLang Inference Engines on TPU.

Motivation

Serving reasoning, multimodal, and tool-using workloads often requires custom model definitions, speculative decoding methods, and quantization schemes. Building the vLLM and SGLang backends on TorchTPU lets you customize and debug your serving stack directly in PyTorch.

Design

Both backends retain standard vLLM and SGLang schedulers, KV cache managers, and API servers while using TorchTPU execution and custom Pallas operators underneath. Most models run their standard PyTorch model definitions, while compute-intensive blocks invoke Pallas kernels inside the compiled graph so hardware-specific kernels can be updated independently.

Key strengths

  • Parallelism, quantization, and speculative decoding: Supports tensor, data, expert, pipeline, and context parallelism across hosts, along with multimodal (image and video) and embedding inputs. Supported speculative decoding methods include multi-token prediction (MTP), EAGLE-3, DFlash, DSpark, and n-gram matching. Supported quantization formats include 8-bit (FP8) KV caches and 8-bit (FP8) and 4-bit (MXFP4 and NVFP4) floating-point model checkpoints.
  • Two-tier prefix caching, tool calling, and structured output: Automatic prefix caching (including for hybrid attention and recurrent models) retains active prefixes in HBM and offloads older prefixes to host memory to avoid recomputation on multi-turn requests. Tool calling and structured output are built in.
  • Overlapped KV cache transfers with TPU Sync: Supports disaggregated prefill and decode serving (separating prompt processing and token generation onto distinct worker pools) across hosts and host KV offloading using TPU Sync. TPU Sync is an open-source C++ library that reads PyTorch device buffers directly on background threads without copying data through host CPU memory.

For more information about vLLM-TorchTPU, see the vLLM-TorchTPU repository, ask questions in the vLLM Slack channel, or report vLLM-TorchTPU issues on GitHub.

For more information about SGLang-TorchTPU, see the SGLang-TorchTPU repository or report SGLang-TorchTPU issues on GitHub.

llm-d: Kubernetes-native distributed inference

llm-d is an open-source Cloud Native Computing Foundation (CNCF) sandbox framework for distributed LLM inference on Kubernetes, developed with Red Hat, IBM, Google Cloud, and a community of more than 500 contributors and 35 organizations.

Motivation

While single-instance inference engines manage per-replica execution, production fleets must handle traffic spikes, multi-tenant priorities, fleet-wide context sharing, and zero-downtime rollouts. llm-d provides the cluster-level routing and operations layer to coordinate a fleet of inference servers.

Design

llm-d runs in front of model servers such as vLLM and SGLang on Kubernetes. llm-d provides the following:

  • Routes requests based on prefix cache state and replica load.
  • Coordinates disaggregated prefill and decode worker pools.
  • Supports MoE sharding and asynchronous batch inference.
  • Manages tiered KV cache offloading across device memory, host memory, and storage.

Key strengths

  • Multi-turn session routing: Routes multi-turn conversations to replicas that hold their cached context. Turn-priority scheduling and stateful Responses API support let agents reuse context across turns to reduce recomputation.
  • Multi-tenant quality of service: Priority and fairness flow control lets latency-sensitive requests and bursty agent workloads share a fleet without resource starvation.
  • Fleet operations and rolling updates: Provides a v1 configuration API, router high availability, and coordinated rolling updates that manage prefill and decode replicas as a single logical unit.
  • Native TorchTPU serving on GKE: You can serve Cloud TPUs alongside other accelerators behind one llm-d endpoint and manage vLLM-TorchTPU and SGLang-TorchTPU replicas with standard Kubernetes tooling.

For more information, ask questions in the llm-d Slack channel, or report issues on GitHub.

Train, fine-tune, and domain libraries

To pre-train, fine-tune, and adapt large models, TorchTPU provides reference training recipes and domain-specific libraries tuned for TPU hardware.

TorchTitan on TPU: Large-scale training reference

TorchTitan on TPU extends TorchTitan, the PyTorch modular reference platform for generative AI training, with TPU sharding passes, custom kernels, and training recipes. For more information, read the blog post Integrating TorchTPU into TorchTitan and Optimizing Models for TPU.

Motivation

Scaling training across multi-host TPU slices requires sharding parameters across the chip interconnect and overlapping communication with computation. Google and Meta co-developed TPU sharding and compilation passes inside TorchTitan so you don't need to design custom TPU sharding strategies or manually schedule communication.

Design

Models are written in single-device PyTorch without accelerator-specific code. TorchTitan configures model layers from configuration files, letting you replace individual layers with custom Pallas kernels (such as Splash Attention) without editing the model definition or training loop.

TorchTitan's SimpleFSDP implementation places parameter sharding operations directly inside the compiled graph so the XLA compiler can overlap device communication with computation, while whole-step graph compilation captures the entire training step as a single XLA program. For multislice runs, hybrid sharding (HSDP) shards within each slice over the inter-chip interconnect (ICI) and replicates across slices over the data center network (DCN). This layout limits cross-slice traffic to a single gradient reduction per step.

Key strengths

  • Pre-training and RL post-training recipes: Includes reference recipes for dense LLMs (Llama 3, Qwen 3, and MuseGlimmer 30B), mixture-of-experts (MoE) models (Qwen 3 235B, validated on up to 8,192 TPU7x chips), and diffusion transformers (FLUX.1). For RL post-training (such as Group Relative Policy Optimization, or GRPO), VeRL (Preview) delegates parameter updates to TorchTitan and pushes updated weights to serving engines to generate training trajectories (rollouts).
  • Preconfigured compiler and parallelism settings: Recipe configurations include preconfigured compiler and parallelism settings so you can tune Model FLOPs Utilization (MFU) through configuration files rather than code changes.
  • Onboarding for new architectures: The Hugging Face Transformers modeling backend runs Hugging Face Hub model checkpoints (including MoE models without expert parallelism) through the TorchTitan trainer. For full model ports, TorchTitan includes a step-by-step onboarding skill for engineers and coding agents.
  • Standalone component reuse: You can adopt TorchTitan end to end or import individual modules, custom operators, and compiler flags into an existing training codebase.

For more information, see the TorchTitan on TPU getting started guide, ask questions in the PyTorch discussion forum, or report issues on GitHub.

TorchRec: Recommendation systems

TorchRec is a PyTorch domain library that provides distributed embedding tables and sharding primitives for recommendation models. On Cloud TPU, TorchRec offloads sparse embedding table lookups to dedicated TPU SparseCore processors while running dense matrix operations on the TPU TensorCores.

For more information about TorchRec, see the TorchRec documentation, ask questions in the PyTorch on TPU discussion forum, or report issues on GitHub.

Cluster orchestration and developer tools

To scale multi-node pipelines and automate model onboarding, the TorchTPU ecosystem integrates with Ray and provides coding agent extensions with access to TPU documentation.

Ray: Distributed training, serving, and RL

Ray scales Python and AI workloads across clusters using Ray Train, Ray Serve, Ray Data, Ray Core, and RL frameworks such as verl.

Motivation

Multislice training and multi-stage pipelines (such as RL workflows that combine data preprocessing, training, and rollout generation) require automated slice provisioning, topology discovery, and fault-tolerant worker coordination. Standalone launch utilities such as torchrun don't provision TPU slices or coordinate multi-stage actor pipelines, so Ray manages cluster-wide scheduling and placement for TorchTPU workloads.

Design

When you submit a job specifying the TPU generation, topology, and worker count, Ray reserves complete TPU slices so all workers in a slice are scheduled together, discovers the chip interconnect addresses, and starts one PyTorch worker per TPU device. Ray Train's PyTorch backend initializes the tpu device and tpu_dist process group automatically before starting the job.

Key strengths

  • Topology-matched placement: Matches worker placement across hosts to the physical TPU generation and slice layout.
  • Automated failure recovery: Runs distributed training jobs across multiple TPU slices on GKE with automated recovery when a worker or slice fails.
  • Direct-to-TPU data prefetching: Ray Data's iter_torch_batches method prefetches batches as PyTorch tensors directly into TPU device memory.
  • Multi-host inference orchestration: Manages multi-host serving deployments for vLLM-TorchTPU and SGLang-TorchTPU.

TPU agent suite: (Preview) Coding agents for TPU

The TPU agent suite is an open-source extension for coding assistants—including Claude Code and Google Antigravity, with more coding agents to follow—focused on transformer inference (dense LLMs and MoE models) on TPU v6e and TPU7x.

Motivation

General-purpose coding assistants often lack current context on TPU hardware topologies, XLA compiler flags, XProf telemetry, and Pallas kernel constraints. The TPU agent suite gives your assistant access to official TPU documentation and structured workflows. This helps prevent invalid flag suggestions and automates iterative profiling and benchmarking.

Design

The suite adds a TPU knowledge base and six task-specific sub-agents to your existing coding assistant. A router agent delegates each prompt to the appropriate sub-agent without running commands itself. Any command execution on your TPU VM or GKE cluster requires your approval, and tuning loops run a bounded number of recorded attempts before returning results.

Key strengths

  • Model onboarding and accuracy verification: The TPU Model Onboarding agent resolves unsupported operators, selects a serving path (vLLM-TorchTPU or standard PyTorch), and compares layer-by-layer outputs against the CPU reference model to verify numerical accuracy.
  • Crash and bottleneck diagnostics: The TPU Diagnostics agent analyzes crash logs, inspects XProf profiles and XLA compiler traces, and isolates runtime bottlenecks.
  • Kernel, compiler, and sharding optimization: The TPU Kernel Development agent assists with custom Pallas kernels, the TPU Performance Optimization agent benchmarks XLA compiler flag combinations, and the TPU Sharding agent plans multislice parallelism layouts.
  • Documentation lookup with citations: The TPU Knowledge Base agent answers technical questions with explicit citations to official TPU documentation and tutorials, and states when a topic isn't covered.

For more information, see the TPU agent suite documentation, share feedback and ask questions in GitHub Discussions, or report issues on GitHub.

What's next