SGLang 推出 AWS Trainium 插件 Beta：让推理与 RL Rollout 跑上专用 AI 芯片

过去一年，SGLang 团队持续扩展其支持的硬件生态，同时也在新开放模型首日支持、以及 Miles 中的 RL rollout 能力上不断推进。如今，SGLang 宣布与 AWS 合作，将 SGLang 带到 AWS Trainium。

随着前沿 AI 逐步走向异构硬件架构，AWS Trainium 提供了一个值得关注的机会。Trainium 是 AWS 为 AI 工作负载专门打造的芯片，并以云服务形式大规模提供，目标是在高吞吐场景下提供更好的 price-performance。开放模型服务与 RL rollout 生成正是典型场景：二者往往需要持续运行，而每个 token 的成本会直接决定团队能够部署多少算力容量。

SGLang 团队表示，其目标是让 SGLang 成为在 Trainium 上运行这些工作负载的一等选择。此次在与 AWS 密切合作、并配合 TorchNeuron 发布的背景下，团队正式宣布推出基于新发布 TorchNeuron 的 SGLang Trainium plugin Beta，该 Beta 同时支持 Trainium 上的推理服务与 RL rollout。

从一开始就基于 TorchNeuron 构建

TorchNeuron 是 AWS 面向 Trainium 的原生 PyTorch 后端。SGLang 团队在 TorchNeuron 私有预览阶段就开始构建 SGLang 的 Neuron backend，而这一技术选择也影响了整个适配工作的路径。

由于 TorchNeuron 采用原生 PyTorch 方式，SGLang 的模型组件和核心抽象大体可以完整迁移到 Trainium。开发者在 Trainium 上服务模型时，依然使用与 GPU 上相同的接口，这意味着团队可以用较小改动在 GPU 与 Trainium 之间迁移工作负载。

在进一步优化性能方面，NKI（Neuron Kernel Interface） 允许开发者在 PyTorch 工作流中编写自定义 Trainium kernel。对于一个覆盖多种硬件平台的开源项目来说，这既提供了熟悉的编程模型，也在性能关键路径上保留了底层硬件控制能力。

Beta 版本包含哪些能力

此次 Beta 支持在 AWS Trainium2 与 Trainium3 上运行两类核心工作负载：LLM inference serving 与用于后训练的 RL rollout generation。首批覆盖范围包括：

• 模型：Qwen3（dense）、Qwen3-MoE-30B-A3B、DeepSeek-V2-Lite、GLM-4，以及 GLM-5.2/5.3。团队表示，更多大型前沿模型的支持正在推进中。
• 特性：Data Parallel、Tensor Parallel、Expert Parallel、Radix Cache、BF16 / FP8、Async scheduler。
• RL rollouts：SGLang 可在 Trainium 上作为异构 RL pipeline 的 rollout engine，支撑后训练流程。

SGLang 团队计划在 Beta 期间继续扩展模型与功能覆盖范围，并预计在今年晚些时候推进更广泛的可用性。

早期预览：异构 RL 工作流

除了推理服务，Trainium 还为 RL 后训练带来了新的部署配置。每一步 RL 通常包含两个不同阶段：其一是 rollout generation，这本质上是以 decoding 为主的推理负载；其二是 policy update，即训练负载。大多数团队会在同一个 GPU 集群上运行这两个阶段，但二者对硬件的需求并不相同。

SGLang Trainium plugin 允许这两个阶段运行在不同硬件上。在首次端到端运行中，SGLang 在 Trainium3 上为 Qwen3-MoE-30B-A3B 生成 rollouts；与此同时，SGLang 团队的 RL 后训练框架 Miles 使用 Megatron，在单个 H200 节点上训练 GRPO policy。Miles 将 Trainium 上的 SGLang 连接为外部 rollout engine，并完成了一次稳定的训练运行。

这种方式为团队规划 RL 算力提供了更高灵活性。不同阶段可以根据各自的 price-performance 选择硬件，Trainium 上的 rollout capacity 也可以独立于 GPU 训练资源进行扩展。

不过，SGLang 团队也强调，这仍是一个早期里程碑。Trainium 上的性能优化仍在继续，团队将在 11 月的 SGLang Summit 上分享后续进展。

如何加入 Beta

当前初始 Beta 仅面向少量团队开放，目的是让 SGLang 团队能够与每位参与者紧密合作。该 Beta 尤其适合以下两类团队：

• 正在或计划在 Trainium2 / Trainium3 上运行 LLM inference 的团队；
• 对异构 RL 感兴趣，希望在 Trainium 上生成 rollouts、并在 GPU 上进行训练的团队。

有意申请访问权限的团队，可以通过 官方表单 联系 SGLang 团队。