vLLM 发布 TorchTPU 后端，支持在 Google Cloud TPU 上运行 vLLM

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

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

通过复用 vLLM 现有的模型代码，让新模型能够轻松迁移到 TPU。
在开源中推动 TPU 推理性能，为注意力、混合专家和量化 matmul 提供 Pallas 内核。
保持 vLLM 体验不变：在其他硬件上使用的相同 CLI、OpenAI 兼容服务器、调度器和 KV 缓存管理器。

工作原理

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

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

快速开始

安装
```
pip install vllm-torchtpu
```
这需要在 Cloud TPU VM 上使用 Python 3.12，并会安装匹配的 vLLM、torch、torch_tpu 和 libtpu 以及该插件。安装指南还涵盖了 uv、源码构建和 Docker。

提供模型服务

vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8 \
  --tensor-parallel-size 8 \
  --enable-expert-parallel \
  --kv-cache-dtype fp8 \
  --max-model-len 16384 \
  --async-scheduling

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

推测解码、前缀缓存、多 LoRA、结构化输出和多模态输入均可通过通常的 vLLM 标志启用。有关更多示例，包括离线批量推理、多模态模型和多 LoRA，请参阅快速入门。

内核

vllm-torchtpu 对受益于在 TPU 上手动调优的复杂操作使用自定义 Pallas 内核。一些示例：

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

横向扩展

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

并行

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 模型，我们通常通过数据并行注意力与专家并行获得最佳吞吐量（例如，在单个 v7x 主机上使用 EP 的 DP=8, TP=1）。这使注意力保持在每个设备本地，并将专家分布在整个切片上。对于长上下文，PCP 和 DCP 将预填充和 KV 缓存拆分到多个设备上。并行指南提供了更多细节和示例配置。

KV Cache 卸载与分离式服务

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

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

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

投机解码与前缀缓存

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

模型、功能与后续计划

vllm-torchtpu 可以服务 vLLM 支持的大多数模型，并针对 Qwen、DeepSeek、Kimi、GLM、Llama、GPT-OSS 和 Gemma 系列进行了 TPU 特定的内核、量化和并行化工作。TPU v7x（Ironwood）是推荐硬件，并在每个拉取请求的 CI 中测试；v6e、v5e 和 v5p 为实验性支持。

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

后续计划？

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

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

欢迎贡献！

面向贡献者的 TPU 容量

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

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

试用

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