AWS Machine Learning BlogT1·算力与基建

SGLang-Neuron 开源支持在 AWS Trainium 上部署前沿模型

作者：Alex Nails原标题：Deploy frontier AI models on AWS Trainium with SGLang-Neuron

62AI 研判分

AWS 宣布 SGLang-Neuron 正式公开可用，可通过 AWS 深度学习容器获取，由 RadixArk 与 Annapurna Labs 合作发布。该插件以树外插件形式将 SGLang 接入 Trainium，底层基于原生 PyTorch，用户沿用现有 SGLang 服务接口与生成 API 即可部署模型并生成强化学习 rollout。

关联知识实体

术语强化学习 rollout 公司Annapurna Labs 工具SGLang-Neuron 工具SGLang 公司亚马逊云科技

报道正文

本文目录3 个章节

今天，Amazon 宣布 SGLang-Neuron 已公开可用，并可通过 AWS Deep Learning Container 访问，该容器由 RadixArk 与 Annapurna Labs 合作发布。借助 SGLang-Neuron，你可以在 AWS Trainium 上使用 SGLang 提供模型服务并生成强化学习（RL）rollout，且沿用你当前使用的同一套 SGLang 接口。Trainium 是 AWS AI 芯片产品组合，可为大规模 AI 训练和推理提供出色的经济性。如今，在 Trainium 上运行 SGLang 遵循你已熟悉的 SGLang 工作流：通过 SGLang 的服务器接口启动模型，并使用相同的生成 API 调用端点。在该接口之下，SGLang-Neuron 运行在原生 PyTorch 上。凭借其树外插件设计，Trainium 支持与 SGLang 核心可独立开发并更快迭代。

SGLang 如何在 Trainium 上运行

SGLang-Neuron 是一个树外插件，用于将 SGLang 与 AWS Trainium 连接起来。安装该插件会通过 SGLang 的插件入口点注册 neuron 平台，设置 SGLANG_PLATFORM=neuron 即可在服务器启动时选择它。SGLang 核心继续提供 API 服务器、请求生命周期、RadixAttention 和前缀缓存簿记。该插件实现了 Trainium 特定的模型、注意力、混合专家（MoE）以及图运行器。

图 1：AWS Trainium 上的 SGLang-Neuron

TorchNeuron 将 Trainium 注册为原生 PyTorch 设备。SGLang-Neuron 在模型执行路径中使用标准 PyTorch 操作，并集成了使用 Neuron Kernel Interface（NKI）编写的内核以提升性能。它使用 torch.compile 将这些操作和内核编译为可执行程序，图运行器在服务期间准备并执行这些程序。

SGLang-Neuron 同样支持 PD 分离。在该配置中，prefill 服务器计算提示词的 KV 缓存，并通过 Elastic Fabric Adapter (EFA) 上的 NIXL 将其传输到 decode 服务器。随后，decode 服务器生成剩余的 token。每个服务器仍然使用 SGLang 的调度器和请求生命周期，而插件负责管理 Trainium 执行和 KV 缓存传输路径。

使用 SGLang-Neuron 在 AWS Trainium 上部署模型

要完成本演练，你需要使用一个搭载 pytorch-inference-sglang-neuronx 容器的 Trainium3 Amazon Elastic Compute Cloud (Amazon EC2) 实例。步骤如下：

启动容器。

BASH

docker run --rm -it \
--privileged --ipc=host \
-v /path/to/work:/work \
public.ecr.aws/neuron/pytorch-inference-sglang-neuron:v0.1.0-neuron-py314-sdk-2.33.0-ubuntu26.04 \
/bin/bash

在容器内启动 SGLang 模型服务器。

BASH

SGLANG_PLATFORM=neuron \
python -m sglang.launch_server \
--model-path openai/gpt-oss-120b \
--tp-size 8 \
--ep-size 8 \
--kv-cache-dtype fp8_e4m3 \
--context-length 10240 \
--chunked-prefill-size 8192 \
--max-running-requests 64 \
--max-total-tokens 655360 \
--mem-fraction-static 0.8 \
--host 127.0.0.1 \
--port 30000

向模型服务器发送请求。

BASH

curl -s http://127.0.0.1:30000/v1/chat/completions \
-H 'Content-Type: application/json' \
-d '{
"model": "openai/gpt-oss-120b",
"messages": [
{"role": "user", "content": "Explain why the sky appears blue."}
],
"temperature": 0,
"max_tokens": 128
}'

如需更多信息，请参阅 SGLang-Neuron 入门指南以了解设置说明，并浏览仓库文档以了解支持的模型、功能和经过验证的配置。

结论

SGLang-Neuron 通过原生 PyTorch 将 SGLang 熟悉的 serving 工作流带到 AWS Trainium。你可以启动 SGLang-Neuron 容器，并按照快速入门来部署受支持的模型，从而开始使用。欢迎在你的工作负载中试用，通过项目仓库参与贡献，并分享反馈，以帮助塑造未来的模型与功能支持。