vLLM 推出 TorchTPU 后端支持 Google TPU 原生推理

出处作者 / 发布主体：Inferact and Google Teams原标题：Introducing vLLM TorchTPU

78AI 研判分

核心综述

vLLM 正式开源 vllm-torchtpu 后端，允许开发者在 Google Cloud TPUs 上直接运行现有模型而无需修改代码。该工具基于 Inferact 与 Google 的合作项目，通过复用 vLLM 的模型定义和标准接口（如 CLI、OpenAI 兼容服务器），显著降低了将大模型部署至 TPU 硬件的门槛。

报道全文原始报道全文

本文目录11 个章节

今天，我们发布了 vLLM TorchTPU（vllm-torchtpu），这是一个用于在 Google Cloud TPU 上运行 vLLM 的新后端。它基于 Google 的 TorchTPU 技术栈构建，并直接复用 vLLM 中现有的模型定义，因此目前能在 vLLM 中运行的绝大多数模型无需修改代码即可部署到 TPU 上。

这项工作源于上个月我们宣布的 Inferact 与 Google 的工程合作。我们的项目目标包括：

通过复用 vLLM 现有的模型代码，轻松将新模型引入 TPU。
在开源社区推动 TPU 推理性能的提升，提供针对注意力机制、混合专家模型（MoE）以及量化矩阵乘法优化的 Pallas 内核。
保持 vLLM 的使用体验一致：沿用你在其他硬件上使用的相同 CLI、OpenAI 兼容服务器、调度器和 KV Cache 管理器。

工作原理

vllm-torchtpu 是一个树外（out-of-tree）的 vLLM 平台插件。安装后，vLLM 会通过 vllm.platform_plugins 入口点自动识别并加载该插件。

模型通过 vLLM 标准的模型运行器执行，并使用 torch.compile(backend="tpu") 进行编译，最终降低为 XLA 代码。每个 TPU 设备拥有独立的 worker 进程，这与 vLLM 在 GPU 上的运行方式相同；设备间的通信则通过 torch.distributed 实现。这意味着调度器、worker、KV Cache 管理器和 API 服务器均未经过修改，完全保留自上游 vLLM。TPU 特有的工作位于这些组件之下，主要体现为一组注册为自定义算子并在编译图中调用的 JAX/Pallas 内核。

快速开始

安装

文本

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 可在重启时复用已编译的产物（参见 Compilation & Warmup Cache）。

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

Kernels

vllm-torchtpu 使用自定义 Pallas kernels 来处理那些在 TPU 上通过手动调优能获益的复杂操作。以下是一些示例：

用于 MHA/GQA 模型的 Ragged Paged Attention v3
用于 DeepSeek、Kimi 和 GLM 的稠密与稀疏 Multi-head Latent Attention (MLA)
用于混合线性注意力模型的 Gated DeltaNet 和 Kimi Delta Attention
支持专家并行（Expert Parallelism）的 Fused Mixture-of-Experts
FP8、INT8、MXFP4 和 NVFP4 量化矩阵乘法

Scaling Out

我们重点关注了当前给服务系统带来最大压力的工作负载，尤其是多轮 agentic 流量和多模态流量，在这些场景中上下文不断增长且工具输出可能很长。

Parallelism

vLLM 的所有六个并行维度均受支持，无论是在单主机还是多主机 TPU 切片上，并且它们可以组合使用：

Dimension	Flag
Tensor Parallelism (TP)	`--tensor-parallel-size`
Data Parallelism (DP)	`--data-parallel-size`
Pipeline Parallelism (PP)	`--pipeline-parallel-size`
Expert Parallelism (EP)	`--enable-expert-parallel`
Prefill Context Parallelism (PCP)	`--prefill-context-parallel-size`
Decode Context Parallelism (DCP)	`--decode-context-parallel-size`

对于大型 MoE 模型，通常采用数据并行注意力结合专家并行（例如在单个 v7x 主机上使用 DP=8, TP=1 并启用 EP）能获得最佳吞吐量。这种方式将注意力保留在每个设备本地，并将专家分布到整个切片中。对于长上下文，PCP 和 DCP 会将预填充（prefill）和 KV cache 拆分到多个设备上。parallelism guide 提供了更多细节和示例配置。

KV Cache Offloading and Disaggregated Serving

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 访问权限。

立即体验

如有任何问题或希望参与进来，请在 vLLM 开发者 Slack 的 #sig-tpu 频道中找到我们。