Google 发布 vLLM-TorchTPU 与 SGLang-TorchTPU 推理引擎

作为我们今日为 PyTorch 发布的更广泛公告的一部分，我们很高兴推出 vLLM(与 Inferact 合作开发)和 SGLang(与 RadixArk 合作开发)的原生 PyTorch 后端。

背景

多年来，TPU 一直以高性能和高效率支撑着大规模 AI 工作负载。随着智能体式和多模态工作负载的兴起，高效服务比以往任何时候都更加重要。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 缓存远离计算它的设备。分离式服务将预填充和解码运行在不同的 TPU 主机上，因此每个请求的 KV 缓存都必须从预填充主机传输到解码主机。分层 KV 卸载将不再能放入 HBM 的那部分 KV 缓存移到主机的 DRAM 以及通过 Google Cloud Managed Lustre 的二级存储，并在后续请求复用该前缀时将其取回。无论哪种方式，传输都必须与计算重叠，否则会直接增加延迟。

两个引擎都为此使用 TPU Sync。它是一个开源 C++ 库，用于在 TPU 上进行 KV 缓存传输和管理。它直接读取 PyTorch 张量背后的设备缓冲区，因此传输无需在主机上进行额外的暂存复制，并且它在 Python 之外的自己的线程上运行传输，同时预填充和解码继续进行。vLLM-TorchTPU 通过 vLLM 的 KV 连接器接口将其暴露出来。SGLang-TorchTPU 将其用作分离式传输后端以及 HiCache 的 I/O 后端。

如今已在生产环境与公开基准中运行

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

在我们于 Ironwood TPU（TPU 7x）上首次提交 Qwen 3.5 的基础上，我们的工程团队进一步优化了技术栈，在高并发配置下实现了总体每百万 token 2.01 美分的最低成本。除了这些改进之外，我们还在 AgentX 上发布了新的最先进（SOTA）性能结果，这是该项目专门用于评估智能体长上下文和多轮工作负载的赛道。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 的 CTO 兼联合创始人 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 上用于新前沿开源模型适配的成熟度也在快速提升，后续提交将通过更大规模的 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