Running SGLang‑Neuron on AWS Trainium: Architecture and Ops Guide

A public Deep Learning Container for SGLang‑Neuron has been released, allowing you to launch SGLang workloads on AWS Trainium without changing the familiar SGLang server interface. Engineers who build or operate large language model services can now target Trainium’s AI‑accelerated hardware using the same generation APIs they already use, which simplifies migration and opens cost‑effective scaling options.

How SGLang‑Neuron integrates with Trainium

SGLang‑Neuron is an out‑of‑tree plugin that registers the Trainium device through the standard SGLang plugin entry points. Setting the environment variable SGLANG_PLATFORM=neuron selects the plugin when the server starts. The core SGLang stack continues to provide the API server, request lifecycle, RadixAttention, and prefix‑cache bookkeeping. The plugin supplies Trainium‑specific implementations for model layers, attention, mixture‑of‑experts, and a graph runner.

Under the hood, TorchNeuron exposes Trainium as a native PyTorch device. SGLang‑Neuron builds its execution path on regular PyTorch operations and injects performance‑critical kernels written with the Neuron Kernel Interface (NKI). These kernels, together with the surrounding PyTorch graph, are compiled via torch.compile. The resulting executable program is handed to the graph runner, which orchestrates inference on Trainium.

When using PD disaggregation, a prefill server builds the prompt KV cache and streams it to a separate decode server over Elastic Fabric Adapter (EFA) using the NIXL transport. Both servers retain the SGLang scheduler, while the plugin manages the Trainium execution and KV‑cache transfer.

Deploying the reference container

The recommended entry point is the pytorch-inference-sglang-neuronx container hosted in Amazon ECR. A typical workflow on a Trainium3 EC2 instance looks like this:

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

Inside the container, start the SGLang server with the desired model and parallelism settings:

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

Once the server is listening, a standard HTTP request can be used to generate text, for example with curl against http://127.0.0.1:30000/v1/chat/completions.

Operational and security implications

• Plugin isolation: Because the Trainium support lives in an out‑of‑tree plugin, updates to SGLang core and Trainium drivers can be applied independently, reducing coordination overhead.
• Container privileges: The reference container runs with --privileged and --ipc=host flags, which broaden the attack surface. Teams should restrict container execution to trusted nodes and monitor host‑level activity.
• Resource allocation: Flags such as --mem-fraction-static 0.8 and KV‑cache datatype fp8_e4m3 affect memory pressure on the Trainium device. Operators need to size instances accordingly and watch for OOM events.
• Network considerations for PD disaggregation: Using EFA and NIXL introduces high‑throughput, low‑latency networking requirements. Security groups and VPC configurations must allow the required traffic while maintaining isolation between prefill and decode servers.
• Observability: The plugin does not alter SGLang’s existing metrics endpoints, so existing monitoring pipelines can continue to collect latency, request count, and GPU/Trainium utilization data.

Related CloudNinjas coverage: AWS.

What This Means For Practitioners

• Evaluate whether your workload benefits from Trainium’s cost model versus existing GPU deployments; the same SGLang API can be reused.
• Plan for privileged container execution and ensure host hardening practices are in place.
• Test memory and KV‑cache settings on a small instance before scaling to larger TP/EP configurations.
• If you need PD disaggregation, verify EFA support in your VPC and incorporate NIXL monitoring.
• Leverage existing SGLang observability hooks; no new instrumentation is required for basic health checks.