Google opens TorchTPU to make TPUs a native PyTorch target across the AI stack

Google has open-sourced TorchTPU, its native PyTorch stack for Tensor Processing Units, broadening the TPU story from a cloud-specific compiler path into a shared foundation for training, serving and orchestration tools across the PyTorch ecosystem.

The launch matters to Red Hat platform teams because two projects central to the company’s open AI stack are already connected to it: vLLM for model serving and llm-d for Kubernetes-native distributed inference. Google says the same TorchTPU foundation also reaches SGLang, Hugging Face Transformers and Diffusers, TorchTitan, VeRL and Ray.

What changed

TorchTPU lets existing PyTorch workloads target TPUs through familiar APIs, using eager mode during development and torch.compile for optimized execution. Google says it co-developed the stack with Meta’s PyTorch team and has now published it on GitHub.

For inference, the integration goes beyond basic device support. The separately released vLLM TorchTPU plug-in reuses vLLM’s upstream scheduler, workers, KV-cache manager and OpenAI-compatible API server while putting TPU-specific compilation and Pallas kernels underneath them. The plug-in supports continuous batching, quantization, multi-chip parallelism, prefix caching, speculative decoding, multimodal inputs and all six vLLM parallelism dimensions.

Google also positions llm-d as the fleet layer above individual inference engines. Its launch post says llm-d can expose a common endpoint on Google Kubernetes Engine while routing workloads to GPUs or TPUs, with prefix-cache-aware and session-aware routing plus separated prefill and decode. A quoted user, Zenteiq, reports raising aggregate throughput from 1.47 billion to 5.97 billion generated tokens per day per 1,000 TPU v6e accelerators with llm-d’s routing; that is a customer-reported result, not an independent benchmark.

Who should care

Teams standardizing on PyTorch gain another accelerator target without adopting an entirely separate application framework. For OpenShift AI and Red Hat AI Inference users following vLLM and llm-d, the immediate implication is architectural rather than a new Red Hat product feature: upstream serving and routing components can now span TPU and GPU fleets behind a more consistent software interface.

The wider ecosystem support is equally important. Hugging Face is adding native TPU paths to Transformers and Diffusers, TorchTitan covers distributed pre-training, VeRL adds a TPU backend for reinforcement-learning workflows, and Ray can schedule heterogeneous pipelines across CPUs, GPUs and TPUs.

What to do

The practical starting point is the vLLM TorchTPU plug-in rather than a wholesale platform migration. Its current documentation recommends TPU v7x, while v6e, v5e and v5p remain experimental. Teams should validate model coverage, compilation warm-up, quantization paths and multi-host behavior against the published support matrices before treating hardware portability as operational equivalence.

The larger test will be whether these integrations stay upstream and interoperable as model architectures change. Open sourcing the base stack creates that possibility; repeatable cross-hardware behavior will determine whether it becomes a durable production layer.