vllm-torchtpu README

The project's README from github.com/vllm-project/vllm-torchtpu. As one markdown file: README.md.

vLLM TorchTPU vLLM TorchTPU

| Documentation | Blog | User Forum | Developer Slack (#sig-tpu) |

About

vLLM TorchTPU is the PyTorch-native TPU backend for 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.

[!NOTE] Looking for the JAX backend? The original TPU backend for vLLM, 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.


Supported hardware

TPU generation Status
TPU v7x (Ironwood) Recommended — validated in CI on every pull request and nightly
TPU v6e (Trillium) Supported — not covered by CI
TPU v5e / v5p Experimental — not covered by CI

See Cloud TPU Setup for project setup, quota, and provisioning commands.


Get started

Install

vllm-torchtpu requires Python 3.12 on a Cloud TPU VM. Install vllm-torchtpu from PyPI and a matching vLLM with a single command:

pip install \
  --extra-index-url https://download.pytorch.org/whl/cpu \
  vllm-torchtpu

Or pull the pre-built container from Docker Hub (vllm/vllm-torchtpu):

docker pull vllm/vllm-torchtpu:latest

For uv, editable source installs, cross-repo development against a local vllm checkout, and Docker images, see the full Installation guide and Docker guide.

Serve a model

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:

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.

The Quickstart covers offline batch inference with vllm.LLM, multi-device serving with TP/EP, and multimodal examples.


Supported models and features

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.


Recipes and deployment


🤝 Contribute

We welcome contributions of all kinds — bug reports, model and feature requests, documentation fixes, and code.

Developer references: Testing & CI · Profiling · Benchmarking · Coding guidance · Optimization best practices


💬 Contact us

vllm-torchtpu (vllm-project/vllm-torchtpu)

Searching Database…

114 ·GitHub
Vllm-torchtpu
What is this?vllm-torchtpu is a PyTorch-native TPU backend plugin for vLLM, letting the same vLLM CLI, OpenAI-compatible server, and model definitions run on Google Cloud TPUs via torch.compile and custom JAX/Pallas kernels, with no model code changes.
Is this a separate app or a plugin for vLLM?It is an out-of-tree vLLM platform plugin, distributed as the vllm-torchtpu package, that you install alongside vLLM rather than a standalone application.
How do I install and run it?Install vllm-torchtpu from PyPI together with a matching vLLM version using a single pip command, or pull the pre-built vllm/vllm-torchtpu image from Docker Hub, then start vLLM's normal OpenAI-compatible server.
What do I need before using this?You need Python 3.12 running on a Google Cloud TPU VM; no GPU is required, but you must set up TPU project access, quota, and provisioning first.
Any catches when running this?The first server startup compiles TPU graphs for the configured shapes, which can take several minutes (you can persist the compiled cache via VLLMXLACACHEPATH), and only TPU v7x is validated in CI while v6e has no CI coverage and v5e/v5p remain experimental.
How does this differ from the project's JAX backend (vLLM JAX)?The original TPU backend, formerly tpu-inference and now rebranded vLLM JAX, serves pure JAX model implementations, while vllm-torchtpu is the recommended path for serving vLLM's native PyTorch model definitions on TPU.
Which models and features are supported?It can serve most PyTorch models vLLM already supports, with first-class TPU kernels and quantization for Qwen3/3.5, DeepSeek-V3/R1/V4, Kimi-K2.5/K3, GLM-5, Llama 3/4, GPT-OSS, Gemma 3/4, Mixtral, Qwen-VL/Gemma-VL, plus diffusion models via vllm-omni.
What custom kernels does it bring for TPU performance?It ships TPU-tuned Pallas kernels including Ragged Paged Attention, dense and sparse MLA, Gated DeltaNet, fused MoE with SparseCore offload, and FP8/INT8/MXFP4 quantized matmuls.
Does it support scaling across multiple TPU hosts?Yes, tensor, data, pipeline, and expert parallelism plus prefill/decode context parallelism compose across single-host and multi-host TPU slices, including disaggregated prefill/decode serving.
What are similar projects?public-apis, free-programming-books, system-design-primer, developer-roadmap, awesome-python, project-based-learning, hermes-agent, Python
Which topics is it in?python
Creator
vllm-project
Size
131.3 MB
Stars
114
Updated
CommitOct 27
Age
4mo
Commits · 90 Days