Today we’re releasing vLLM TorchTPU (vllm-torchtpu), a new backend for running vLLM on Google Cloud TPUs. It’s built on Google’s TorchTPU stack and runs the model definitions that already live in vLLM, so most models that work in vLLM today can be served on TPU without changes to the model code.

This work comes out of the engineering partnership between Inferact and Google we announced last month. Our goals for the project are to:

  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 you use on other hardware.

How It Works

vllm-torchtpu is an out-of-tree vLLM platform plugin. Once installed, vLLM picks it up 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, the same way vLLM runs on GPUs, and communication between devices goes through torch.distributed. This means the scheduler, workers, KV cache manager, and API server are all unmodified upstream vLLM. The TPU-specific work sits underneath them, mostly in a set of JAX/Pallas kernels that are registered as custom ops and called from inside the compiled graph.

Getting Started

  1. Install

     pip install vllm-torchtpu

    This requires Python 3.12 on a Cloud TPU VM, and installs a matching vLLM, torch, torch_tpu, and libtpu alongside the plugin. The installation guide also covers uv, source builds, and Docker.

  2. Serve a model

     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. Set VLLM_XLA_CACHE_PATH to reuse compiled artifacts across restarts (see Compilation & Warmup Cache).

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

Kernels

vllm-torchtpu uses custom Pallas kernels for complex operations that benefit from hand-tuning on TPU. Some examples:

  • 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

Scaling Out

We’ve focused heavily on the workloads that put the most pressure on serving systems today, especially multi-turn agentic and multimodal traffic, where contexts keep growing and tool outputs can be long.

Parallelism

All six of vLLM’s parallelism dimensions are supported, on both single-host and multi-host TPU slices, and they can be combined:

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 MoE models, we usually get the best throughput from data-parallel attention with expert parallelism (for example, DP=8, TP=1 with EP on a single v7x host). This keeps attention local to each device and spreads the experts across the whole slice. For long contexts, PCP and DCP split prefill and the KV cache across devices. The parallelism guide has more detail and example configurations.

KV Cache Offloading and Disaggregated Serving

vllm-torchtpu uses TPU Sync, an open-source C++ library that can read TPU device buffers directly, outside the Python runtime. Through vLLM’s KV connector interface, this gives us:

  • Zero-copy KV cache offloading to host DRAM, so the prefix cache for long multi-turn sessions isn’t limited by HBM.
  • 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 transfer.

See the disaggregated serving guide to set it up.

Speculative Decoding and Prefix Caching

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

Models, Features, and What’s Next

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; v6e, v5e, and v5p are experimental.

The Supported Models & Features page has the full list of validated checkpoints and features, with how to enable each one. The nightly support matrices track TPU v7x correctness and benchmark results from CI.

What’s Next?

We want new open models from the major families (Qwen, DeepSeek, GLM, Kimi, Gemma, GPT-OSS) to run on TPU on the day they’re released. Some of what’s coming, from the public roadmap:

  • 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, and 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 welcome!

TPU Capacity for Contributors

We’re also opening a Community Contribution TPU Capacity Pool. If you’re building on vllm-torchtpu, you can apply for free TPU time, and core vLLM maintainers at Inferact will help with design and PR reviews.

Apply for contributor TPU access here.

Try It Out

If you have questions or want to get involved, find us in #sig-tpu on the vLLM Developer Slack.