RadixArk 联合 Google、Meta 开源 SGLang-TorchTPU，将 SGLang 推理栈带到 TPU

SGLang-TorchTPU

TL;DR

SGLang-TorchTPU 今日开源，由 RadixArk 的 SGLang 团队与 Google 和 Meta 合作构建。
首批支持的模型包括：Muse Spark Open、GLM-5.3 Flash、GLM-5.3、Kimi K3、DeepSeek V4、Ling 3.1 Flash、Qwen3.8、Gemma 4 和 MiniMax H3。
已可用于生产环境的 agentic 工作负载和大模型服务。
已在 Meta 的生产环境中运行。
可作为 TPU 上 RL 训练的 rollout 引擎，通过 Miles 与 TorchTitan 配合使用。

引言

今天我们推出 SGLang-TorchTPU，这是一个原生 PyTorch 后端，将 SGLang 的完整服务栈带到 Google Cloud TPU。它让同一个 SGLang 引擎既能在 GPU 上服务模型，也能在 TPU 上服务模型，因此你可以选择最适合每项工作负载及其部署经济性的硬件。

为什么选择 TorchTPU

如今大多数模型都用 PyTorch 编写和提供服务，服务栈也构建在其生态之上，从 torch.compile 到 torch.distributed。我们希望 TPU 融入同样的工作流。

TorchTPU 让 TPU 成为原生 PyTorch 设备，可以 eager 执行，也可以通过 XLA 编译。SGLang-TorchTPU 构建于其上，在针对 TPU 调优的 Pallas 内核之上，将 SGLang 的服务特性带到 TPU。

Google 构建了 TorchTPU，并与我们合作开发 TPU 后端本身，从 Pallas 内核到 TPU Sync 传输路径。RadixArk 的 SGLang 团队构建了服务运行时、其 TPU 集成以及大部分模型支持。Meta 工程师贡献了编译和分离式服务方面的工作，并带来了生产工作负载，从一开始就塑造了这个引擎。

“随着 agentic 工作负载增长，服务效率比以往任何时候都更重要，开发者希望在不离开 PyTorch 的情况下在 TPU 上实现这一点。与 RadixArk 和 Meta 共同开发 SGLang-TorchTPU，为他们提供了面向热门开源模型的生产级引擎，并具备 Cloud TPU 的成本效率。”

Qi Zhou，Google AI 推理技术负责人

“SGLang-TorchTPU 目前已在 Meta 的生产环境中运行，实现这一目标离不开 Meta、Google 和 SGLang 团队之间的紧密三方合作，使整个技术栈达到生产就绪。它为我们提供了一个非常成熟、PyTorch 原生的服务栈，可以轻松集成到我们现有的 AI 基础设施中，并具备在 TPU 上高效运行前沿工作负载所需的性能、可靠性和灵活性。”

Zongwei Zhou，Meta PyTorch-TPU 负责人

工作原理

The SGLang-TorchTPU stack: SGLang core (API server, router, scheduler, RadixAttention, model code) on top of the sglang-tpu plugin (TPU layers, Pallas kernels, KV cache pools, TPU Sync), on top of TorchTPU (tpu device, torch.compile + XLA, torch.distributed), on Cloud TPU v7x.

图 1：SGLang-TorchTPU 技术栈。

每个设备运行自己的 PyTorch 进程，因此 SGLang 的调度器、RadixAttention 和 API 服务器无需改动即可运行。TPU 相关工作位于它们之下。KV 缓存存放在为 Pallas ragged paged attention 布局的分页池中，默认使用 FP8 KV，attention 和 MoE grouped matmul 的 Pallas 内核作为 PyTorch 自定义算子，在 torch.compile 图中被调用。所有这些都以 sglang-tpu 插件形式发布，当你传入 --device tpu 时 SGLang 会加载它。

