WORLDTECH亚马逊为AWS Trainium公开发布SGLang-Neuron
由 WORLDTECH 译自英文。事实、数字、名称和来源均与英文原文一致。 English
亚马逊宣布,SGLang-Neuron现已通过AWS深度学习容器向公众开放,开发者可借此在Trainium芯片上部署模型并生成强化学习的rollout数据。该插件由RadixArk与Annapurna Labs合作发布。
此次发布提供了什么
SGLang-Neuron是一个树外(out-of-tree)插件,用于将SGLang框架与AWS Trainium连接起来。亚马逊将Trainium描述为一系列AI芯片,能够以出色的经济性支持大规模训练和推理。用户可通过RadixArk与Annapurna Labs合作发布的AWS深度学习容器获取该插件。
插件安装后,会通过SGLang的插件入口点注册一个neuron平台。服务器启动时,设置SGLANG_PLATFORM=neuron即可选用该平台。
SGLang核心与插件的分工
SGLang核心仍负责API服务器、请求生命周期、RadixAttention以及前缀缓存的管理。插件则提供面向Trainium的模型、注意力机制、专家混合(MoE)组件和图运行器(graph runner)。
在这一层之下,SGLang-Neuron基于原生PyTorch运行。TorchNeuron将Trainium注册为原生PyTorch设备。模型执行依赖标准PyTorch运算,以及为提升性能而使用神经内核接口(Neuron Kernel Interface,NKI)编写的内核。这些运算和内核通过torch.compile编译为可执行程序,由图运行器在服务过程中进行准备和运行。
与现有SGLang工作流的衔接
在Trainium上运行SGLang保留了开发者已经习惯的工作流程。模型通过SGLang的服务器接口启动,端点使用相同的生成API进行调用。
亚马逊还介绍了一种分离式(disaggregated)部署方式:预填充(prefill)服务器构建提示词的KV缓存,并通过NIXL和弹性结构适配器(Elastic Fabric Adapter,EFA)将其发送至解码(decode)服务器,由后者生成其余的token。