Deploy frontier AI models on AWS Trainium with SGLang-Neuron

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:

1. 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

2. 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

3. 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