今天我们发布vLLM TorchTPU（vllm-torchtpu），这是一个用于在 Google Cloud TPU 上运行 vLLM 的新后端。它构建在 Google 的TorchTPU技术栈之上，运行的是已经存在于 vLLM 中的模型定义，因此目前能在 vLLM 中运行的大多数模型无需修改模型代码即可在 TPU 上提供服务。

这项工作源自Inferact 与 Google 的工程合作，我们上个月宣布了这项合作。我们对该项目的目标是：

  1. 通过复用 vLLM 现有的模型代码，让将新模型引入 TPU 变得简单。
  2. 在开源中推动 TPU 推理性能，为注意力、混合专家和量化矩阵乘法提供 Pallas 内核。
  3. 保持 vLLM 体验不变：与其他硬件上相同的 CLI、OpenAI 兼容服务器、调度器和 KV 缓存管理器。

工作原理

vllm-torchtpu是一个 out-of-tree 的 vLLM 平台插件。安装后，vLLM 会通过vllm.platform_plugins入口点自动加载它。

模型通过 vLLM 的标准模型运行器运行，并使用torch.compile(backend="tpu")进行编译，后者会降级到 XLA。每个 TPU 设备都有自己的工作进程，与 vLLM 在 GPU 上的运行方式相同，设备间通信通过torch.distributed进行。这意味着调度器、工作进程、KV 缓存管理器和 API 服务器都是未经修改的上游 vLLM。TPU 特定的工作位于它们之下，主要是一组 JAX/Pallas 内核，这些内核被注册为自定义算子，并从编译后的图中调用。

快速开始

  1. 安装

    pipinstallvllm-torchtpu

    这需要在 Cloud TPU VM 上使用 Python 3.12，并会安装匹配的 vLLM、torch、torch_tpu，以及libtpu与插件一起。安装指南还涵盖了uv、源码构建以及Docker。

  2. 提供模型服务

    vllmserveQwen/Qwen3-Coder-30B-A3B-Instruct-FP8\
--tensor-parallel-size8\
--enable-expert-parallel\
--kv-cache-dtypefp8\
--max-model-len16384\
--async-scheduling

    首次启动会为配置的形状编译 TPU 图，这可能需要几分钟。设置VLLM_XLA_CACHE_PATH以在重启之间复用编译产物（参见编译与预热缓存）。

投机解码、前缀缓存、多 LoRA、结构化输出和多模态输入都可以通过常用的 vLLM 标志启用。参见快速入门了解更多示例，包括离线批量推理、多模态模型和多 LoRA。

内核

vllm-torchtpu使用自定义 Pallas 内核来执行那些在 TPU 上受益于手工调优的复杂操作。一些例子：

  • 面向 MHA/GQA 模型的 Ragged Paged Attention v3
  • 面向 DeepSeek、Kimi 和 GLM 的稠密与稀疏多头潜在注意力（MLA）
  • 面向混合线性注意力模型的 Gated DeltaNet 和 Kimi Delta Attention
  • 带专家并行的融合混合专家
  • FP8、INT8、MXFP4 和 NVFP4 量化矩阵乘法

横向扩展

我们重点聚焦于当前对服务系统压力最大的工作负载，尤其是多轮智能体和多模态流量，这类场景中上下文持续增长，工具输出也可能很长。

并行

vLLM 的全部六个并行维度均受支持，单主机和多主机 TPU 切片上都可以，并且可以组合使用：

维度标志
张量并行（TP）--tensor-parallel-size
数据并行（DP）--data-parallel-size
流水线并行（PP）--pipeline-parallel-size
专家并行（EP）--enable-expert-parallel
预填充上下文并行（PCP）--prefill-context-parallel-size
解码上下文并行（DCP）--decode-context-parallel-size

对于大型 MoE 模型，我们通常通过数据并行注意力配合专家并行获得最佳吞吐（例如，DP=8, TP=1在单台 v7x 主机上启用 EP）。这样可以让注意力保持在每个设备本地，同时把专家分散到整个切片上。对于长上下文，PCP 和 DCP 会把预填充和 KV 缓存拆分到多个设备上。并行指南中有更多细节和示例配置。

KV 缓存卸载与分离式服务

vllm-torchtpu使用TPU Sync，这是一个开源 C++ 库，可以在 Python 运行时之外直接读取 TPU 设备缓冲区。通过 vLLM 的 KV 连接器接口，我们获得了：

  • 零拷贝 KV 缓存卸载到主机 DRAM，因此长多轮会话的前缀缓存不再受 HBM 限制。
  • 跨 TPU 主机的分离式预填充/解码服务，KV 传输开销很低。预填充和解码可以使用不同的布局，例如预填充在PCP=4, TP=2上、解码在DP=8, TP=1上，并且 KV 缓存在传输过程中会重新分片。

参见分离式服务指南进行配置。

投机解码与前缀缓存

投机解码支持 MTP、EAGLE-3、DFlash、DSpark、n-gram 和后缀解码（指南）。自动前缀缓存同样适用于 Qwen3.5 和 Kimi K3 等混合模型，其中线性注意力状态会与 KV 缓存一起被缓存。

模型、特性与下一步

vllm-torchtpu可服务 vLLM 支持的大多数模型，并针对 Qwen、DeepSeek、Kimi、GLM、Llama、GPT-OSS 和 Gemma 系列做了 TPU 专属 kernel、量化和并行化工作。推荐使用 TPU v7x（Ironwood）硬件，且每个 pull request 都会在 CI 中测试；v6e、v5e 和 v5p 为实验性支持。

该支持的模型与特性页面列出了完整的已验证 checkpoint 和特性清单，以及如何启用每一项。每夜支持矩阵跟踪 TPU v7x 的正确性和来自 CI 的基准测试结果。

下一步是什么？

我们希望各大系列（Qwen、DeepSeek、GLM、Kimi、Gemma、GPT-OSS）的新开源模型能在发布当天就在 TPU 上运行。来自公开路线图的部分即将推出的内容：

  • 一条自动化 CI 流水线，可在新的开源权重模型发布后 24 小时内完成接入、基准测试并发布支持矩阵。
  • 前缀感知、缓存感知和会话感知的路由，以及 llm-d 的 TPU 内存分层路由。
  • 用于Tokamaxkernel 的 PyTorch 接口，与 TorchTPU 生态中的其他项目共享。
  • 通过vllm-omni提供扩散模型服务。

欢迎贡献！

面向贡献者的 TPU 容量

我们还将开放一个社区贡献 TPU 容量池。如果你正在基于vllm-torchtpu进行构建，可以申请免费 TPU 时长，Inferact 的核心 vLLM 维护者将协助设计和 PR 评审。

在此申请贡献者 TPU 访问权限。

试用

如果你有问题或想参与，请在#sig-tpu上的vLLM 开发者 Slack找到我们。