Deploy frontier AI models on AWS Trainium with SGLang-Neuron

Amazon announced that SGLang-Neuron is publicly available through an AWS Deep Learning Container released by RadixArk and Annapurna Labs. The release lets teams serve models and run reinforcement learning rollouts on AWS Trainium using standard SGLang interfaces, as detailed in AWS Machine Learning Blog's report. The architecture runs on native PyTorch, linking Trainium hardware directly to SGLang.

This report details architectural components, hardware compilation pipelines, disaggregated serving mechanics, and deployment parameters for high-throughput inference systems. Infrastructure operators can evaluate container configuration flags, hardware execution pipelines, and network transfers across AWS AI silicon. Readers will find operational commands required to launch models across distributed topologies.

01Out of tree container release by RadixArk

SGLang-Neuron connects SGLang with AWS Trainium through an out-of-tree plugin design. Setting the SGLANG_PLATFORM environment variable to neuron selects the registered runtime when starting the server. SGLang core continues providing the API server, request lifecycle management, RadixAttention handling, and prefix-cache bookkeeping.

The plugin implements Trainium-specific models, attention mechanisms, mixture of experts, and graph runner orchestration directly. This modular structure separates accelerator support from core code so Annapurna Labs and upstream SGLang teams iterate independently. Engineering teams can adopt newer kernel updates without waiting for mainline serving engine releases.

Amazon and RadixArk distribute the runtime as an AWS Deep Learning Container. Operators pull the preconfigured image from Amazon Elastic Container Registry, eliminating manual driver builds and dependency conflicts. This packaging delivers identical SGLang command interfaces across Trainium silicon without changing external application endpoints.

02Native PyTorch integration with TorchNeuron and NKI

TorchNeuron registers AWS Trainium as a native PyTorch device, establishing standard PyTorch operations as the primary execution path. SGLang-Neuron incorporates custom kernels written with the Neuron Kernel Interface to optimize compute-heavy operations. These targeted kernels accelerate matrix transformations across Trainium compute engines without breaking framework consistency.

The execution stack compiles these standard operations and specialized kernels into executable programs using torch.compile. During model initialization, the graph runner prepares compiled artifacts and schedules execution across available hardware resources. This compilation sequence extracts hardware acceleration while retaining standard PyTorch abstractions throughout serving.

Because the execution path relies on native PyTorch, developers avoid building custom runtime bridges for inference. The graph runner manages compiled programs during active inference, minimizing framework overhead while executing incoming user requests. Engineering teams preserve standard debugging workflows alongside hardware-level performance optimizations.

03Prefill and decode disaggregation over EFA

SGLang-Neuron supports disaggregated prefill and decode execution configurations. In this multi-node topology, a prefill server calculates prompt key-value caches and hands them off to a decode server. The decode server receives the cache data and generates subsequent output tokens across dedicated accelerator resources.

Inter-node key-value cache transfers travel over the Neuron Intermittent eXchange Layer using AWS Elastic Fabric Adapter networking. This direct fabric connection transfers intermediate tensor data between servers with minimal communication latency. Offloading KV cache transmission to high-bandwidth fabric adapters removes network stalls during long-context processing.

Both prefill and decode servers run SGLang schedulers and request lifecycle managers to maintain consistent routing. The plugin manages Trainium-specific execution pipelines and coordinates physical key-value transfers across nodes. This division optimizes hardware utilization across intensive prompt processing and memory-bound token generation phases.

04Trainium3 deployment parameters and runtime flags

Running the environment requires a Trainium3 Amazon Elastic Compute Cloud instance hosting the pytorch-inference-sglang-neuron container. Administrators start privileged Docker sessions with host inter-process communication using the public image tag public.ecr.aws/neuron/pytorch-inference-sglang-neuron:v0.1.0-neuron-py314-sdk-2.33.0-ubuntu26.04. The container includes Neuron SDK version 2.33.0 alongside Python 3.14 and Ubuntu 26.04.

Launching the model server begins by assigning SGLANG_PLATFORM to neuron and executing sglang.launch_server. Command line flags configure an openai/gpt-oss-120b model path with tensor parallelism set to eight and expert parallelism set to eight. The server enables an fp8_e4m3 key-value cache data type to conserve memory bandwidth.

Operational settings designate a 10240 token context length, an 8192 token chunked prefill size, and a limit of 64 concurrent requests. Total token capacity reaches 655360 tokens with static memory fraction established at 0.8. The service listens on local IP address 127.0.0.1 using port 30000.

05Model server verification and API invocation

After the model server initializes on port 30000, engineers verify endpoint operation using standard HTTP requests. A curl command submits a JSON payload to the local v1 chat completions route for model openai/gpt-oss-120b. The test payload transmits a prompt asking why the sky appears blue.

The request sets temperature to zero for deterministic output and caps generation at 128 tokens. SGLang-Neuron processes the input through the graph runner, executing compiled operations on Trainium chips to return standard response data. The resulting output conforms to standard generation interfaces without requiring custom client adaptation.

This setup provides a direct deployment path for frontier model inference and reinforcement learning rollouts on AWS silicon. Repository documentation details supported model architectures, feature sets, and validated cluster profiles. Annapurna Labs and RadixArk encourage engineers to test workloads and contribute feedback through the repository.