Google 在 TPU 上发布 PyTorch 原生 vLLM 与 SGLang 推理引擎

作为我们今天 PyTorch 系列公告的一部分，我们很高兴推出 vLLM（与 Inferact 合作开发）和 SGLang（与 RadixArk 合作开发）的原生 PyTorch 后端。

背景

多年来，TPU 一直以高性能和高效率支撑着大规模 AI 工作负载。随着智能体（agentic）和多模态工作负载的兴起，高效服务变得比以往任何时候都更加重要。vLLM 和 SGLang 已经通过基于 JAX 的后端在 TPU 上运行，但我们不断听到开发者表示，他们希望用 PyTorch 来自定义和调试自己的服务栈。

于是我们正式启动了 TorchTPU 项目，让 PyTorch 原生运行在 TPU 上。随后，我们在其之上构建了新的 vLLM 和 SGLang 引擎。

我们的目标是让 PyTorch 社区能够将模型带到 TPU 上，并使用开发者已经熟悉的工具让模型跑得更快。这些引擎自带经过优化的、生产就绪的服务功能以及自定义 Pallas 内核。

推理栈

大规模服务需要在技术栈的每一层进行协同设计，从芯片到编译器再到编排。今天发布的新推理引擎，正处在一个集成、开放的 TPU 生态系统的核心位置。

引擎概览

TorchTPU 让 TPU 成为原生的 PyTorch 设备。张量作为标准 PyTorch 张量在 tpu 设备上管理，没有包装类型；代码可以即时执行，也可以通过 torch.compile 使用 XLA 编译；集合通信则通过 torch.distributed 进行。

由于每个设备运行自己的 PyTorch 进程，两个引擎都保留了各自框架的 worker 模型、调度器、KV 缓存管理器和 API 服务器等，并在底层添加了 TPU 实现。

大多数模型使用 vLLM 或 SGLang 中标准的上游 PyTorch 模型定义，并在内核和层实现中加入针对 TPU 的优化。用 Pallas 编写的自定义内核作为自定义算子，在编译后的计算图中直接从 PyTorch 调用。这种平台插件式设计尽可能复用了各框架的现有能力，同时让 TPU 优化能够按自己的节奏推进。

工作负载支持

两个引擎都为全范围的生产工作负载而构建。

以下是支持的功能列表：

智能体与多轮对话：自动前缀缓存，包括针对混合模型；通过 TPU Sync 将 KV 缓存卸载到主机内存和 Managed Lustre；工具调用和结构化输出
交互式与低延迟：使用 MTP、EAGLE-3、DFlash、DSpark 和 n-gram 草稿器的投机解码
高吞吐批处理：结合专家并行的数据并行注意力
长上下文：用于预填充和解码的上下文并行；FP8 KV 缓存；主机 KV 卸载
分离式服务：通过 TPU Sync 在不同主机上分别进行预填充和解码，包括在不同并行布局之间
多模态：支持图像和视频输入的视觉语言模型；嵌入模型
超大型模型：混合与稀疏注意力架构，跨主机的张量、数据、专家、流水线和上下文并行；FP8、FP4（MXFP4、NVFP4）、INT8 和 INT4 检查点

尤其是智能体工作负载，会同时需要上述多项功能。由于智能体会反复重发不断扩大的多轮对话，要高效地服务它们，就需要找到已缓存的前缀，而不是重新计算。两个引擎都将可复用的前缀保存在 HBM 中，并在开启主机卸载后，自动将较旧的前缀移到主机内存作为第二层。

使用 TPU Sync 迁移 KV 缓存

其中两项特性将 KV 缓存移出计算它的设备。分离式服务将 prefill 和 decode 运行在独立的 TPU 主机上，因此每个请求的 KV 缓存都必须从 prefill 主机传输到 decode 主机。分层 KV 卸载将 HBM 中已放不下的那部分 KV 缓存迁移到主机的 DRAM 以及借助 Google Cloud Managed Lustre 的二级存储，并在后续请求复用该前缀时将其取回。无论哪种方式，传输都必须与计算重叠，否则会直接增加延迟。

两个引擎都为此使用 TPU Sync。它是一个用于 TPU 上 KV 缓存传输与管理的开源 C++ 库。它直接读取 PyTorch 张量背后的设备缓冲区，因此传输无需在主机上额外做一次暂存拷贝，并且它在 Python 之外的独立线程上运行传输，同时 prefill 和 decode 继续进行。vLLM-TorchTPU 通过 vLLM 的 KV connector 接口暴露它。SGLang-TorchTPU 将其用作分离式传输后端以及 HiCache 的 I/O 后端。

今日已在生产环境和公开基准测试中运行

