vLLM 与 SGLang 新增 TPU 原生后端

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

多年来，TPU 一直以高性能和高效率支撑着大规模 AI 工作负载。随着 Agent 和多模态工作负载兴起，高效的推理服务变得比以往更加重要。vLLM 和 SGLang 已经可以通过基于 JAX 的后端在 TPU 上运行，但开发者不断向我们反馈：他们希望能在 PyTorch 中定制和调试自己的推理服务栈。

因此，我们开始全力投入 TorchTPU，让 PyTorch 在 TPU 上原生运行。随后，我们在它的基础上构建了全新的 vLLM 和 SGLang 引擎。

我们的目标是让 PyTorch 社区能够使用开发者已经熟悉的工具，把模型带到 TPU 上，并让它们高效运行。这两个引擎提供了经过优化、可用于生产环境的推理服务功能，以及定制的 Pallas kernel。

要实现大规模推理服务，就需要从芯片、编译器到编排系统，对技术栈的每一层进行协同设计。今天发布的新推理引擎，正是这个一体化、开放的 TPU 生态系统的核心组成部分。

TorchTPU 让 TPU 成为 PyTorch 的原生设备。张量以标准 PyTorch tensor 的形式在 tpu 设备上管理，不需要包装类型；代码既可以采用即时执行模式，也可以通过 torch.compile 使用 XLA 编译；集合通信则通过 torch.distributed 完成。

由于每个设备都运行着独立的 PyTorch 进程，两个引擎都保留了各自框架的 worker 模型、调度器、KV cache 管理器和 API server 等组件，并在底层加入 TPU 实现。

大多数模型直接使用 vLLM 或 SGLang 上游的标准 PyTorch 模型定义，针对 TPU 的优化则放在 kernel 和层实现中。用 Pallas 编写的自定义 kernel，会作为编译图中的自定义算子，由 PyTorch 直接调用。这种平台插件设计尽可能复用各框架已有的能力，同时让 TPU 优化可以按自身节奏推进。

这两个引擎面向完整的生产工作负载而构建。

以下是支持的功能：

• Agent 与多轮对话：自动前缀缓存，包括对混合架构模型的支持；通过 TPU Sync 将 KV cache 卸载到主机内存和 Managed Lustre；tool calling 和结构化输出。
• 交互与低延迟：使用 MTP、EAGLE-3、DFlash、DSpark 和 n-gram 草稿生成器进行推测解码。
• 高吞吐批处理：结合 expert parallelism 的 data-parallel attention。
• 长上下文：支持 prefill 和 decode 阶段的 context parallelism；FP8 KV cache；将 KV cache 卸载到主机。
• 分离式推理服务：通过 TPU Sync，在不同主机上分别执行 prefill 和 decode，也支持两者采用不同的并行布局。
• 多模态：支持图像和视频输入的视觉语言模型；embedding 模型。
• 超大模型：混合架构和稀疏注意力架构；跨主机的 tensor、data、expert、pipeline 和 context parallelism；FP8、FP4（MXFP4、NVFP4）、INT8 和 INT4 checkpoint。

尤其是 Agent 工作负载，往往需要同时使用其中多项功能。由于 Agent 会反复发送不断增长的多轮对话，要高效处理这些请求，就必须找到已缓存的前缀，避免重新计算。两个引擎都会把可复用的前缀保存在 HBM 中；开启主机卸载后，还会自动将较旧的前缀移到主机内存，作为第二级缓存。

使用 TPU Sync 移动 KV cache

上述功能中，有两项需要把 KV cache 从计算它的设备上移走。分离式推理服务会在不同的 TPU 主机上分别执行 prefill 和 decode，因此，每个请求的 KV cache 都必须从 prefill 主机传输到 decode 主机。分层 KV 卸载则会把 HBM 容纳不下的那部分 KV cache 移到主机 DRAM，以及由 Google Cloud Managed Lustre 提供的二级存储中；后续请求复用该前缀时，再将其取回。无论哪种情况，传输都必须与计算重叠进行，否则就会直接增加延迟。