TPU 设备与主机内存之间的 KV 缓存传输可以运行在 TPU Sync 上，这是 Google 的开源传输库。对于 KV 缓存卸载，它将 KV 页移动到 DMA 固定的主机内存池。对于 prefill-decode 分离，decode 侧会在每个 prefill 步骤提交时逐块从 prefill 拉取 KV，因此 prefill 结束时只剩最后一块需要移动。当两侧使用不同的张量并行规模时，KV 会在传输途中重新分片，例如从 16 个 prefill 分片到 8 个 decode 分片。

KV cache transfer with TPU Sync: in prefill-decode disaggregation, KV moves chunk by chunk from prefill (TP 16) to decode (TP 8); in KV cache offloading, KV is evicted from TPU HBM to host memory and reloaded.

图 2：使用 TPU Sync 进行 KV 缓存传输。

功能与模型覆盖

以下是 TPU 上首发可用的 SGLang 服务功能：

领域	支持的功能
核心服务	兼容 OpenAI 的 API；连续批处理；分块预填充；结构化输出与函数调用（实验性）
智能体与多轮对话	RadixAttention 前缀缓存，包括混合滑动窗口与线性注意力模型；KV 缓存卸载至主机内存；每晚在智能体工具调用负载上测试
交互式、低延迟	使用 EAGLE-3、DFlash 和 MTP 草稿模型的投机解码；DFlash2 与 DSPark（实验性）
高吞吐批处理	数据并行注意力与专家并行
长上下文	FP8 KV 缓存；预填充与解码的上下文并行（实验性）
分离式服务	预填充与解码部署在不同主机上，支持不同张量并行规模之间的 KV 重分片
多模态	支持图像输入的视觉语言模型：Qwen3-VL、Qwen3.5-VL、Kimi MoonViT 和 Gemma 4 vision
超大规模模型	跨主机的张量并行、专家并行与数据并行注意力；FP8、MXFP4 和 INT4 检查点；流水线并行与 NVFP4 MoE（实验性）

我们一直在努力将最热门的开源模型在发布时带到 TPU 上。少数最新模型仍处于实验阶段，我们正在对其进行全面测试并将其加入 CI。

模型	类型
Muse Spark Open	LLM
GLM-5.3 Flash、GLM-5.3	LLM
Kimi K3	LLM
DeepSeek V4	LLM
Ling 3.1 Flash、Ling 3.0 系列	LLM
Qwen3.8	LLM
Gemma 4	LLM
MiniMax H3	扩散模型
Wan2.2	扩散模型

此外还支持：gpt-oss、Qwen3-Coder-480B、Qwen3.5、DeepSeek-V3.1 和 R1、Kimi-K2.6、MiMo-V2.5-Pro 以及 Qwen3-VL。

我们将持续添加新模型并优化已支持的模型，首先从 DeepSeek V4.1 Flash 和 Reflection Beam 开始。

包含支持层级和逐模型说明的完整列表见 SUPPORT.md。

可靠性

可靠性被设计进 SGLang-TorchTPU 的构建方式中。每个标记为支持的模型和功能都绑定到真实 TPU 硬件上的 CI 测试，支持矩阵由这些测试注册生成，因此我们的声明始终反映实际运行的内容。仍在验证中的功能会被标记为实验性，直到其测试通过。

SGLang-TorchTPU support tiers: Experimental moves to Supported with a nightly test, and Supported moves to Optimized with accuracy and performance tests, all on TPU v7x every PR and every night.

图 3：模型和功能如何在支持层级之间流转。

拉取请求会在单主机 TPU v7x 切片上运行 TPU 测试套件，结合每晚测试套件，CI 在单主机和多主机 TPU v7x 上覆盖 15 个模型，目前注册了 193 项 TPU 测试。它们检查：

内核正确性
在 GSM8K、GPQA、AIME、MMLU-Pro 和 MMMU 等标准评测上的端到端服务与模型准确率
投机解码
前缀缓存与 KV 缓存卸载
多主机执行与 PD 分离，包括 KV 重分片
多主机部署上的智能体工具调用负载

使用 TPU 在 Miles 上进行训练

服务只是闭环的一半。要用强化学习（RL）训练模型，训练器也必须保持同样的 GPU 或 TPU 选择，因此我们使用相同的 PyTorch 技术栈在 TPU 上进行了端到端 RL 训练。

