inference/research
verifiedhigh · 0.85updated 2026-10-04

vLLM TPU inference plugin (tpu-inference)

tpu

Fields

id
engines/alt-tpu-inference-vllm-plugin
Repo
vllm-project/tpu-inference
Languages
Python
License
Apache-2.0
First release
2025
Design goal
tpu-only-servingunify-jax-and-pytorchserve-vLLM-semantics-on-XLA
Scheduling
continuous-batchingpaged-kvprefix-cachechunked-prefill
Backends
tpu
Formats
safetensors
Notable features
THE ANSWER TO 'how do you actually serve on TPU'. Google's own JetStream server was retired into this plugin (see [[alt-jetstream]]), so there is one supported TPU serving path and it lives inside vLLM rather than beside itOne lowering path for two frameworks. The README's stated contribution is 'an expressive and powerful new hardware plugin unifying JAX and PyTorch under a single lowering path', so a PyTorch model definition runs on TPU with no code changes, while native JAX support is retained. Announced 2025/10Publishes an honest support matrix instead of claiming parity. There is a public dashboard of per-model Unit Test / Correctness Test / Performance Test status, auto-generated from the in-repo support_matrices directory and stamped with a last-updated time (2026-08-27 when read). Rolling performance-tested status up across hardware generations is published with explicit rules: a feature is Failing if it fails on EITHER v6e or v7x, and Untested if either generation is untested - so the matrix never flatters itself by testing only the newest chipThe same honesty about incompleteness: the README states that while the unified backend makes out-of-the-box high performance possible for any vLLM-supported model, 'the reality is that we're still in the process of implementing a few core components', and points to Recommended Models and Recommended Features pages as the validated set. That is a vendor-family project (Google is funding vLLM here) choosing not to overstate coverageExplicit generation split that shows the constraint is generational, not generational-universal: recommended v7x (Ironwood), v5e, v6e (Trillium); experimental v3, v4, v5p. Older and larger parts (v4, v5p) are experimental while the smaller inference parts are recommended, which is the opposite of the datacenter-part patternRecipe-driven deployment: published v7x (Ironwood) and v6e (Trillium) recipes under AI-Hypercomputer/tpu-recipes/inference, so the supported path is a pinned recipe rather than free assemblyThe XLA cost is moved, not removed. Serving on TPU means the graph is compiled ahead of the request, which is why the repo's TPU throughput numbers elsewhere in this database ([[alt-tpu-v5e-bf16-tps]], [[alt-tpu-v6e-bf16-tps]]) carry different prefill/decode economics than a GPU engine. The plugin's own reporting (unit/correctness/performance test columns) is how it discloses that cost rather than hiding it453 stars, 329 forks, 531 open issues, Apache-2.0, created 2025-05-06, last push 2026-10-04. The open-issue count exceeds the star count, which for a hardware-first plugin means the support burden is carried by a small user base with a large upstream to track
Hardware caveats
TPU-only. No CUDA, ROCm, Metal or CPU fallback exists in this plugin, and unlike a CUDA backend it cannot even be compiled for another target - the lowering path goes through XLATPU-only means Cloud TPU or a TPU appliance with the matching generation. The generation split above is not cosmetic: experimental on v3/v4/v5p means you are on an unsupported path thereThe XLA compilation model is the dominant cost. Static shapes are mandatory, tensors are padded to the MXU tile size, and every new shape class triggers recompilation - see [[alt-tpu-xla-static-shapes]] for the tile arithmetic. On a TPU this is not a tuning detail, it is the price of the hardware and it shows up as warm-up latency and as wasted compute on paddingModel coverage is explicitly partial and the project says so, unlike most plugin claims. Check the support matrix for your exact checkpoint before planning around it531 open issues is a large queue for a plugin with 453 stars. Some is upstream-tracking friction, some is genuine TPU-specific work, and it is a fair signal of how much of the XLA model surface still needs hardware-specific handlingVery young: created 2025-05-06, so it is absorbing the serving responsibilities of a system (JetStream) that itself only existed from 2024. The TPU serving path in open source is roughly two years old and has changed owner once
Best for
Serving on TPU v5e, v6e or v7x with real vLLM semantics - paged KV, continuous batching, prefix caching - which is what JetStream never quite hadTeams with PyTorch model code that need to move to TPU without a JAX rewrite, which is the plugin's headline claimAnyone deciding whether TPU serving is ready, because the support matrix with per-model correctness AND performance test status is the best public evidence available on that questionReading a vendor-funded hardware plugin that publishes its own failures - the status legend explicitly includes Failing, Unplanned and Untested
Avoid for
Any non-TPU hardwarev3, v4 or v5p deployments without accepting that you are on the experimental pathAssuming 'supported by vLLM' means 'performance-tested on TPU' - the plugin goes out of its way to separate those two columns, and so should youAny workload needing a model or feature absent from the support matrix, since the project itself says the core components are still being implemented
notes
This is the serving-path record for TPU; the compiler-side records are [[xla]], [[pallas]] and [[iree]] and the constraint mechanics are in [[alt-tpu-xla-static-shapes]]. Google's own answer to 'does a TPU-native server exist outside JAX' is now a qualified no: JetStream was retired into this plugin, which does support PyTorch natively but is still XLA underneath. Cloud TPU's own inference guidance is in [[alt-tpu-inference]] and its capacity story in [[supply/gcp-cloud-tpu-v4-v5e-v5p]]. Metadata from the GitHub API on 2026-10-04. ||| MAINTENANCE (dated fact, GitHub API, 2026-10-04): last commit 2026-10-04 (0.0 months before this record), latest release v0.30.0 published 2026-09-30, 160 contributors (distinct committers including anonymous), open_issues_count 532 (issues PLUS pull requests - see [[ev-github-open-issues-count-includes-pull-requests]]), not archived, 453 stars. See [[ev-maintenance-status-is-four-dated-facts]] for how to read this fact set.

Notes

This is the serving-path record for TPU; the compiler-side records are [[xla]], [[pallas]] and [[iree]] and the constraint mechanics are in [[alt-tpu-xla-static-shapes]]. Google's own answer to 'does a TPU-native server exist outside JAX' is now a qualified no: JetStream was retired into this plugin, which does support PyTorch natively but is still XLA underneath. Cloud TPU's own inference guidance is in [[alt-tpu-inference]] and its capacity story in [[supply/gcp-cloud-tpu-v4-v5e-v5p]]. Metadata from the GitHub API on 2026-10-04. ||| MAINTENANCE (dated fact, GitHub API, 2026-10-04): last commit 2026-10-04 (0.0 months before this record), latest release v0.30.0 published 2026-09-30, 160 contributors (distinct committers including anonymous), open_issues_count 532 (issues PLUS pull requests - see [[ev-github-open-issues-count-includes-pull-requests]]), not archived, 453 stars. See [[ev-maintenance-status-is-four-dated-facts]] for how to read this fact set.

Cross-references

This record references no other records — it is a leaf. It is cited by 0.