宣布在 TPU 上推出全新的 vLLM 和 SGLang 推理引擎

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

背景

多年来，TPU 一直以高性能和高效率支撑大规模 AI 工作负载。随着 agentic 和多模态工作负载的兴起，高效服务比以往任何时候都更加重要。这两个引擎已经通过基于 JAX 的后端在 TPU 上运行，但我们不断收到开发者的反馈，他们希望在 PyTorch 中自定义和调试其服务栈。

因此我们正式开始开发 TorchTPU，它可以在 TPU 上原生运行 PyTorch。随后我们在此基础上构建了新的 vLLM 和 SGLang 引擎。

我们的目标是让 PyTorch 社区能够将模型迁移到 TPU，并使用开发者已经熟悉的工具使其运行得更快。这些引擎配备了经过优化的、可用于生产环境的服务功能以及自定义 Pallas 内核。

推理栈

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

引擎概述

TorchTPU 使 TPU 成为 PyTorch 的原生设备。张量作为标准 PyTorch 张量在 tpu 设备上管理，无需包装类型；代码以 eager 模式运行或通过 torch.compile 使用 XLA 进行编译；集合通信通过 torch.distributed 进行。

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

大多数模型在两个引擎中都使用标准的上游 PyTorch 模型定义，TPU 特定的优化则体现在内核和层实现中。用 Pallas 编写的自定义内核作为编译图中的自定义算子直接从 PyTorch 调用。这种平台插件设计尽可能多地复用了每个框架，同时让 TPU 优化按自己的节奏推进。

工作负载支持

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

以下是支持的功能列表：

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

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

使用 TPU Sync 移动 KV 缓存

这些特性中有两个会将 KV 缓存移出计算它的设备。分离式服务（Disaggregated serving）在独立的 TPU 主机上运行预填充（prefill）和解码（decode），因此每个请求的 KV 缓存都必须从预填充主机传输到解码主机。分层 KV 卸载（Tiered KV offload）会将不再容纳于 HBM 中的那部分 KV 缓存移至主机的 DRAM 和基于 Google Cloud Managed Lustre 的二级存储中，并在后续请求复用前缀时将其取回。无论哪种方式，传输都必须与计算重叠，否则会直接增加延迟。

两个引擎都使用 TPU Sync 来完成此操作。它是一个用于 TPU 上 KV 缓存传输与管理的开源 C++ 库。它直接读取 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，当我们跨多种芯片推出 SGLang 支持时，TorchTPU 极大地提升了我们的开发速度。它为我们提供了一条轻松实现原生 TPU 加速的途径，将巨大的效率提升和灵活的算力选择直接传递给我们的客户。”RadixArk 联合创始人兼 CTO Banghua Zhu 表示。}

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

我们共同为 2.8T 参数的 Kimi K3 模型集成了高性能 Pallasmegakernel，将整个解码步骤作为一个内核运行，在 Ironwood TPU 上为单个用户每秒提供超过 700 个 token。在此基础上，Inferact 最近主导了我们首次初步的AgentX 提交，针对 vLLM-TorchTPU 技术栈上的 Kimi K3。这一初始里程碑已经展示了在超高交互性工作负载中卓越的每 TCO 性能，以及 vLLM 和 PyTorch 生态系统在 TPU 上用于新前沿开源模型 bring-up 的成熟度快速提升，后续提交将通过更大规模的 prefill/decode 分离来释放更高的效率。

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

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

申请加入 vLLM TPU 贡献者池

路线图

随着模型变得越来越大，我们正在改进两个引擎，使大规模 TPU 推理部署兼具高性能和高效率。面向生产级 agentic 工作负载的主要功能已经得到支持。更多性能优化和高级功能正在开发中——打开一个 Github issue 来请求协作解决你的需求！

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

关于我们的 JAX 运行时

vLLM-JAX 和 SGLang-JAX 仍是重要的项目。两者将继续得到维护、支持，并供社区使用和改编。

快速开始

Github 项目：

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