Announcing New vLLM and SGLang Inference Engines on TPU

在TPU上推出新的vLLM和SGLang推理引擎 - Google开发者博客

Google Tag Manager (noscript)

结束Google Tag Manager (noscript)

HTML

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

2026年10月20日

周琪

首席软件工程师

陈志

产品管理总监

分享

• Facebook

• Twitter

• LinkedIn

• 邮件

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

背景

多年来，TPU一直以高性能和高效率支持大规模AI工作负载。随着智能体和多模态工作负载的兴起，高效服务比以往任何时候都更加重要。虽然vLLM和SGLang目前通过基于JAX的后端在TPU上运行，但我们不断收到来自开发者的反馈，他们希望能够在PyTorch中自定义和调试服务堆栈。

因此我们正式启动了TorchTPU项目，该项目可在TPU上原生运行PyTorch。随后，我们在其基础上构建了新的vLLM和SGLang引擎。

我们的目标是让PyTorch社区能够将模型部署到TPU上，并使用开发者熟悉的工具使其运行加速。这些引擎包含经过优化的、生产就绪的服务功能和自定义的Pallas内核。

推理堆栈

大规模服务需要在技术堆栈的每一层（从芯片到编译器到编排）进行协同设计。今天推出的新型推理引擎是集成开放TPU生态系统的核心组成部分。

引擎概述

TorchTPU将TPU作为原生PyTorch设备。张量作为标准PyTorch张量在tpu设备上进行管理，无需包装类型，代码可以急切执行或通过torch.compile使用XLA进行编译，集合操作通过torch.distributed进行。

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

大多数模型使用vLLM或SGLang中标准的上游PyTorch模型定义，但在内核和层实现中包含特定于TPU的优化。用Pallas编写的自定义内核可以直接从PyTorch的编译图中作为自定义操作符调用。这种平台插件设计在尽可能重用每个框架现有功能的同时，使TPU优化能够按照自己的节奏推进。

工作负载支持

两个引擎都针对完整的生产工作负载范围进行了构建。

以下是支持的功能列表：

• 智能体与多轮对话：自动前缀缓存（包括混合模型）；通过TPU Sync将KV缓存卸载到主机内存和管理式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实现这一功能。这是一个用于TPU上KV缓存传输和管理的开源C++库。它可以直接读取PyTorch张量背后的设备缓冲区，因此传输无需在主机上进行额外的中转复制，且可以在独立线程上运行传输操作（与预填充和解码过程并行）。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基础设施中，具备我们高效运行前沿工作负载所需的性能、可靠性和灵活性。" Meta的PyTorch-TPU负责人Zongwei Zhou表示。

"在 RadixArk，TorchTPU 极大地提升了我们的开发速度，使我们能够跨多种芯片架构快速部署 SGLang 支持。它为我们提供了一条无缝实现原生 TPU 加速的路径，将巨大的效率提升和灵活的计算选择直接传递给我们的客户，" RadixArk 的 CTO 和联合创始人朱邦华表示。

vLLM-TorchTPU 与 Inferact（由 vLLM 的创建者创立）联合开发

我们共同集成了高性能的 Pallas 兆核，用于 2.8T 参数的 Kimi K3 模型，该模型将整个解码步骤作为单一内核运行，在 Ironwood TPU 上每秒可为单个用户处理超过 700 个 token。基于这项工作，Inferact 最近主导了我们首次在 vLLM-TorchTPU 架构上为 Kimi K3 提交的初步 AgentX 方案。这一初步里程碑已在超高交互性工作负载中展现出卓越的每总拥有成本（TCO）性能，而 vLLM 和 PyTorch 生态系统在 TPU 上进行新前沿开源模型部署的成熟度迅速提升，后续提交将通过更大规模的预填充/解码分离实现更高效率。

"Inferact 的使命一直是让高性能推理在整个 AI 生态系统中都具备可访问性、可靠性和成本效益。我们非常高兴能与 RadixArk 合作推出 TorchTPU，并为 TPUs 提供一流的 vLLM 支持。通过将 PyTorch 优化的 TPU 能力直接集成到 vLLM 中，开发者可以立即在 TPUs 上运行生产工作负载，并获得我们堆栈所保证的高吞吐量、低延迟和内存效率，" Inferact 的 CEO 和联合创始人莫Simon 表示。

为了促进持续的开源贡献，vLLM 还推出了首个社区贡献 TPU 资源池，为积极贡献两个引擎的开发者提供免费的 TPU 计算资源。立即申请

路线图

随着模型规模的增大，我们正在改进两个引擎，以实现大规模 TPU 推理部署的高性能和高效率。生产级代理工作负载的主要功能已经得到支持。更多的性能优化和高级功能正在开发中 – 在 GitHub 上打开问题以请求协作！

在两个引擎中，大多数流行的开源模型系列都已得到支持，包括 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

相关公告：

• Inferact 宣布

• Radixark 宣布

发布在：

• AI

• 云

• 公告

• 最佳实践

• 文档

• 社区

• 学习

• 探索

• 影响

上一篇

下一篇

相关文章

列表

移动

AI

公告

行业趋势

统一边缘 AI：Coral NPU 的开放标准如何将便携、节能的智能从模型带到芯片

案例研究

学习

使用 Model Explorer 将可视化图调试引入自定义边缘编译器

云

最佳实践

将 TorchTPU 集成到 TorchTitan 中并为 TPU 优化模型

导航点