训练器是 TorchTitan，PyTorch 的原生训练库，运行在 TorchTPU 上，无需 TPU 专用模型代码。与服务一样，每个 TPU 设备运行自己的 PyTorch 进程。Miles（RadixArk 的 RL 框架）将两侧连接起来。SGLang-TorchTPU 生成 rollout，TorchTitan 基于这些 rollout 进行训练。每一步或每隔几步之后，Miles 会将新权重从训练器同步到引擎。

RL training with Miles on one TPU v7x host: SGLang-TorchTPU generates rollouts on 2 chips, TorchTitan trains on 2 chips, and new weights move back to the engine through shared memory, disk or ICI (planned).

图 4：在单台 TPU v7x 主机上使用 Miles 进行 RL 训练。

Miles 可以通过三种方式将新权重从训练器移动到引擎：

共享内存（当前）：训练器和引擎位于同一主机上；我们的运行使用的就是这种方式。
磁盘：写入完整检查点，由引擎重新加载；它可跨主机工作，但速度最慢。
ICI（计划中）：通过 TPU 互连在设备间移动权重；我们的测试测得它比共享内存快约 50 倍。

我们在单台 TPU v7x 主机上用 4 个芯片，在 DAPO-Math-17k 上以 GRPO 训练 Qwen3-4B：一半用于训练器，一半用于 rollout。作为参照，我们在 6 块 H200 GPU 上运行了相同的配方和超参数。每个平台运行两个种子以进行一致性检查。

Qwen3-4B GRPO on DAPO-Math-17k: 6x H200 GPU vs. 4 TPU v7x chips

图 5：在 DAPO-Math-17k 上对 Qwen3-4B 进行 GRPO，6 块 H200 GPU 对比 4 个 TPU v7x 芯片。

TPU 上的奖励和响应长度曲线与 GPU 在种子间波动范围内一致，AIME-24 保持在评估噪声范围内。

当前的局限与未来计划：

Miles TPU 训练处于预览阶段。我们在小规模实验中展示了端到端 RL 工作流，以及与 GPU 基线相似的早期学习行为。

更多模型：更大的稠密模型和 MoE 模型，例如 Qwen3-30B-A3B，并支持路由重放。
多主机训练：扩展到单台 TPU v7x 主机之外。
更快的权重传输：通过 ICI 而非主机共享内存同步权重。
性能：缩小与 GPU 的步时差距。
更大规模验证：在更长运行和长时程智能体任务上进行一致性检查。
训练与推理一致性：训练与推理之间的对数概率按位完全相同。

开始使用

拉取镜像并在 TPU v7x 主机上启动容器：

docker pull lmsysorg/sglang:v0.1.0-tpu

docker run -it --rm --privileged --net=host --shm-size=64g \
  -e SGLANG_PLATFORM=tpu_oot \
  lmsysorg/sglang:v0.1.0-tpu

在容器内，使用单主机配方启动 GPT-OSS-20B，然后发送请求：

python -m benchmark.tpu.launch_server --config recipe/configs/gpt_oss_20b/server.py

curl http://localhost:30000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model": "gpt-oss-20b", "messages": [{"role": "user", "content": "Hello"}]}'

如需改用 pip 或从源码安装，请参阅 README。

GitHub：sgl-project/sglang-torchtpu
支持的模型和功能：SUPPORT.md，每一行对应的测试见 SUPPORT_MATRIX.md
Cookbook：cookbook.sglang.io

我们正努力在 2027 年第一季度前为新开放模型实现 Day 0 TPU 支持。

RadixArk 正在开放一个社区 TPU 算力池，为向 SGLang-TorchTPU 贡献的开发者提供免费算力。
如需提问，请加入 SGLang Slack 上的 #torchtpu。如需请求模型或提出功能建议，请在 GitHub 上提交 issue。
Google、Meta 和 RadixArk 将在 SGLang Summit 2026 上分享更多内容，该峰会于 11 月 12 日至 13 日在旧金山 Fort Mason 举行（sglang.io/summit）。