两个引擎都使用 TPU Sync 来完成这些工作。这是一个开源 C++ 库，用于 TPU 上的 KV cache 传输和管理。它直接读取 PyTorch tensor 底层的设备缓冲区，因此传输时不需要在主机上额外复制一份数据作为中转。它还会在 Python 之外，使用自己的线程执行传输，让 prefill 和 decode 可以继续运行。vLLM-TorchTPU 通过 vLLM 的 KV connector 接口接入它；SGLang-TorchTPU 则将它用作分离式服务的传输后端，以及 HiCache 的 I/O 后端。

已在生产环境运行，并公开基准测试结果

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

在最初提交 Ironwood TPU（TPU 7x）上运行 Qwen 3.5 的结果之后，我们的工程团队进一步优化了整个技术栈，在高并发配置下，将每百万 tokens 的成本降至 2.01 美分，达到整体最低水平。除了这些改进，我们还向 AgentX 发布了新的业界最佳（SOTA）性能结果。AgentX 是该项目专门用于评估 Agent 长上下文和多轮交互工作负载的赛道。Ironwood TPU 在 Pareto 前沿的大部分区间都处于领先位置。随着更多模型和配置进入排行榜，我们期待持续更新公开的性能表现数据。

与社区共同构建

这些后端并不是我们闭门造车的产物。通过拥抱 PyTorch，并直接与开源项目维护者合作，我们确保这些引擎始终与上游保持一致，并把社区放在首位：

SGLang-TorchTPU 与 RadixArk 联合开发，后者由 SGLang 的创建者创立

这项合作让 SGLang 的优秀团队参与 TPU 开发，交付的引擎已经在 Meta 的生产环境运行，并在要求严苛的真实场景中，经受了大规模部署的检验。

vLLM-TorchTPU 与 Inferact 联合开发，后者由 vLLM 的创建者创立

我们共同为拥有 2.8T 参数的 Kimi K3 模型集成了一个高性能 Pallas megakernel，将整个 decode 步骤合并为单个 kernel 执行，在 Ironwood TPU 上为单个用户提供超过 700 tokens/s 的生成速度。在此基础上，Inferact 最近牵头完成了我们首份基于 vLLM-TorchTPU 技术栈运行 Kimi K3 的 AgentX 初步结果提交。这个初步里程碑已经展现出：在交互要求极高的工作负载中，该方案按总拥有成本（TCO）衡量的性能十分出色；同时，TPU 上的 vLLM 和 PyTorch 生态也在迅速成熟，能够更快地适配新的前沿开源模型。后续提交将通过更大规模的 prefill/decode 分离，进一步提升效率。

为了鼓励持续的开源贡献，vLLM 还推出了首个面向社区贡献的 TPU 算力池，为积极参与这两个引擎贡献的开发者提供免费 TPU 算力。现在即可申请。

随着模型规模不断扩大，我们也在持续演进这两个引擎，让大规模 TPU 推理部署兼具性能与效率。生产级 Agent 工作负载所需的主要功能已经得到支持。更多性能优化和高级功能正在开发中——欢迎提交 GitHub issue，就你需要的功能发起合作！

这两个引擎已经支持大多数热门开放模型家族，包括 Qwen、DeepSeek、GLM、Kimi、GPT-OSS、Gemma4 等。团队正在投入资源，让更多模型能够在 TPU 上原生运行，包括在新开放模型发布当天就提供支持，也就是 Day 0 支持。我们还希望为社区提供工具和 TPU 算力，让开发者能够针对自己的工作负载和配置调优性能。

关于我们的 JAX runtime

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

vLLM-TorchTPU

SGLang-TorchTPU

TorchTPU

相关公告：

Inferact 公告

Radixark 公告