两个引擎都已与预览客户一起在生产环境中大规模运行。为了提供透明、可复现的性能证明，我们通过社区标准的 SemiAnalysis InferenceX 项目公开进行基准测试。

在我们于 Ironwood TPU（TPU 7x）上首次提交 Qwen 3.5 的基础上，我们的工程团队进一步优化了技术栈，在高并发配置下实现了整体最低的每百万 token 成本，为 2.01 美分。伴随这些改进，我们还向 AgentX 发布了新的最先进（SOTA）性能结果，AgentX 是该项目中专门用于评估智能体长上下文和多轮工作负载的赛道。Ironwood TPU 在帕累托前沿的大部分区域领先。随着我们将更多模型和配置带上排行榜，我们期待持续更新我们的公开性能档案。

我们并非在真空中构建这些后端。通过拥抱 PyTorch 并直接与开源项目维护者合作，我们确保这些引擎保持与上游一致并以社区为先：

SGLang-TorchTPU 是与 RadixArk（由 SGLang 的创造者创立）共同开发的

这一合作让 SGLang 的优秀团队为 TPU 构建，交付了一个已在 Meta 生产环境中运行、并在严苛的真实场景中经过大规模实战检验的引擎。

^{“SGLang-TorchTPU 今天已在 Meta 的生产环境中运行，而实现这一目标离不开 Meta、Google 和 SGLang 团队之间的紧密三方协作，使该技术栈达到生产就绪。它为我们提供了一个非常成熟、PyTorch 原生的服务栈，可以毫不费力地集成到我们现有的 AI 基础设施中，并具备我们在 TPU 上高效运行前沿工作负载所需的性能、可靠性和灵活性。”Meta 的 PyTorch-TPU 负责人 Zongwei Zhou 表示。}

^{“在 RadixArk，TorchTPU 在我们于多种芯片上推出 SGLang 支持的过程中极大提升了我们的开发速度。它为我们提供了一条通往原生 TPU 加速的无痛路径，将巨大的效率提升和灵活的计算选择直接传递给我们的客户，”RadixArk 首席技术官兼联合创始人 Banghua Zhu 表示。}

vLLM-TorchTPU 是与 Inferact（由 vLLM 的创造者创立）共同开发的

我们携手为 2.8T 参数的 Kimi K3 模型集成了一个高性能 Pallas megakernel，将整个解码步骤作为单个 kernel 运行，在 Ironwood TPU 上为单个用户提供每秒超过 700 个 token 的服务。在此基础上，Inferact 最近在 vLLM-TorchTPU 栈上主导了我们首个针对 Kimi K3 的初步 AgentX 提交。这一初步里程碑已在超高交互性工作负载中展现出卓越的每 TCO 性能，同时 vLLM 和 PyTorch 生态在 TPU 上用于新前沿 OSS 模型 bring-up 的成熟度也在快速提升，后续提交将通过更大规模的 prefill/decode 分离来释放更高的效率。

^{“Inferact 的使命始终是让高性能推理在整个 AI 生态中可及、可靠且经济高效。我们非常高兴能合作推出 TorchTPU，并为 TPU 带来对 vLLM 的一流支持。通过将 PyTorch 优化的 TPU 能力直接桥接到 vLLM，开发者可以立即在 TPU 上运行生产工作负载，获得他们期望从我们栈中获得的高吞吐、低延迟和内存效率，”Inferact 首席执行官兼联合创始人 Simon Mo 表示。}

为促进持续的开源贡献，vLLM 还推出了首个社区贡献 TPU 容量池，为积极贡献两个引擎的开发者提供免费 TPU 算力。请在此处申请。

路线图

随着模型规模增大，我们正在演进两个引擎，使大规模 TPU 推理部署兼具高性能与高效率。面向生产级 agentic 工作负载的主要功能已得到支持。更多性能优化和高级功能正在推进中——欢迎提交 Github issue 申请就你的需求开展协作！

在两个引擎中，大多数流行的开放模型系列均已得到支持，包括 Qwen、DeepSeek、GLM、Kimi、GPT-OSS、Gemma4 等。团队正投入资源让更多模型可在 TPU 上原生运行，包括对新开放模型的 Day 0 支持。我们还希望为社区提供工具和 TPU 算力，以便针对自身工作负载和配置调优性能。

关于我们的 JAX 运行时

vLLM-Jax 和 SGLang-JAX 仍是重要项目。两者都将继续维护、支持，并可供社区使用和适配。

快速开始

立即开始：

vLLM-TorchTPU：https://github.com/vllm-project/vllm-torchtpu
SGLang-TorchTPU：https://github.com/sgl-project/sglang-torchtpu
TorchTPU：https://github.com/google-pytorch/torch_tpu