面向 TPU 的全新 vLLM 与 SGLang 推理引擎正式发布：原生支持 PyTorch，加速生产级 Agent 与大模型服务

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

背景

多年来，TPU 一直以高性能和高效率为大规模 AI 工作负载提供算力支持。随着 Agentic（智能体）与多模态工作负载的兴起，高效模型服务比以往任何时候都更加关键。尽管 vLLM 和 SGLang 此前已通过基于 JAX 的后端在 TPU 上运行，但我们一直收到开发者的反馈，希望能在 PyTorch 中自定义和调试他们的服务堆栈。

因此，我们全力推进了 TorchTPU 的研发，使 PyTorch 能够原生运行在 TPU 上。随后，我们在此基础上构建了全新的 vLLM 和 SGLang 引擎。

我们的目标是让 PyTorch 社区能够利用开发者已经熟悉的工具，将模型迁移至 TPU 并实现极致性能。这些引擎配备了经过优化的生产级服务特性以及自定义 Pallas 内核（kernels）。

推理技术栈

实现大规模模型服务需要跨越技术栈的每一层进行协同设计——从芯片、编译器到编排系统。今天发布的全新推理引擎正处于这个集成的、开放的 TPU 生态系统的核心位置。

引擎概览

TorchTPU 将 TPU 变成了一个原生 PyTorch 设备。张量在 tpu 设备上直接作为标准 PyTorch 张量进行管理，无需任何包装类型；代码既支持 eager 模式直接运行，也可以利用 XLA 通过 torch.compile 进行编译；集合通信则统一经由 torch.distributed 执行。

由于每个设备都运行各自独立的 PyTorch 进程，因此两个引擎都完整保留了各自框架原有的 worker 模型、调度器、KV 缓存管理器和 API 服务器等，仅在底层增加了 TPU 实现。

大多数模型直接使用 vLLM 或 SGLang 上游的标准 PyTorch 模型定义，仅在内核与网络层实现中加入 TPU 专有优化。用 Pallas 编写的自定义内核可以直接作为编译图内的自定义算子从 PyTorch 中调用。这种平台插件式设计最大程度地复用了各框架既有成果，同时让 TPU 专属优化能够按自身的节奏独立迭代。

工作负载支持

两款引擎均面向全套生产级工作负载而构建。

以下是支持的功能列表：

• Agentic（智能体）与多轮对话： 自动前缀缓存（Automatic Prefix Caching），包括对混合模型的支持；通过 TPU Sync 将 KV 缓存卸载至主机内存及托管 Lustre（Managed Lustre）；工具调用（tool calling）与结构化输出（structured output）
• 交互式与低延迟： 采用 MTP、EAGLE-3、DFlash、DSpark 及 n-gram 草稿模型的投机解码（Speculative decoding）
• 高吞吐批处理： 结合专家并行（expert parallelism）的数据并行注意力机制
• 长上下文： 用于 Prefill 和 Decode 阶段的上下文并行（Context parallelism）；FP8 KV 缓存；主机 KV 缓存卸载
• 分离式服务（Disaggregated Serving）： 通过 TPU Sync 在不同主机间分离 Prefill 与 Decode 阶段，包括跨不同并行拓扑架构的处理
• 多模态： 支持图像与视频输入的视觉-语言模型（VLM）；嵌入（embedding）模型
• 超大模型： 混合与稀疏注意力架构；跨主机的张量并行、数据并行、专家并行、流水线并行及上下文并行；支持 FP8、FP4（MXFP4、NVFP4）、INT8 和 INT4 检查点（checkpoints）

特别是 Agentic 工作负载，往往需要同时具备上述多项特性。由于 Agent 会反复发送持续扩展的多轮对话内容，高效提供服务的关键在于查找已缓存的前缀而非重复计算。两个引擎都会在 HBM 中保留可复用的前缀；若开启主机卸载功能，还会自动将较旧的前缀迁移至主机内存作为二级存储缓存。

利用 TPU Sync 实现 KV 缓存传输

上述两项功能需要将 KV 缓存移离生成它的物理设备。分离式服务将 Prefill 和 Decode 分别运行在不同的 TPU 主机上，因此每个请求的 KV 缓存必须从 Prefill 主机传输到 Decode 主机。分层 KV 卸载则将 HBM 无法容纳的多余 KV 缓存转存至主机的 DRAM 以及通过 Google Cloud Managed Lustre 挂载的二级存储中，并在后续请求复用该前缀时重新加载。无论哪种方式，传输过程都必须与计算重叠（overlap），否则将直接增加延迟。

两款引擎均采用 TPU Sync 来解决这一问题。它是一个用于 TPU 上 KV 缓存传输与管理的开源 C++ 库。TPU Sync 可以直接读取 PyTorch 张量底层的设备缓冲区，因此传输过程无需在主机上进行额外的暂存复制；同时，它在 Python 外部独立的线程中执行传输，Prefill 与 Decode 运算得以不受阻塞地继续进行。vLLM-TorchTPU 通过 vLLM 的 KV 连接器（KV connector）接口将其暴露；SGLang-TorchTPU 则将其用作分离式传输后端以及 HiCache 的 I/O 后端。

生产落地现状与公开基准测试

这两款引擎目前已在预览版客户的生产环境中大规模运行。为了提供透明、可复现的性能佐证，我们通过社区标准的 SemiAnalysis InferenceX 项目公开发布了基准测试结果。

