AWS Trainium 上能用 SGLang 了吗？SGLang-Neuron 发布要点与上手路径

SGLang-Neuron 已通过 AWS 深度学习容器公开可用，它让 SGLang 的服务流程跑在 AWS Trainium 上，底层是原生 PyTorch。已经在用 SGLang、又想评估 Trainium 推理路径的开发者，可以沿用原有启动方式和生成接口；如果现有 GPU 方案稳定且没有算力成本压力，短期不必迁移。

这次发布的是什么

据 AWS 机器学习博客，Amazon 宣布 SGLang-Neuron 公开可用，可通过 AWS Deep Learning Container 获取，由 RadixArk 与 Annapurna Labs 合作发布。来源标注的发布时间为 2026 年 11 月 12 日（UTC），文中“Today”即指公告当日。原文链接可在 AWS 机器学习博客找到，标题为《Deploy frontier AI models on AWS Trainium with SGLang-Neuron》，读者可据此核对发布时间与细节。

需要区分两个概念。SGLang 是开源推理框架，负责 API 服务器、请求生命周期、RadixAttention（一种通过复用前缀缓存降低重复计算的开源注意力调度机制）和前缀缓存记账。SGLang-Neuron 则是连接 SGLang 与 Trainium 的插件，负责把 Trainium 特有的模型实现、注意力、混合专家（MoE）和图执行器接进来。Trainium 是 AWS 自研的 AI 芯片系列，面向训练和推理场景。

插件如何选择 Trainium 平台

理解机制有助于排查启动问题。SGLang-Neuron 采用 out-of-tree 插件设计，即不把 Trainium 支持塞进 SGLang 主干，而是独立开发和迭代。安装插件后，它通过 SGLang 的插件入口注册 neuron 平台；启动服务时设置环境变量 SGLANG_PLATFORM=neuron，SGLang 就会选中该平台。

在 Trainium 一侧，TorchNeuron 把 Trainium 注册为原生 PyTorch 设备。模型执行路径使用标准 PyTorch 算子，性能关键部分用 Neuron Kernel Interface（NKI）编写内核，再经 torch.compile 编译成可执行程序，由图执行器在服务期调度。这个分层意味着：SGLang 的接口和调度逻辑不变，变化发生在设备与算子层。

PD 分离与 KV 缓存传输

SGLang-Neuron 支持 PD 分离（prefill 与 decode 分离）。在这种配置下，prefill 服务器先计算提示词的 KV 缓存，再通过 NIXL over Elastic Fabric Adapter（EFA）把缓存传给 decode 服务器，由后者继续生成剩余 token。两个服务器各自仍使用 SGLang 的调度器和请求生命周期，插件负责 Trainium 执行与 KV 缓存传输。

PD 分离适合长提示、高并发或希望把预填充与解码资源分开规划的场景，但它依赖 EFA 网络和跨节点缓存传输，部署复杂度高于单机服务。建议先跑通基础服务模式，再评估是否拆分。

三步跑通示例服务

以下流程来自 AWS 博客示例，适用于 Trainium3 EC2 实例与 pytorch-inference-sglang-neuronx 容器。镜像标签、SDK 版本和实例规格请以官方镜像仓库与实例文档为准，示例中的路径 /path/to/work 需替换为本机目录。

第一步：启动容器

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

--privileged 与 --ipc=host 用于让容器访问设备与共享内存，-v 挂载工作目录以便存放模型或日志。生产环境应评估是否必须使用特权模式，并按最小权限原则调整。

第二步：拉起 SGLang 模型服务

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

关键参数含义：--tp-size 是张量并行度，--ep-size 是专家并行度，--kv-cache-dtype 指定 KV 缓存数据类型，--context-length 限制上下文长度，--max-running-requests 控制并发请求数，--mem-fraction-static 决定静态显存占比。--chunked-prefill-size 控制分块预填充的块大小，--max-total-tokens 是 KV 缓存与请求可用的总 token 预算，具体语义与调整方式以 SGLang 官方文档为准。以上数值均为示例值，需按实际模型和芯片数量调整。

第三步：发送请求验证

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
  }'

能返回符合 OpenAI 兼容格式的补全结果，说明服务链路已通。若启动失败，优先检查 SGLANG_PLATFORM 是否生效、模型是否在支持清单内、并行度与芯片数量是否匹配。

对部署决策的实际影响

这次发布的价值主要在接口一致性，而非性能数字。来源没有给出吞吐、延迟或成本数据，因此无法据此判断 Trainium 是否比现有 GPU 方案更划算。对已经在 SGLang 上构建服务的团队，迁移工作量集中在更换平台选择和容器镜像；对做强化学习 rollout 的团队，同一套接口兼顾在线服务与 rollout 生成，可减少工具链分裂。

反过来，如果现有 GPU 推理已经满足业务，且没有明确的成本或供给压力，短期没有必须迁移的理由。Trainium 生态的支持模型、已验证配置和功能覆盖仍需逐项核对。

部署前需要自行核实的清单

• 目标模型是否在 SGLang-Neuron 官方仓库的支持模型与已验证配置列表中。
• Trainium3 实例的可用区域、规格与计费方式，以 AWS 官方文档为准。
• 容器镜像标签与 Neuron SDK 版本是否与实例环境匹配。
• 是否需要 PD 分离；若需要，EFA 网络配置与跨节点传输是否已规划。
• 在自身模型和真实负载下记录吞吐与延迟，再与现有方案比较。

操作提醒：示例中的 --privileged 会放宽容器权限，生产部署前应确认能否用更细粒度的设备映射替代；修改并行度或缓存参数后，先小流量验证再扩大并发。

常见问题

SGLang-Neuron 需要重写现有 SGLang 代码吗？

按来源描述，服务启动方式和生成 API 保持一致，主要变化是设置 SGLANG_PLATFORM=neuron 并使用对应容器。模型侧是否兼容取决于支持清单，不能一概而论。

没有 Trainium 实例能先试用吗？

示例流程要求 Trainium3 EC2 实例。没有实例时无法完成设备侧验证，可先阅读官方仓库文档了解支持范围。

PD 分离是必须开启的吗？

不是。它是可选配置，适合需要分开规划预填充与解码资源的场景。单机基础模式更容易先跑通。

能直接拿来做强化学习 rollout 吗？

来源称可用同一 SGLang 接口生成 RL rollout，但具体框架集成、采样配置和验证方式需参考项目仓库文档，本文不展开未经验证的细节。

如何判断是否值得从 GPU 迁移？

先在目标模型和真实请求分布下测吞吐、延迟与单位成本，再与现有方案对比。缺少自身基准数据时，任何迁移结论都只是推测。

下一步建议

先在单个 Trainium3 实例上按三步流程跑通示例模型，确认容器、平台选择和请求链路无误；再替换为业务模型，记录并发与上下文长度变化下的表现；最后才考虑 PD 分离和多实例扩展。对关注海外云服务器与部署成本的读者，可结合自身机房位置、访问来源和预算，在 UU云 的相关部署文章与 AWS 官方产品页之间交叉核对后再做决定。