Article · AI translation
今天，Amazon 宣布 SGLang-Neuron 已公开提供，可通过由 RadixArk 与 Annapurna Labs 合作发布的 AWS Deep Learning Container 使用。借助 SGLang-Neuron，你可以在 AWS Trainium 上通过现有的 SGLang 接口使用 SGLang 提供模型服务并生成强化学习（RL）轨迹。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 分离。在这种配置下，预填充服务器计算提示词的 KV 缓存，并通过 Elastic Fabric Adapter（EFA）上的 NIXL 将其传输到解码服务器。解码服务器随后生成剩余的词元。每台服务器仍使用 SGLang 的调度器和请求生命周期，而插件负责管理 Trainium 执行及 KV 缓存传输路径。使用 SGLang-Neuron 在 AWS Trainium 上提供模型服务 按照本演练操作时，你需要使用一个运行 pytorch-inference-sglang-neuronx 容器的 Trainium3 Amazon Elastic Compute Cloud（Amazon EC2）实例。步骤如下：启动容器。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 模型服务器。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 向模型服务器发送请求。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 服务工作流程带到 AWS Trainium。你可以启动 SGLang-Neuron 容器，并按照快速入门指南提供受支持模型的服务。请使用你的工作负载试用，通过项目代码仓库贡献，并分享反馈，帮助确定未来的模型和功能支持。作者简介 Alex Nails Alex 是 RadixArk 的技术人员。他负责 MultiHardware，也是 SGLang 的核心开发者。Yihao Wang Yihao 是 RadixArk 的技术人员。他从事多模态／扩散模型和多硬件方面的工作。Lingyan Hao Lingyan 是 RadixArk 的产品经理。她负责 SGLang 的采用和多硬件支持。Srecharan Selvam Srecharan 是 Amazon 旗下 Annapurna Labs 的软件工程师，从事 AI 推理技术工作。Zeyu Wang Zeyu 是 Amazon 旗下 Annapurna Labs 的软件工程师，从事 AI 推理技术工作。Vinesh Ravuri Vinesh 是 Amazon 旗下 Annapurna Labs 的软件工程师，从事 AI 推理技术工作。Yuchen Jiang Yuchen 是 Amazon 旗下 Annapurna Labs 的软件工程师，从事 AI 推理技术工作。Pinak Panigrahi Pinak 在 Amazon 旗下 Annapurna Labs 从事 AWS Trainium 上 AI 模型的性能优化工作。Chi Wang Chi 是 Amazon 旗下 Annapurna Labs 的软件开发经理，从事 AI 推理技术工作。