Amazon announced that SGLang-Neuron is publicly available for serving AI models on AWS Trainium chips, according to a company blog post. The software, released by RadixArk in collaboration with Annapurna Labs, can be accessed through an AWS Deep Learning Container.
SGLang-Neuron allows developers to serve models and generate reinforcement learning rollouts using the same SGLang interface and workflow they currently use, while running on Trainium hardware. AWS Trainium is Amazon’s portfolio of AI chips designed for training and inference.
The implementation uses an out-of-tree plugin design that connects SGLang with Trainium. Installing the plugin registers the Neuron platform through SGLang’s plugin entry points, and setting SGLANG_PLATFORM=neuron selects it when the server starts. SGLang core continues to provide the API server, request lifecycle, RadixAttention, and prefix-cache bookkeeping, while the plugin implements Trainium-specific models, attention, mixture of experts, and the graph runner.
SGLang-Neuron runs on native PyTorch, with TorchNeuron registering Trainium as a native PyTorch device. The system uses standard PyTorch operations for model execution and incorporates kernels written with the Neuron Kernel Interface for performance, compiling these into executable programs using torch.compile.
The software also supports prefill-decode disaggregation, where a prefill server computes the prompt key-value cache and transfers it to a decode server using NIXL over Elastic Fabric Adapter. The decode server then generates the remaining tokens. Each server uses SGLang’s scheduler and request lifecycle, while the plugin manages Trainium execution and key-value cache transfer.
Amazon says the software is available for use on Trainium3 EC2 instances with a pytorch-inference-sglang-neuronx container.