vLLM TorchTPU: a new vLLM backend for Google Cloud TPU inference

vLLM TorchTPU is a new out-of-tree vLLM platform plugin that lets vLLM serve models on Google Cloud TPUs while reusing vLLM's existing model code, CLI, scheduler, KV cache manager, and OpenAI-compatible server unchanged. It matters because it means most models that already work in vLLM can run on TPU without edits to model definitions, and because the project puts hand-written Pallas kernels for attention, Mixture-of-Experts, and quantized matmul into open source.

What vLLM TorchTPU is

vLLM TorchTPU is the package vllm-torchtpu, a backend for running vLLM on Google Cloud TPUs. It is built on Google's TorchTPU stack, which runs PyTorch natively on TPUs, and it executes the model definitions that already live in vLLM. As a result, most models that work in vLLM today can be served on TPU without changes to the model code.

The project comes out of the engineering partnership between Inferact and Google that was announced the month before the release. The release post states three goals for the project:

  1. Make it easy to bring new models to TPU by reusing vLLM's existing model code.
  2. Push TPU inference performance in open source, with Pallas kernels for attention, Mixture-of-Experts, and quantized matmul.
  3. Keep the vLLM experience the same: the same CLI, OpenAI-compatible server, scheduler, and KV cache manager used on other hardware.

How the TPU backend plugs into vLLM

vllm-torchtpu is an out-of-tree vLLM platform plugin. Once installed, vLLM picks it up automatically through the vllm.platform_plugins entry point, so no forked or patched copy of vLLM is required.

Models run through vLLM's standard model runner and are compiled with torch.compile(backend="tpu"), which lowers the graph to XLA. Each TPU device gets its own worker process, the same way vLLM runs on GPUs, and communication between devices goes through torch.distributed.

The practical consequence: the scheduler, workers, KV cache manager, and API server are all unmodified upstream vLLM. The TPU-specific work sits underneath those components, mostly in a set of JAX/Pallas kernels that are registered as custom ops and called from inside the compiled graph.

Getting started: install and serve a model

Installation is a single pip command on a Cloud TPU VM running Python 3.12:

pip install vllm-torchtpu

The package installs a matching vLLM, torch, torch_tpu, and libtpu alongside the plugin. The installation guide also covers uv, source builds, and Docker.

A representative serving command from the announcement, using an FP8 Qwen3-Coder MoE model across eight TPU chips with expert parallelism and FP8 KV cache:

vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8 \
  --tensor-parallel-size 8 \
  --enable-expert-parallel \
  --kv-cache-dtype fp8 \
  --max-model-len 16384 \
  --async-scheduling

The first startup compiles TPU graphs for the configured shapes, which can take a few minutes. Setting VLLM_XLA_CACHE_PATH reuses compiled artifacts across restarts, avoiding a recompilation on every boot.

Speculative decoding, prefix caching, multi-LoRA, structured outputs, and multimodal inputs are all enabled with the usual vLLM flags. The quickstart covers additional examples, including offline batch inference, multimodal models, and multi-LoRA.

Pallas kernels: what is hand-tuned for TPU

vLLM TorchTPU uses custom Pallas kernels for operations that benefit from hand-tuning on TPU rather than relying entirely on the compiler. The announced kernel set includes:

  • Ragged Paged Attention v3 for MHA and GQA models.
  • Dense and sparse Multi-head Latent Attention (MLA) for DeepSeek, Kimi, and GLM.
  • Gated DeltaNet and Kimi Delta Attention for hybrid linear-attention models.
  • Fused Mixture-of-Experts with expert parallelism.
  • FP8, INT8, MXFP4, and NVFP4 quantized matmuls.

Scaling out: six parallelism dimensions

All six of vLLM's parallelism dimensions are supported, on both single-host and multi-host TPU slices, and they can be combined in one deployment.

Dimension Flag
Tensor Parallelism (TP) --tensor-parallel-size
Data Parallelism (DP) --data-parallel-size
Pipeline Parallelism (PP) --pipeline-parallel-size
Expert Parallelism (EP) --enable-expert-parallel
Prefill Context Parallelism (PCP) --prefill-context-parallel-size
Decode Context Parallelism (DCP) --decode-context-parallel-size

For large Mixture-of-Experts models, the project reports that the best throughput usually comes from data-parallel attention combined with expert parallelism, for example DP=8, TP=1 with EP enabled on a single v7x host. That configuration keeps attention local to each device while spreading the experts across the whole slice. For long contexts, PCP and DCP split prefill and the KV cache across devices.

The stated focus is on the workloads that put the most pressure on serving systems today, particularly multi-turn agentic and multimodal traffic, where contexts keep growing and tool outputs can be long.

KV cache offloading and disaggregated serving

vllm-torchtpu builds its KV movement features on TPU Sync, an open-source C++ library that can read TPU device buffers directly, outside the Python runtime. Wired through vLLM's KV connector interface, this yields two capabilities:

  • Zero-copy KV cache offloading to host DRAM, so the prefix cache for long multi-turn sessions is not limited by HBM capacity.
  • Disaggregated prefill/decode serving across TPU hosts with low KV transfer overhead. Prefill and decode can use different layouts — for example prefill on PCP=4, TP=2 and decode on DP=8, TP=1 — and the KV cache is resharded during the transfer.

Speculative decoding and prefix caching

Speculative decoding supports MTP, EAGLE-3, DFlash, DSpark, n-gram, and suffix decoding. Automatic prefix caching also works for hybrid models such as Qwen3.5 and Kimi K3, where linear-attention state is cached alongside the KV cache.

Supported models and hardware

vLLM TorchTPU can serve most models supported by vLLM, with TPU-specific kernel, quantization, and parallelism work for the Qwen, DeepSeek, Kimi, GLM, Llama, GPT-OSS, and Gemma families.

TPU v7x (Ironwood) is the recommended hardware and is tested in CI on every pull request. TPU v6e, v5e, and v5p are described as experimental. A nightly support matrix tracks TPU v7x correctness and benchmark results from CI, and a Supported Models & Features page lists validated checkpoints and how to enable each feature.

Roadmap: what's next

The stated ambition is that new open models from the major families — Qwen, DeepSeek, GLM, Kimi, Gemma, and GPT-OSS — should run on TPU on the day they are released. The public roadmap lists:

  • An automated CI pipeline that onboards, benchmarks, and publishes support matrices for new open-weight models within 24 hours of release.
  • Prefix-, cache-, and session-aware routing, plus llm-d TPU memory-tier routing.
  • PyTorch interfaces for Tokamax kernels, shared with other projects in the TorchTPU ecosystem.
  • Diffusion model serving through vllm-omni.

Contributions are explicitly welcomed through the project's contributing guide.

Contributor TPU capacity pool

Alongside the launch, the project is opening a Community Contribution TPU Capacity Pool. Developers building on vllm-torchtpu can apply for free TPU time, and core vLLM maintainers at Inferact will help with design and PR reviews. Applications go through a public Google Form linked from the announcement.

Where to find the code and documentation

  • Quickstart, Cloud TPU setup, and installation (pip, uv, source, Docker) documentation live in the vllm-torchtpu repository.
  • Deployment resources include TPU v7x (Ironwood) serving recipes, vLLM TorchTPU recipes on Ironwood in the tpu-recipes repository, and a Kubernetes deployment guide.
  • Discussion happens in the #sig-tpu channel on the vLLM Developer Slack.

Sources