‹ 返回博客‹ 返回博客

目录

SGLang-TorchTPU：在 TPU 上实现 PyTorch 原生生产级服务RadixArk、Google 和 Meta2026 年 10 月 20 日

SGLang-TorchTPU

摘要

  • 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。
  • 该插件将 SGLang 的完整功能集带到 TPU，可满足智能体工作负载的需求。
  • 它已在 Meta 的大规模生产环境中得到验证。
  • 在 TPU 上使用 SGLang-TorchTPU 和 Miles 中的 TorchTitan 进行强化学习训练可以正常工作。
  • 贡献者可申请 RadixArk 的 社区 TPU 算力池.

引言

今天我们推出 SGLang-TorchTPU，这是一个原生 PyTorch 后端，将 SGLang 的完整服务栈带到 Google Cloud TPU。SGLang 的服务引擎和模型实现均基于 PyTorch 构建，借助 TorchTPU，同一套代码库可同时服务于 GPU 和 TPU，因此你可以选择最适合每种工作负载及其部署经济性的硬件。

TPU 专属组件仍隔离在插件中。因此，添加到 SGLang 的新模型和新功能只需一次实现即可在 TPU 上可用。

工作原理

SGLang core (scheduler, RadixAttention, API server) on top of the sglang-tpu plugin (Pallas ragged paged attention, paged KV cache pools, Pallas grouped matmul MoE, a graph runner that compiles once per bucket), on TorchTPU (one PyTorch process per TPU core, torch.compile to XLA), on Cloud TPU v7x.

图 1：sglang-tpu 插件如何融入 SGLang。

SGLang-TorchTPU 是一个树外插件。安装 sglang-tpu 包会通过 sglang.srt.platforms 入口点向 SGLang 注册 TPU 平台，而 SGLANG_PLATFORM=tpu_oot 会选择它。调度器、RadixAttention 和 API 服务器均为 SGLang 自有。插件补齐了 SGLang 本就要求各硬件平台提供的部分：注意力后端、KV 缓存池、MoE 内核和图运行器。

每个 TPU 核心通过 TorchTPU 运行自己的 PyTorch 进程，正如每个 GPU 在 CUDA 上所做的那样。模型在 torch.compile 下运行并下沉到 XLA。这里没有图捕获：运行器改为将每次前向填充到固定的一组 token 数桶（用于预填充）和批大小桶（用于解码）。每个桶编译一次，后续步骤按形状分派到已缓存的程序。

注意力和 MoE 运行在 Pallas 内核上，这些内核注册为 torch.library 自定义算子，Dynamo 在编译图中将其视为不透明调用：在分页 KV 缓存上执行不规则分页注意力，以及对排序后的专家分组执行分组矩阵乘法以实现 MoE。

当 KV 需要离开设备时——无论是将 KV 缓存卸载到主机内存，还是跨主机进行预填充-解码分离——它可以通过 TPU Sync（Google 的开源 TPU DMA 库）传输。卸载通过直接 PJRT DMA 在 HBM 和固定主机池之间复制整个 KV 页。对于分离，传输采用消费者拉取模式：每个分块预填充步骤提交其页时，预填充端会通告这些页，解码端在下一个分块计算的同时拉取该分块，因此预填充完成后只需移动最后一个分块。解码端已从其自身前缀缓存中持有的页永远不会被发送。当预填充和解码使用不同的张量并行规模时，TPU Sync 会在传输中重新分片：每个预填充 rank 将其 KV 头切片按步长推入解码 rank 的页布局。

KV cache transfer with TPU Sync: decode pulls each prefill chunk while the next computes, so only the last chunk waits; TP 16 to TP 8 reshards strided KV heads; offload moves KV pages between TPU HBM and a host pool by DMA.

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

功能与模型覆盖

以下是发布时在 TPU 上可用的 SGLang 服务功能：

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

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

模型类型
Muse Spark OpenLLM
GLM-5.3 Flash、GLM-5.3LLM
Kimi K3LLM
DeepSeek V4LLM
Ling 3.1 Flash、Ling 3.0 系列LLM
Qwen3.8LLM
Gemma 4LLM
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 开始，并且我们正努力在 2027 年第一季度前实现新开源模型的 Day 0 TPU 支持。

完整列表（含支持层级和每个模型的说明）见 SUPPORT.md.

可靠性

每个标记为支持的模型和功能都在真实 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 缓存卸载、DP 注意力、多主机执行，以及带 KV 重分片的 PD 分离
  • 功能组合，例如投机解码与 PD 分离组合、前缀缓存与 KV 缓存卸载组合
  • 多主机部署上的智能体工具调用工作负载

每晚运行都会对每个受支持模型重复这些检查，因此支持矩阵始终保持最新。

在 Miles 上使用 TPU 进行训练

服务只是循环的一半。要用强化学习（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 倍。

我们在单台配备 4 个芯片的 TPU v7x 主机上，使用 GRPO 在 DAPO-Math-17k 上训练了 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 训练处于预览阶段。这一小规模实验展示了 TPU 上的端到端 RL 工作流，其早期学习行为与 GPU 基线相似。接下来，我们计划开展以下工作：

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

由 Google 和 Meta 共同构建

SGLang-TorchTPU 由 Google、Meta 和 RadixArk 的 SGLang 团队联合开发。

“随着智能体工作负载增长，服务效率比以往任何时候都更重要，而开发者希望在不离开 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 负责人

开始使用

在 TPU v7x 主机上，启动预构建镜像、启动模型并发送请求：

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

# Inside the container, launch GPT-OSS-20B
python -m benchmark.tpu.launch_server \
  --config recipe/configs/gpt_oss_20b/server.py

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

要使用 pip 或从源码安装，请参阅 README。 TPU 指南 涵盖部署， recipes 提供可直接运行的服务器配置， SUPPORT.md 列出了受支持的模型和功能。

社区

RadixArk 正在为对 SGLang-TorchTPU 做出贡献的开发者提供免费 TPU 算力。 在此申请社区 TPU 算力池.

如有问题，请加入 SGLang Slack上的 #torchtpu。要请求模型或提出功能建议， 请在 GitHub 上提交 issue.

Google、Meta 和 RadixArk 将在 SGLang Summit 2026上分享更多内容，该活动于 11 月 12 日至 13 日在旧金山 Fort Mason 举行。