基于我们此前在 Ironwood TPU（TPU 7x）上提交的 初步基准数据（运行 Qwen 3.5），我们的工程团队进一步对技术栈进行了优化，将整体每百万 token 成本压降至行业领先的 2.01 美分。除了这些性能改进之外，我们还在 AgentX 上发布了最新的 SOTA（业界顶尖）性能结果——AgentX 是该项目专门用于评估 Agentic 长上下文和多轮对话工作负载的测试赛道。Ironwood TPU 在大部分帕累托前沿（Pareto frontier）上均保持领先。随着我们将更多模型和配置引入排行榜，我们期待持续更新我们的公开性能数据档案。

我们并非闭门造车构建这些后端。通过全面拥抱 PyTorch 并与开源项目维护者直接协作，我们确保这些引擎始终与上游保持同步并坚持社区优先：

SGLang-TorchTPU 由我们与 RadixArk（由 SGLang 创始团队创立）联合开发

这一合作为 TPU 带来了 SGLang 出色的技术团队，打造出的引擎现已在 Meta 的生产环境中稳定运行，并在严苛的实际真实场景中经受住了大规模考验。

“SGLang-TorchTPU 如今已在 Meta 的生产环境中运行，这一成果源于 Meta、Google 和 SGLang 团队之间的紧密三方协作，共同打造出了生产就绪的技术栈。它为我们提供了一个极其成熟、原生支持 PyTorch 的模型服务堆栈，能够毫不费力地集成到我们现有的 AI 基础设施中，并具备我们在 TPU 上高效运行前沿工作负载所需的性能、可靠性与灵活性。” Meta 的 PyTorch-TPU 负责人 Zongwei Zhou 表示。

“在 RadixArk，随着我们在多种异构芯片上推广 SGLang 支持，TorchTPU 极大地提升了我们的开发速度。它为我们提供了一条无痛获得原生 TPU 加速能力的路径，直接将巨大的效率提升和灵活的算力选择传递给我们的客户。” RadixArk 首席技术官兼联合创始人 Banghua Zhu 表示。

vLLM-TorchTPU 由我们与 Inferact（由 vLLM 创始团队创立）联合开发

我们共同为 2.8T（2.8万亿）参数的 Kimi K3 模型集成了高性能 Pallas 巨型内核（megakernel），该内核将整个解码（decode）步骤整合为单一内核运行，在 Ironwood TPU 上为单个用户提供超过每秒 700 tokens 的推理吞吐。基于此项工作，Inferact 最近主导了基于 vLLM-TorchTPU 堆栈的 Kimi K3 首个初步 AgentX 评测提交。这一初始里程碑不仅展示了在超高交互性工作负载下卓越的性价比（TCO），也证明了 TPU 上的 vLLM 与 PyTorch 生态在适配全新前沿开源模型方面的成熟度正在迅速跃升。后续提交还将通过更大规模的 Prefill/Decode 分离式架构进一步释放更高的计算效率。

“Inferact 的使命始终是让整个 AI 生态系统都能获得高性能、可靠且高性价比的推理能力。我们非常高兴能参与 TorchTPU 的发布，并为 TPU 带来对 vLLM 的一流原生支持。通过将 PyTorch 针对 TPU 优化的能力直接桥接到 vLLM 中，开发者可以立即在 TPU 上运行生产工作负载，同时享受到符合预期的极高吞吐量、低延迟和高内存效率。” Inferact 首席执行官兼联合创始人 Simon Mo 表示。

为了促进开源社区的持续贡献，vLLM 还推出了首创的“社区贡献 TPU 算力池”，为积极向这两个引擎贡献代码的开发者提供免费 TPU 算力支持。开发者可在此处申请。

路线图

随着模型体量不断扩大，我们正在持续推进这两款引擎的发展，以确保大规模 TPU 推理部署的高性能与高效能。目前针对生产级 Agentic 工作负载的核心特性已全部就绪。更多性能优化与高级特性正在紧锣密鼓地研发中——欢迎在 GitHub 上提交 issue 提出合作需求！

横跨这两大引擎，目前已支持大多数主流开源模型家族，包括 Qwen、DeepSeek、GLM、Kimi、GPT-OSS、Gemma4 等。团队正持续投入让更多模型能够在 TPU 上原生运行，包括对最新开源模型的“首日支持（Day 0 support）”。我们还希望为社区提供充足的工具与 TPU 算力，以便开发者针对自身业务和特定配置进行性能调优。

关于 JAX 运行时的说明

vLLM-Jax 和 SGLang-JAX 依然是非常重要的项目。两者将继续获得维护、支持，并供社区使用和改造。

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

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.

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.

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

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 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.

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.

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:

• 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
• Interactive & Low Latency: Speculative decoding with MTP, EAGLE-3, DFlash, DSpark and n-gram drafters
• High-throughput Batch: Data-parallel attention with expert parallelism
• Long Context: Context parallelism for prefill and decode; FP8 KV cache; host KV offloading
• 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
• 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

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.

Moving the KV Cache with TPU Sync

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.

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

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.

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 to an industry leading 2.01 cents. 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.

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-developedwith RadixArk (founded by the creators of SGLang)

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 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.

"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 was co-developedwith Inferact (founded by the creators of vLLM)

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.

"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.

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

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!

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.

A note on our JAX runtimes

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:

Related announcements: