vLLM TorchTPU 發布 PyTorch 原生 TPU 後端

| 檔案 | 部落格 | 使用者論壇 | 開發者 Slack (#sig-tpu) |

關於

vLLM TorchTPU 是 vLLM 的 PyTorch 原生 TPU 後端。它是一個離線 vLLM 平臺外掛，打包為 vllm-torchtpu，在 Google Cloud TPUs 上透過 TorchTPU (torch_tpu, via torch.compile(backend="tpu")) 執行 vLLM 的標準 PyTorch 模型定義，並提供自訂 JAX/Pallas 核心，用於注意力、Mixture-of-Experts 與量化矩陣乘法。

目標很簡單：相同的 vLLM，您已經在使用——相同的 CLI、相同的 OpenAI 相容伺服器、相同的模型定義——以 TPU 速度執行，且不需要更改模型程式碼。

PyTorch 首先。 模型透過 vLLM 原生 PyTorch 執行器執行；不需要在其他框架重新實作即可在 TPU 上提供服務。
TPU 調校 kernel。 Ragged Paged Attention、密集與稀疏 MLA、Gated DeltaNet、融合 MoE（包含 SparseCore offload）以及 FP8/INT8/MXFP4 量化矩陣乘法，均以 Pallas 編寫。
內建 Scale-out。 TP、DP、PP、EP 以及預填/解碼上下文並行性可跨單主機與多主機 TPU 切片組合，並支援分離的預填/解碼服務。
標準 vLLM 體驗。 透過常規 vLLM 旗標，支援推測解碼、字首快取、多 LoRA、結構化輸出與多模態輸入。

注意

尋找 JAX 後端嗎？ 原始 TPU 後端 for vLLM, tpu-inference, 已重新命名為 vLLM JAX。它仍可用於在 TPU 上提供純 JAX 模型實作。vllm-torchtpu 是在 TPU 上提供 PyTorch 模型定義的推薦路徑。

支援硬體

TPU 代數	狀態
TPU v7x (Ironwood)	推薦 — 在每一次 PR 與 nightly CI 中已驗證
TPU v6e / v5e / v5p	實驗性 — 未被 CI 覆蓋

請參閱 Cloud TPU Setup 以瞭解專案設定、配額與配置命令。

開始

安裝

vllm-torchtpu 需要在 Cloud TPU VM 上執行 Python 3.12。從 PyPI 安裝 vllm-torchtpu 並使用單一指令安裝相容的 vLLM：

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

或從 Docker Hub (vllm/vllm-torchtpu) 拉取預建容器：

docker pull vllm/vllm-torchtpu:latest

對於 uv，可編輯來源安裝、對本地 vllm 檢出進行跨倉庫開發，以及 Docker 映像，請參閱完整 Installation guide 與 Docker guide。

提供模型

以與任何其他 vLLM 平臺相同的方式啟動 OpenAI 相容伺服器：

vllm serve meta-llama/Llama-3.2-1B-Instruct \
  --tensor-parallel-size 1 \
  --max-model-len 4096

然後傳送請求：

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
  }'

提示

首次啟動時會為設定的形狀編譯 TPU 圖形，可能需要數分鐘。設定 VLLM_XLA_CACHE_PATH 以在重啟時持久化編譯產物 — 請參閱 Compilation & Warmup Cache。

Quickstart 涵蓋離線批次推論 vllm.LLM、TP/EP 的多裝置服務，以及多模態範例。

支援的模型與功能

vllm-torchtpu 可為 vLLM 支援的大多數 PyTorch 模型提供服務，並具備 Qwen3 / Qwen3.5、DeepSeek-V3 / R1 / V4、Kimi-K2.5 / K3、GLM-5、Llama 3 / 4、GPT-OSS、Gemma 3 / 4、Mixtral、Qwen-VL 與 Gemma VL 系列，以及透過 vllm-omni 的擴散模型的首屈一指 TPU kernel、量化與並行支援。

支援模型與功能 — 已驗證的模型系列、功能狀態與量化格式。
每日支援矩陣 — TPU v7x 單元、正確性與基準測試狀態，自動從每日 CI 更新。

範例與部署

🤝 參與貢獻

我們歡迎各類貢獻——錯誤回報、模型與功能請求、檔案修正以及程式碼。

🐞 報告錯誤或請求功能：開啟一個新問題。
👀 審核 Pull Request：協助我們保持高品質，審核未結束 PR。
💻 撰寫程式碼：先從 貢獻指引開始，裡面說明環境設定、DCO 簽署、ready 標籤 CI 門檻與 PR 生命週期。Good first issues 是開始的好地方。

開發者參考：測試 & CI · 分析 · 基準測試 · 編碼指引 · 最佳化實務

💬 聯絡我們

如有技術問題與功能請求，請開啟 GitHub 問題。
若想與其他使用者討論，請使用 vLLM Forum 上的 TPU 支援主題。
為協調貢獻與開發，請加入 vLLM Developer Slack (#sig-tpu)。