Article · Original
Today, Amazon announced that SGLang-Neuron is publicly available and can be accessed through an AWS Deep Learning Container, released by RadixArk in collaboration with Annapurna Labs. With SGLang-Neuron, you can serve models and generate reinforcement learning (RL) rollouts with SGLang on AWS Trainium using the same SGLang interface you use today. Trainium is a portfolio of AWS AI chips that delivers strong economics for AI training and inference at scale. Now, running SGLang on Trainium follows the SGLang workflow you already know: you launch models through SGLang’s server interface and invoke the endpoint using the same generation APIs. Underneath that interface, SGLang-Neuron runs on native PyTorch. With its out-of-tree plugin design, Trainium support and SGLang core develop independently and iterate faster. How SGLang runs on Trainium SGLang-Neuron is an out-of-tree plugin that connects SGLang with AWS 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. The plugin implements the Trainium-specific models, attention, mixture of experts (MoE), and the graph runner. Figure 1: SGLang-Neuron on AWS Trainium TorchNeuron registers Trainium as a native PyTorch device. SGLang-Neuron uses standard PyTorch operations for the model execution path and incorporates kernels written with the Neuron Kernel Interface (NKI) for performance. It compiles these operations and kernels into executable programs using torch.compile, which the graph runner prepares and executes during serving. SGLang-Neuron also supports PD disaggregation. In this configuration, a prefill server computes the prompt KV cache and transfers it to a decode server using NIXL over Elastic Fabric Adapter (EFA). The decode server then generates the remaining tokens. Each server still uses SGLang’s scheduler and request lifecycle, while the plugin manages the Trainium execution and KV-cache transfer path. Serving models on AWS Trainium using SGLang-Neuron To follow this walkthrough, you use a Trainium3 Amazon Elastic Compute Cloud (Amazon EC2) instance with a pytorch-inference-sglang-neuronx container. Here are the steps: Start the container. 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 Launch the SGLang model server inside the container. 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 Send a request to the model server. 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 }' For additional information, please see the SGLang-Neuron getting started guide for setup instructions, and explore the repository documentation for supported models, features, and validated configurations. Conclusion SGLang-Neuron brings SGLang’s familiar serving workflow to AWS Trainium through native PyTorch. You can get started by launching the SGLang-Neuron container and following the quickstart to serve a supported model. Try it with your workloads, contribute through the project repository, and share feedback to help shape future model and feature support. About the authors Alex Nails Alex is a Member of Technical Staff at RadixArk. He leads MultiHardware and is a core developer of SGLang. Yihao Wang Yihao is a Member of Technical Staff at RadixArk. He works on both multi-modal/diffusion models and multi-hardware. Lingyan Hao Lingyan is a Product Manager at RadixArk. She works on SGLang adoption and multi-hardware enablement. Srecharan Selvam Srecharan is a software engineer working on AI Inference technology at Annapurna Labs, Amazon. Zeyu Wang Zeyu is a software engineer working on AI Inference technology at Annapurna Labs, Amazon. Vinesh Ravuri Vinesh is a software engineer working on AI Inference technology at Annapurna Labs, Amazon. Yuchen Jiang Yuchen is a software engineer working on AI Inference technology at Annapurna Labs, Amazon. Pinak Panigrahi Pinak works on performance optimization of AI models on AWS Trainium at Annapurna Labs, Amazon. Chi Wang Chi is a Software Development Manager working on AI Inference technology at Annapurna Labs, Amazon.