Google 发布 vLLM 与 SGLang 的 TorchTPU 原生推理引擎

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

两个引擎都为此使用 TPU Sync。它是一个用于 TPU 上 KV cache 传输与管理的开源 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 模型引入的成熟度正在快速提升，后续提交将通过更大规模的 prefill/decode 分离来释放更高的效率。

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

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

路线图

随着模型规模增大，我们正在演进这两个引擎，使大规模 TPU 推理部署兼具高性能和高效率。生产级智能体工作负载的主要功能已得到支持。更多性能优化和高级功能正在推进中——请提交 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