推出 TPU 上全新的 vLLM 和 SGLang 推理引擎
原题Announcing New vLLM and SGLang Inference Engines on TPU
读前速览
Google 推出适用于 vLLM 和 SGLang 的原生 PyTorch 后端,让 TPU 能够为智能体和多模态工作负载提供高性能、可定制的推理能力。
原生支持 PyTorch
这些引擎基于 TorchTPU 构建,开发者可以使用标准 PyTorch 语法和工具编写、调试推理代码,无需依赖 JAX 或专用封装。
全面支持各类工作负载
借助 KV 缓存卸载和数据并行注意力等功能,这些引擎能够处理多轮对话、高吞吐量批处理和长上下文任务等多种复杂工作负载。
与社区协作
这些引擎的开发得到了 vLLM 原始创建团队 Inferact 和 SGLang 原始创建团队 RadixArk 的密切合作,确保后端与上游项目深度集成并保持一致。
公开验证性能
Google 在 SemiAnalysis InferenceX 排行榜上公布了测试结果,显示新引擎实现了最低的每 token 成本和出色的交互性评分,证明其已具备投入生产的能力。
完整摘要
基于 TorchTPU 构建的全新 vLLM 和 SGLang 引擎,为在 Google TPU 上运行 AI 推理提供了完整集成的开源方案。这些引擎支持广泛的生产工作负载,包括智能体、多模态和长上下文任务,并具备自动前缀缓存、推测解码和解耦式服务等功能。此次发布还得到了社区伙伴的大力支持,并公布了基准测试结果,展现出业界领先的性能。
目录 · 10 节
作为我们今天 PyTorch 系列重要公告的一部分,我们很高兴推出 vLLM(与 Inferact 合作开发)和 SGLang(与 RadixArk 合作开发)的原生 PyTorch 后端。As part of our broader announcements for PyTorch today, we are excited to launch native PyTorch backends for vLLM (developed in partnership with Inferact) and SGLang (developed in partnership with RadixArk).
背景Background
多年来,TPU 一直凭借出色的性能和效率,为大规模 AI 工作负载提供支持。随着智能体和多模态工作负载兴起,高效地提供服务比以往任何时候都更加重要。vLLM 和 SGLang 已经通过基于 JAX 的后端在 TPU 上运行,但我们不断听到开发者表示,他们希望能在 PyTorch 中定制和调试自己的服务技术栈。For years, TPUs have powered large-scale AI workloads with high performance and efficiency. With the rise of agentic and multimodal workloads, serving efficiently is more important than ever. vLLM and SGLang already run on TPU through JAX-based backend, but we kept hearing from developers that they wanted to customize and debug their serving stack in PyTorch.
于是我们正式着手开发 TorchTPU,让 PyTorch 能在 TPU 上原生运行。随后,我们在此基础上构建了新的 vLLM 和 SGLang 引擎。So we started in earnest on TorchTPU, which runs PyTorch natively on TPU. We then built new vLLM and SGLang engines on top of it.
我们的目标是让 PyTorch 社区能够使用开发者已经熟悉的工具,将模型部署到 TPU 上并提升其运行速度。这些引擎配备了经过优化、可用于生产环境的服务功能以及定制的 Pallas 内核。Our goal is for the PyTorch community to bring models to TPU and make them fast using the tools developers are already familiar with. The engines come with optimized, production-ready serving features and custom Pallas kernels.
推理栈The Inference Stack
要实现大规模服务,需要对技术栈的每一层进行协同设计,从硅芯片到编译器,再到编排。今天推出的新推理引擎,是集成式开放 TPU 生态系统的核心组成部分。Serving at scale requires co-design across every layer of the technology stack, from silicon to compilers to orchestration. Today's launch of new inference engines is at the central part of an integrated, open TPU ecosystem.
引擎概述Engine Overview
TorchTPU 让 TPU 成为 PyTorch 原生设备。张量作为标准 PyTorch 张量在 tpu 设备上进行管理,无需包装类型;代码可即时执行,也可通过 torch.compile 使用 XLA 编译,集合通信则通过 torch.distributed 完成。TorchTPU makes the TPU a native PyTorch device. Tensors are managed as standard PyTorch tensors on a tpu device with no wrapper types, code runs eagerly or compiles via torch.compile using XLA, and collectives go through torch.distributed.
由于每台设备都运行各自的 PyTorch 进程,因此两个引擎都保留了各自框架的 worker 模型、调度器、KV 缓存管理器和 API 服务器等,并在底层添加了 TPU 实现。Because each device runs its own PyTorch process, both engines keep each framework’s worker model, scheduler, KV cache manager and API server etc, and add TPU implementations underneath.
大多数模型使用 vLLM 或 SGLang 中标准的上游 PyTorch 模型定义,并在内核和层实现中加入针对 TPU 的优化。使用 Pallas 编写的自定义内核,会作为自定义算子由 PyTorch 在编译图中直接调用。这种平台插件设计尽可能复用各框架的现有功能,同时让 TPU 优化按照自身节奏推进。Most models use the standard upstream PyTorch model definitions in vLLM or SGLang, with TPU-specific optimizations in kernels and layer implementations. Custom kernels written in Pallas are called directly from PyTorch as custom operators inside the compiled graph. This platform-plugin design reuses as much of each framework as possible while letting TPU optimizations move at their own pace
工作负载支持Workload support
这两款引擎均针对各种生产工作负载而构建。Both engines are built for the full range of production workloads.
以下是支持的功能列表:Here’s the list of supported features:
- 智能体式与多轮交互:自动前缀缓存(包括混合模型);通过 TPU Sync 将 KV 缓存卸载到主机内存和 Managed Lustre;工具调用和结构化输出Agentic & Multi-Turn: Automatic prefix caching, including for hybrid models; KV cache offloading to host memory and Managed Lustre through TPU Sync; tool calling and structured output
- 交互式与低延迟:使用 MTP、EAGLE-3、DFlash、DSpark 和 n-gram 草稿模型进行推测解码Interactive & Low Latency: Speculative decoding with MTP, EAGLE-3, DFlash, DSpark and n-gram drafters
- 高吞吐批处理:采用专家并行的数据并行注意力机制High-throughput Batch: Data-parallel attention with expert parallelism
- 长上下文:用于预填充和解码的上下文并行;FP8 KV 缓存;主机端 KV 卸载Long Context: Context parallelism for prefill and decode; FP8 KV cache; host KV offloading
- 解耦式推理服务:通过 TPU Sync 在不同主机上执行预填充和解码,支持不同并行布局之间的切换。Disaggregated Serving: Prefill and decode on separate hosts over TPU Sync, including between different parallel layouts
- 多模态:支持图像和视频输入的视觉语言模型;嵌入模型Multimodal: Vision-language models with image and video input; embedding models
- 超大型模型:混合式和稀疏注意力架构;跨主机的张量、数据、专家、流水线和上下文并行;FP8、FP4(MXFP4、NVFP4)、INT8 和 INT4 模型检查点Very Large Models: Hybrid and sparse-attention architectures, tensor, data, expert, pipeline and context parallelism across hosts; FP8, FP4 (MXFP4, NVFP4), INT8 and INT4 checkpoints
智能体工作负载尤其需要同时具备其中多项特性。由于智能体会反复重新发送不断扩展的多轮对话,要高效地为其提供服务,就必须找到缓存的前缀,而不是重新计算。两种引擎都会将可复用的前缀保存在 HBM 中;启用主机卸载后,还会自动将较旧的前缀移至主机内存,作为第二级存储。Agentic workloads in particular demand several of these features simultaneously. Because agents repeatedly resend expanding multi-turn conversations, serving them efficiently requires finding the cached prefix instead of recomputing it. Both engines keep reusable prefixes in HBM and, with host offload turned on, automatically move older prefixes to host memory as a second tier.
使用 TPU Sync 移动 KV 缓存Moving the KV Cache with TPU Sync
其中两项功能会将 KV 缓存从计算它的设备上移走。解耦式服务会在不同的 TPU 主机上分别执行预填充和解码,因此每个请求的 KV 缓存都必须从预填充主机传输到解码主机。分层 KV 卸载会将不再能容纳于 HBM 中的 KV 缓存部分移至主机的 DRAM 和二级存储(使用 Google Cloud Managed Lustre),并在后续请求复用该前缀时将其取回。无论哪种方式,传输都必须与计算重叠,否则就会直接增加延迟。Two of these features move the KV cache away from the device that computed it. Disaggregated serving runs prefill and decode on separate TPU hosts, so each request's KV cache has to travel from the prefill host to the decode host.Tiered KV offload moves the part of the KV cache that no longer fits in HBM to the host's DRAM and secondary storage with Google Cloud Managed Lustre, and brings it back when a later request reuses the prefix. Either way, the transfer has to overlap computation or it adds directly to latency.
两个引擎都使用 TPU Sync 来实现这一功能。它是一个开源 C++ 库,用于在 TPU 上传输和管理 KV 缓存。它会直接读取 PyTorch 张量背后的设备缓冲区,因此传输无需在主机上进行额外的暂存复制;同时,它会在 Python 之外通过自身线程执行传输,而预填充和解码过程可以继续进行。vLLM-TorchTPU 通过 vLLM 的 KV 连接器接口提供该库。SGLang-TorchTPU 则将其用作解耦传输后端,以及 HiCache 的 I/O 后端。Both engines use TPU Sync for this. It is an open-source C++ library for KV cache transfer and management on TPU. It reads the device buffers behind PyTorch tensors directly, so a transfer needs no extra staging copy on the host, and it runs transfers on its own threads outside Python while prefill and decode continue. vLLM-TorchTPU exposes it through vLLM's KV connector interface. SGLang-TorchTPU uses it as a disaggregation transfer backend and as the I/O backend for HiCache.
目前已在生产环境和公开基准测试中运行Running today in production & public benchmarks
两款引擎都已在面向预览客户的生产环境中实现规模化运行。为透明、可复现地证明我们的性能,我们通过社区标准项目 SemiAnalysis InferenceX 公开进行基准测试。Both engines are already running at scale in production with preview customers. To provide transparent, reproducible proof of our performance, we benchmark in public through the community-standard SemiAnalysis InferenceX project.
在最初提交 Qwen 3.5 在 Ironwood TPU(TPU 7x)上的成绩后,我们的工程团队进一步优化了技术栈,在高并发配置下实现了整体最低的每百万 token 成本——2.01 美分。与此同时,我们还在 AgentX 上发布了全新的最先进(SOTA)性能结果。AgentX 是该项目专门用于评估智能体长上下文和多轮工作负载的赛道。Ironwood TPU 在帕累托前沿的大多数区间均处于领先地位。随着我们向排行榜提交更多模型和配置,我们期待持续更新公开的性能档案。Building on our initial submission of Qwen 3.5 on Ironwood TPU (TPU 7x), our engineering teams have further optimized the stack, bringing the lowest cost per million tokens overall at 2.01 cents at high-concurrency configurations. Alongside these improvements, we have also published new state-of-the-art (SOTA) performance results to AgentX, the project's dedicated track for evaluating agentic long-context and multi-turn workloads. Ironwood TPU leads across most of the Pareto frontier. We look forward to continuously updating our public performance profiles as we bring more models and configurations to the leaderboard.
与社区携手建设Building with the community
我们并非闭门造车地开发这些后端。通过采用 PyTorch 并与开源项目维护者直接合作,我们确保这些引擎始终与上游保持一致,并以社区为先:We did not build these backends in a vacuum. By embracing PyTorch and working directly with open-source project maintainers, we ensure these engines remain upstream-aligned and community-first:
SGLang-TorchTPU was co-developed with RadixArk (founded by the creators of SGLang)
这项合作将由 SGLang 的卓越团队为 TPU 构建引擎。该引擎已在 Meta 投入生产,并在各种严苛的真实场景中经受了大规模实战检验。This partnership is bringing SGLang's stellar team to build for TPU, delivering an engine that is already running in production at Meta and battle-tested at scale across demanding real-world scenarios.
“如今,SGLang-TorchTPU 已在 Meta 投入生产使用。要实现这一目标,Meta、Google 和 SGLang 团队进行了紧密的三方协作,共同打造出可用于生产环境的技术栈。它为我们提供了一套非常成熟、原生支持 PyTorch 的服务技术栈,能够轻松融入现有的 AI 基础设施,并具备我们在 TPU 上高效运行前沿工作负载所需的性能、可靠性和灵活性。”Meta 的 PyTorch-TPU 负责人 Zongwei Zhou 表示。"SGLang-TorchTPU is running in production at Meta today, and getting there was a close three-way collaboration between Meta, Google, and the SGLang team to make the stack production-ready. It gives us a remarkably mature, PyTorch-native serving stack that integrates effortlessly into our existing AI infrastructure, with the performance, reliability, and flexibility we need to run frontier workloads efficiently on TPUs. " said Zongwei Zhou, PyTorch-TPU lead at Meta.
RadixArk 的 CTO 兼联合创始人朱邦华表示:“TorchTPU 大幅提升了我们的开发速度,让我们能够在各种不同的芯片上推出对 SGLang 的支持。它为我们提供了一条轻松实现 TPU 原生加速的途径,并将显著的效率提升和灵活的计算选择直接带给客户。”"At RadixArk, TorchTPU has supercharged our development velocity as we roll out SGLang support across diverse silicon. It’s given us a painless path to native TPU acceleration, passing massive efficiency gains and flexible compute choices directly down to our customers," said Banghua Zhu, CTO and co-founder of RadixArk.
vLLM-TorchTPU 由 Inferact(由 vLLM 的创作者创立)共同开发。vLLM-TorchTPU was co-developed with Inferact (founded by the creators of vLLM)
我们共同为拥有 2.8 万亿参数的 Kimi K3 模型集成了一个高性能 Pallas megakernel,可将整个解码步骤作为单个 kernel 运行,在 Ironwood TPU 上为单个用户提供每秒超过 700 个 token 的服务。在此基础上,Inferact 最近牵头提交了我们首个针对 vLLM-TorchTPU 技术栈上 Kimi K3 的 AgentX 初步方案。这一初步里程碑已经展现出超高交互性工作负载下卓越的单位 TCO 性能,也体现了 TPU 上 vLLM 和 PyTorch 生态系统在适配新一代前沿开源模型方面日益成熟。后续提交方案将通过更大规模的预填充/解码分离,进一步提升效率。Together, we integrated a high-performance Pallas megakernel for the 2.8T-parameter Kimi K3 model that runs the entire decode step as one kernel, serving more than 700 tokens per second to a single user on Ironwood TPU. Building on this work, Inferact recently led our first preliminary AgentX submission for Kimi K3 on the vLLM-TorchTPU stack. This initial milestone already demonstrates exceptional performance per TCO in ultra-high-interactivity workloads, and the rapid increase in the maturity of the vLLM and PyTorch ecosystem on TPU for doing new frontier OSS models bring-up, with follow-up submissions set to unlock even greater efficiency through larger scale prefill/decode disaggregation.
“Inferact 的使命始终是让整个 AI 生态系统都能获得高性能、可靠且经济高效的推理服务。我们很高兴能携手推出 TorchTPU,并为 TPU 带来对 vLLM 的一流支持。通过将 PyTorch 优化的 TPU 能力直接集成到 vLLM 中,开发者现在就能在 TPU 上运行生产级工作负载,享受他们对我们的技术栈所期待的高吞吐量、低延迟和高内存效率。”Inferact 首席执行官兼联合创始人 Simon Mo 表示。"Our mission at Inferact has always been to make high-performance inference accessible, reliable, and cost-effective across the entire AI ecosystem. We are thrilled to partner on the launch of TorchTPU and bring first-class support for vLLM to TPUs. By bridging PyTorch’s optimized TPU capabilities directly into vLLM, developers can immediately run production workloads on TPUs with the high throughput, low latency, and memory efficiency they expect from our stack," said Simon Mo, CEO and co-founder of Inferact.
为促进持续的开源贡献,vLLM 还将推出首个社区贡献 TPU 算力池,为积极参与两个引擎开发的开发者免费提供 TPU 算力。请在此申请。To foster continuous open-source contribution, vLLM is also launching a first-of-its-kind community contribution TPU capacity pool to provide free TPU compute to developers who actively contribute to both engines. Request here.
路线图Roadmap
随着模型规模不断增大,我们正在升级这两种引擎,以提升大规模 TPU 推理部署的性能和效率。面向生产级智能体工作负载的主要功能现已得到支持。更多性能优化和高级功能正在开发中——欢迎提交 GitHub issue,申请协作支持!As the models get larger, we are evolving both engines to make large-scale TPU inference deployments performant and efficient. The major features for production-grade agentic workloads are already supported. More performance optimizations and advanced features are under way – open a Github issue to request collaboration on yours!
这两款引擎均支持大多数热门开源模型系列,包括 Qwen、DeepSeek、GLM、Kimi、GPT-OSS、Gemma4 等。团队正致力于让更多模型能够在 TPU 上原生运行,包括在新开源模型发布首日便提供支持。我们还希望为社区提供相应工具和 TPU 算力,让大家能够针对自己的工作负载和配置调优性能。Across the two engines, most of the popular open model families are supported, including Qwen, DeepSeek, GLM, Kimi, GPT-OSS, Gemma4 and more. The team is investing in making more models available to run natively on TPUs, including Day 0 support for new open models. We also want to give the community the tools and TPU capacity to tune performance for their own workloads and configurations.
关于我们的 JAX 运行时的一点说明A note on our JAX runtimes
vLLM-Jax 和 SGLang-JAX 仍是重要项目。两者都将继续维护和提供支持,并向社区开放使用和改编。vLLM-Jax, and SGLang-JAX remain important projects. Both will continue to be maintained, supported, and available for the community to use and adapt.
入门Getting started
立即开始:Get started today:
- 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
相关公告:Related announcements:
- Inferact 公告
- Radixark 公告