vLLM TorchTPU brings vLLM model serving to Google Cloud TPUs

A plugin built on Google's TorchTPU stack

vllm-torchtpu is built on Google's TorchTPU stack and comes out of the engineering partnership between Inferact and Google, which the vLLM team says it announced last month. It is an out-of-tree platform plugin: once installed, vLLM detects it automatically through the vllm.platform_plugins entry point.

Models run through vLLM's standard model runner and are compiled with torch.compile(backend="tpu"), which lowers to XLA. Each TPU device gets its own worker process, as on GPUs, and devices communicate through torch.distributed. According to the vLLM team, the scheduler, workers, KV cache manager and API server remain unmodified upstream vLLM. The TPU-specific work sits underneath, mostly in JAX/Pallas kernels registered as custom ops.

Installation is done with pip install vllm-torchtpu. It requires Python 3.12 on a Cloud TPU VM and installs a matching vLLM, torch, torch_tpu and libtpu. The first startup compiles TPU graphs for the configured shapes, which can take a few minutes. Setting VLLM_XLA_CACHE_PATH lets users reuse compiled artifacts across restarts.

Kernels, parallelism and serving features

The project uses custom Pallas kernels for operations that benefit from hand-tuning on TPU. Examples listed in the announcement:

• Ragged Paged Attention v3 for MHA/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

All six of vLLM's parallelism dimensions are supported on single-host and multi-host TPU slices, and they can be combined: tensor, data, pipeline, expert, prefill context and decode context parallelism. For large MoE models, the team says it usually gets the best throughput from data-parallel attention with expert parallelism, for example DP=8, TP=1 with EP on a single v7x host.

The team says it focused on multi-turn agentic and multimodal traffic, where contexts keep growing and tool outputs can be long. For this, vllm-torchtpu uses TPU Sync, an open-source C++ library that reads TPU device buffers directly, outside the Python runtime. Through vLLM's KV connector interface, it enables zero-copy KV cache offloading to host DRAM and disaggregated prefill/decode serving across TPU hosts. Prefill and decode can use different layouts, and the KV cache is resharded during transfer.

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 along with the KV cache.

Supported hardware and models

vllm-torchtpu can serve most models supported by vLLM. It includes 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, while v6e, v5e and v5p are described as experimental. A Supported Models & Features page lists validated checkpoints, and nightly support matrices track v7x correctness and benchmark results.

The announcement's example serves Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8 with tensor parallel size 8, expert parallelism, an fp8 KV cache and async scheduling, using the usual vllm serve command.

Roadmap and what it means for users

The team wants new open models from the major families to run on TPU on the day they are released. Items on the public roadmap include 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, llm-d TPU memory-tier routing, PyTorch interfaces for Tokamax kernels, and diffusion model serving through vllm-omni.

The project is also 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 pull request reviews.

For teams that already run vLLM, the practical point is continuity: according to the project, the CLI, OpenAI-compatible server, scheduler and KV cache manager are the same as on other hardware. That should make TPUs a more direct option for serving open models, with Ironwood as the best-supported target for now.