介绍 vLLM TorchTPU

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

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

1. 通过复用 vLLM 现有的模型代码，让新模型轻松移植到 TPU。
2. 在开源中推动 TPU 推理性能，包括用于注意力机制、Mixture-of-Experts 和量化 matmul 的 Pallas 内核。
3. 保持 vLLM 体验不变：与其他硬件上相同的 CLI、OpenAI 兼容服务器、调度器和 KV cache 管理器。

工作原理

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

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

入门

1. 安装

   pip install vllm-torchtpu

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

2. 部署模型

   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
• 带专家并行性的融合专家混合（MoE）
• FP8、INT8、MXFP4 和 NVFP4 量化矩阵乘法

横向扩展

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

并行性

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

对于大型 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 专用的内核、量化和并行方面的工作。TPU v7x (Ironwood) 是推荐的硬件，并在每次拉取请求时于 CI 中进行测试；v6e、v5e 和 v5p 为实验性支持。

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

后续计划？

我们希望各大模型家族（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 访问权限。

动手试试

• 快速入门
• Cloud TPU 设置
• TPU v7x (Ironwood) 服务化配方
• Ironwood 上的 vLLM TorchTPU 配方
• Kubernetes 部署
• 贡献指南

如果您有问题或想参与其中，请在 #sig-tpu 上找到我们： vLLM 开发者 Slack.