Google 宣布在 TPU 上原生支援 vLLM 與 SGLang 推理引擎

Google 宣布在 PyTorch 公告中推出 vLLM 和 SGLang 的原生 PyTorch 後端。TPU 過去已透過基於 JAX 的後端運行 vLLM 和 SGLang，但開發者希望使用 PyTorch 自訂和除錯其服務堆疊。透過 TorchTPU，TPU 成為原生 PyTorch 裝置，張量在 tpu 裝置上作為標準 PyTorch 張量管理，無包裝類型。

此次推出的原生後端由 Google 與合作夥伴共同開發。vLLM 由 Google 與 Inferact 合作開發。SGLang 由 Google 與 RadixArk 合作開發。在架構上，每個裝置運行自己的 PyTorch 進程，引擎保留各自框架的工作者模型、排程器、KV 快取管理器和 API 伺服器。大多數模型使用 vLLM 或 SGLang 中的標準上游 PyTorch 模型定義，TPU 特定的優化在內核和層實現中，使用 Pallas 編寫的自訂內核作為編譯圖內的自訂運算子直接從 PyTorch 呼叫。

在支援功能方面，引擎支援自動前綴快取，包括混合模型的前綴快取。引擎支援 KV 快取卸載至主機記憶體和透過 TPU Sync 的 Managed Lustre。引擎支援工具呼叫和結構化輸出。引擎支援使用 MTP、EAGLE-3、DFlash、DSpark 和 n-gram 草稿生成器的推測解碼。引擎支援具有專家並行性的資料平行注意力。引擎支援預填和解碼的上下文並行性、FP8 KV 快取和主機 KV 卸載。引擎支援在 TPU Sync 上於不同主機上進行預填和解碼，包括不同平行佈局之間。引擎支援具有影像和影片輸入的視覺語言模型和嵌入模型。引擎支援混合和稀疏注意力架構，以及在主機間進行張量、資料、專家、管線和上下文並行性，支援 FP8、FP4 (MXFP4, NVFP4)、INT8 和 INT4 檢查點。

關於快取管理，引擎將可重複使用的前綴保留在 HBM 中，並在開啟主機卸載時自動將較舊的前綴移至主機記憶體作為第二層。TPU Sync 是一個用於 TPU 上 KV 快取傳輸和管理的開源 C++ 函式庫。

Google 表示，隨著模型變大，Google 正在演進這兩個引擎，以使大規模 TPU 推理部署具有效能和效率。生產級代理工作負載的主要功能已獲得支援。更多效能最佳化和進階功能正在進行中。Google 鼓勵開啟 Github 問題以請求在特定需求上進行合作。在兩個引擎中，大多數受歡迎的開放模型系列已獲得支援，包括 Qwen、DeepSeek、GLM、Kimi、GPT-OSS、Gemma4 等。Google 希望給予社區工具和 TPU 容量，以調整其自身工作負載和配置的效能。

此外，vLLM 正在啟動一個首創的社區貢獻 TPU 容量池，以向積極為兩個引擎做出貢獻的開發者提供免費的 TPU 運算資源。vLLM 提供申請入口。相關公告包括 Inferact 公告和 Radixark 公告。

儘管推出原生 PyTorch 後端，vLLM-Jax 和 SGLang-JAX 仍然是重要的專案。vLLM-Jax 和 SGLang-JAX 將繼續獲得維護、支援，並供社區使用和調整。

開發者可透過以下 GitHub 連結取得相關專案： vLLM-TorchTPU 的 GitHub 連結為 SGLang-TorchTPU 的 GitHub 連結為 TorchTPU 的 GitHub 連結為