PyTorch-native TPU back end for vLLM

VLLM released vllm-torchtpu, a PyTorch-native TPU backend packaged as an out-of-tree vLLM platform plugin that runs vLLM's standard PyTorch model definitions on Google Cloud TPUs via TorchTPU (torch_tpu, through torch.compile(backend="tpu")) plus custom JAX/Pallas kernels for attention, Mixture-of-Experts and quantized matmul. The plugin requires Python 3.12 on a Cloud TPU VM, is recommended and CI-validated on TPU v7x (Ironwood) while TPU v6e/v5e/v5p remain experimental and outside CI, and keeps the same CLI, OpenAI-compatible server and model definitions with no model code changes. The release also rebrands the original TPU backend tpu-inference as vLLM JAX, which remains available for pure JAX model implementations, with vllm-torchtpu now the recommended path for PyTorch models on TPU.

| Documentation https://vllm-torchtpu.readthedocs.io/en/latest/ | Blog https://blog.vllm.ai/ | User Forum https://discuss.vllm.ai/c/hardware-support/google-tpu-support/27 | Developer Slack https://slack.vllm.ai sig-tpu | vLLM TorchTPU is the PyTorch-native TPU backend for vLLM https://github.com/vllm-project/vllm . It is an out-of-tree vLLM platform plugin, packaged as vllm-torchtpu , that runs vLLM's standard PyTorch model definitions on Google Cloud TPUs through TorchTPU torch tpu , via torch.compile backend="tpu" and a library of custom JAX/Pallas kernels for attention, Mixture-of-Experts, and quantized matmul. The goal is simple: the same vLLM you already use — same CLI, same OpenAI-compatible server, same model definitions — running at TPU speed with no model code changes. - PyTorch first. Models run through vLLM's native PyTorch runner; no re-implementation in another framework is required to serve on TPU. - TPU-tuned kernels. Ragged Paged Attention, dense and sparse MLA, Gated DeltaNet, fused MoE including SparseCore offload , and FP8/INT8/MXFP4 quantized matmuls written in Pallas. - Scale-out built in. TP, DP, PP, EP, and prefill/decode context parallelism compose across single-host and multi-host TPU slices, with disaggregated prefill/decode serving. - Standard vLLM experience. Speculative decoding, prefix caching, multi-LoRA, structured output, and multimodal inputs work through the normal vLLM flags. Note Looking for the JAX backend? The original TPU backend for vLLM, tpu-inference https://github.com/vllm-project/tpu-inference , has been rebranded as vLLM JAX . It remains available for serving pure JAX model implementations on TPU. vllm-torchtpu is the recommended path for serving PyTorch model definitions on TPU. | TPU generation | Status | |---|---| | TPU v7x Ironwood | Recommended — validated in CI on every pull request and nightly | | TPU v6e / v5e / v5p | Experimental — not covered by CI | See Cloud TPU Setup https://vllm-torchtpu.readthedocs.io/en/latest/getting started/tpu setup/ for project setup, quota, and provisioning commands. vllm-torchtpu requires Python 3.12 on a Cloud TPU VM. Install the plugin and a matching vLLM with a single command: pip install \ --extra-index-url https://download.pytorch.org/whl/cpu \ vllm-torchtpu For uv , editable source installs, cross-repo development against a local vllm checkout, and Docker images, see the full Installation guide https://vllm-torchtpu.readthedocs.io/en/latest/getting started/installation/ and Docker guide https://vllm-torchtpu.readthedocs.io/en/latest/getting started/docker/ . Start an OpenAI-compatible server exactly as you would on any other vLLM platform: vllm serve meta-llama/Llama-3.2-1B-Instruct \ --tensor-parallel-size 1 \ --max-model-len 4096 Then send a request: A stage may instead span a prefill context-parallel group --prefill-context-parallel-size , which splits each chunk's attention over the stage's cores; MoE models then need --enable-expert-parallel . This needs a torch tpu that compiles a stage's kernels and collectives for the stage's devices only the 2026-09-28 build of google-pytorch/torch tpu 4387 https://github.com/google-pytorch/torch tpu/pull/4387 ; on a torch tpu that builds world-size meshes only, the platform refuses such stages. A two-core stage must not be the two cores of one chip: set VLLM TPU DEVICE ORDER to the device each local rank binds, e.g. 0,2,1,3,4,6,5,7 for four stages of two cores on one host. curl http://localhost:8000/v1/chat/completions \ -H "Content-Type: application/json" \ -d '{ "model": "meta-llama/Llama-3.2-1B-Instruct", "messages": {"role": "user", "content": "Explain how TPUs accelerate matrix multiplication."} , "max tokens": 128 }' Tip The first startup compiles TPU graphs for the configured shapes and can take several minutes. Set VLLM XLA CACHE PATH to persist compiled artifacts across restarts — see Compilation & Warmup Cache https://vllm-torchtpu.readthedocs.io/en/latest/user guide/compilation cache/ . The Quickstart https://vllm-torchtpu.readthedocs.io/en/latest/getting started/quickstart/ covers offline batch inference with vllm.LLM , multi-device serving with TP/EP, and multimodal examples. vllm-torchtpu can serve most PyTorch models supported by vLLM, with first-class TPU kernel, quantization, and parallelism support for Qwen3 / Qwen3.5, DeepSeek-V3 / R1 / V4, Kimi-K2.5 / K3, GLM-5, Llama 3 / 4, GPT-OSS, Gemma 3 / 4, Mixtral, Qwen-VL and Gemma VL families, and diffusion models via vllm-omni . - Supported Models & Features https://vllm-torchtpu.readthedocs.io/en/latest/user guide/supported models/ — validated model families, feature status, and quantization formats. - Nightly support matrices https://github.com/vllm-project/vllm-torchtpu/blob/main/support matrices/nightly — TPU v7x unit, correctness, and benchmark status, updated automatically from nightly CI. We welcome contributions of all kinds — bug reports, model and feature requests, documentation fixes, and code. - 🐞 Report bugs or request features: open a new issue https://github.com/vllm-project/vllm-torchtpu/issues/new/choose . - 👀 Review pull requests: help us keep quality high by reviewing open PRs https://github.com/vllm-project/vllm-torchtpu/pulls . - 💻 Write code: start with the Contributing Guide https://github.com/vllm-project/vllm-torchtpu/blob/main/CONTRIBUTING.md , which covers environment setup, DCO sign-off, the ready -label CI gate, and the PR lifecycle. Good first issues https://github.com/vllm-project/vllm-torchtpu/issues?q=is%3Aissue+is%3Aopen+label%3A%22good+first+issue%22 are a great place to begin. Developer references: Testing & CI https://vllm-torchtpu.readthedocs.io/en/latest/developers guide/testing and ci/ · Profiling https://vllm-torchtpu.readthedocs.io/en/latest/developers guide/profiling/ · Benchmarking https://vllm-torchtpu.readthedocs.io/en/latest/developers guide/benchmarking/ · Coding guidance https://vllm-torchtpu.readthedocs.io/en/latest/developers guide/coding guidance/ · Optimization best practices https://vllm-torchtpu.readthedocs.io/en/latest/developers guide/optimization best practices/ - For technical questions and feature requests, open a GitHub issue https://github.com/vllm-project/vllm-torchtpu/issues . - For discussion with fellow users, use the TPU support topic on the vLLM Forum https://discuss.vllm.ai/c/hardware-support/google-tpu-support/27 . - For coordinating contributions and development, join the vLLM Developer Slack https://slack.vllm.ai sig-tpu .