vllm-project/tpu-inference

TPU inference for vLLM, with unified JAX and PyTorch support.

What it is

vLLM TPU (powered by the tpu‑inference plugin) is an open‑source backend that lets you run large language models on Google TPUs. It unifies the JAX and PyTorch ecosystems under a single lowering path, so the same vLLM code you use on GPUs can run on TPUs without changes.

Core capabilities

  • Unified JAX + PyTorch support – write your model once (in either framework) and execute it on TPU.
  • vLLM‑compatible user experience – same CLI, telemetry, and API as the regular vLLM server.
  • Performance features – async scheduler, chunked prefill, KV‑cache offload, LoRA support, multimodal inputs, speculative decoding (Eagle3, N‑gram, DFlash), prefix caching, single‑program‑multiple‑data, and more (all marked Passing in the support matrix).
  • Parallelism – pipeline (PP), data (DP), expert (EP), tensor (TP) parallelism work on single‑host; multi‑host variants are still experimental.
  • Quantization – early support for FP4/FP8 compressed tensors on v7 TPUs (currently untested).

Who should use it

  • Researchers and engineers who want to serve or fine‑tune state‑of‑the‑art LLMs on TPU hardware.
  • Teams already using vLLM on GPUs and looking to extend to TPUs with minimal code changes.
  • Anyone interested in pushing the performance limits of TPU inference in open source.

Getting started

  1. Follow the quick‑start guide in the documentation: https://docs.vllm.ai/projects/tpu/en/latest/getting_started/quickstart/.
  2. Install the tpu‑inference plugin (instructions are in the docs).
  3. Choose a supported TPU generation – recommended v7x, v5e, v6e; older v3/v4/v5p are experimental.
  4. Use one of the recommended models (see the linked page) to verify your setup.

Supported hardware & models

  • TPU generations: v7x (Ironwood) and v6e (Trillium) are the primary targets; v5e also recommended.
  • Tested models (all passing unit, correctness, and performance tests) include:
    • Google Gemma‑4 (various sizes)
    • Google Gemma‑3‑27B‑it
    • Meta Llama‑3.1‑8B‑Instruct, Llama‑3.3‑70B‑Instruct
    • Qwen 3‑30B‑A3B, Qwen 3‑32B, Qwen 3‑4B, Qwen 3‑Coder‑480B‑A35B‑Instruct, Qwen 3.5‑397B‑A17B, etc.
  • Experimental or untested models are listed in the support matrix with ❓ or ❌ markers.

Community & contribution

  • Issue tracker shows 12 open bugs, 8 good‑first‑issues, and other labels to help newcomers.
  • Join the vLLM Slack (#sig-tpu) or the user forum for help.
  • Contribute via the GitHub repo – many enhancement and auto‑generated issues are labeled for easy entry.

Additional resources

Related

  • Project
    vllm-project/vllm-metal

    vLLM Metal is a plugin that enables the vLLM LLM serving system to run on Apple Silicon Macs using the MLX/Metal backend. It provides Metal‑optimized attention kernels (paged varlen, M5 NAX prefill) and supports large models such as Qwen‑3.8‑27B. Install via Homebrew or a curl script, then use the regular `vllm` command on macOS 15+.

  • Project
    sgl-project/sglang-jax

    A high-performance JAX-based inference engine optimized for Google TPUs, providing high throughput and low latency for LLMs and multimodal models.

  • Project
    baidu/vLLM-Kunlun

    A hardware plugin that enables vLLM to run on Kunlun XPU, providing optimized kernels and support for a wide range of LLMs and multimodal models.

  • Project
    vllm-project/vllm-ascend

    vLLM Ascend is a hardware plugin that enables the vLLM LLM serving engine to run on Huawei Ascend NPUs. It supports a range of Atlas NPU hardware, requires Linux, Python 3.10‑3.12, CANN 9.1.0, PyTorch 2.10.0 + TorchNPU, and matches specific vLLM versions. The project provides regular releases, detailed docs, and an active community (Slack, forum, weekly meetings). It’s licensed under Apache 2.0.

  • Project
    xLLM-AI/xllm

    An efficient LLM inference framework optimized for Chinese AI accelerators to enable high-throughput, low-latency enterprise deployment.