Skip to content

Releases

v0.1.0

v0.1.0 Latest
Latest

Choose a tag to compare

View all tags
@alexnails alexnails released this 12 Nov 21:38
v0.1.0
4e24326

Component Release Notes for the SGLang-Neuron plug-in

Latest version: 0.1.0 (Public Preview)

The release notes for SGLang-Neuron. Read them for the details about the changes, improvements, and bug fixes for all release versions of the AWS Neuron SDK.

SGLang-Neuron (11/12/2026 Public Preview Release)

SGLang-Neuron 0.1.0 (sglang-neuron) is the first release of the plug-in. It brings the SGLang serving stack to AWS Trainium as a native PyTorch plug-in. In this public preview, large language models (LLMs) run on Trn3 instances on one fully asynchronous runtime, and multimodal diffusion models run on Trn2 instances through SGLang Diffusion. To run the plug-in, use the SGLang-Neuron container image.

Note: This release is a public preview. Features, APIs, and supported models can change in later releases. SGLang-Neuron is released by RadixArk in collaboration with the AWS Neuron team. SGLang-Neuron enables developers to serve open models and generate reinforcement learning (RL) rollouts with SGLang on AWS Trainium, using the same SGLang interface they use today.

Highlights

  • SGLang loads the plug-in through a single plug-in entry point. The plug-in runs on TorchNeuron (torch-neuronx).
  • All supported LLMs use one fully asynchronous execution path with prefill-decode disaggregation.
  • GPT-OSS-120B decodes with a giga kernel on Trainium3. The giga kernel runs every layer of a decode step in one kernel call.
  • Qwen3.5-397B-A17B runs on a whole-model megakernel on Trainium3, in aggregated and disaggregated serving.
  • GLM-5.2 (MXFP4) serves a 1M-token context on Trainium3, with prefill-decode disaggregation and multi-token prediction (MTP) speculative decoding.
  • Kimi-K3 serves a 1M-token context on Trainium3, with prefill-decode disaggregation and DSPARK speculative decoding.
  • HiCache L2 keeps prefixes in host memory, so the server can reuse them across requests. A remote DRAM tier below HiCache L2 extends the cache over NIXL (libfabric).
  • You can serve multimodal diffusion models on Trn2. FLUX.1-dev, FLUX.3-Action, MiniMax-H3, Wan2.2 A14B, and LingBot-World-V2 run on a new Neuron platform for SGLang Diffusion, with NKI kernels. For each model, see Multimodal model support and optimizations.

Multimodal model support and optimizations

The multimodal models run on Trn2 instances through SGLang Diffusion on the Neuron platform:

  • FLUX.1-dev generates 1024 x 1024 and 2048 x 2048 images from a text prompt. It uses fused K/V all-gather attention and const-max ring attention kernels.
  • Wan2.2 A14B generates 81-frame videos at 480p, 720p, and 1080p from a text prompt or an input image. It uses classifier-free guidance (CFG) parallelism, a const-max ring attention kernel that keeps K/V resident in SBUF (the on-chip state buffer), and a patch-parallel variational autoencoder (VAE).
  • MiniMax-H3 generates a 124-frame 768p video with its soundtrack, through text-to-video-and-audio (t2va) or first-and-last-frame conditioning (fl2va). It uses a guarded const-max ring attention kernel, an adaptive layer normalization (AdaLN) cache, and a distributed video and audio VAE.
  • FLUX.3-Action returns the next 32 robot actions from three camera views, the robot state, and an instruction. It serves all three DROID checkpoints (base, gd, and sd) over the SGLang OpenPI WebSocket and HTTP action API. It uses const-max ring attention, a video VAE encoder sharded across chips, and UniPC sampling on the device.
  • LingBot-World-V2 streams a world in real time at 480p from one image along a camera path. It uses a flash attention kernel over paged KV history and a streaming VAE that is sharded by space. For interactive control, a session accepts prompt updates and camera actions for each chunk while it runs.

Running SGLang-Neuron with Neuron

In this public preview, you run the plug-in from the plug-in container image. The image has SGLang, the plug-in, and the patched NIXL installed.

Prerequisites

  • For the LLMs, a trn3.48xlarge instance. Kimi-K3 needs a Trn3 UltraServer: at least two trn3.48xlarge instances connected by a switch.
  • For the diffusion models, a trn2.48xlarge instance.
  • The Neuron driver for Neuron SDK 2.33.0 (aws-neuronx-dkms 2.31.7.0) on the host.
  • Docker, and a host directory that holds the model checkpoints.

Step 1: Pull the image

docker pull public.ecr.aws/neuron/pytorch-inference-sglang-neuron:v0.1.0-neuron-py314-sdk-2.33.0-ubuntu26.04

Step 2: Start the server

These commands come from the quickstart in the plug-in repository README. They download Qwen3-8B, start a server on four NeuronCores, wait until it is ready, and send a request. Run them in bash, because they rely on bash to split $DEVICES into separate arguments:

IMAGE=public.ecr.aws/neuron/pytorch-inference-sglang-neuron:v0.1.0-neuron-py314-sdk-2.33.0-ubuntu26.04
MODELS=/path/to/models     # holds Qwen3-8B, mounted at /models
CACHE=/path/to/cache       # the compile cache: the first start compiles, later starts reuse it
docker pull $IMAGE
docker run --rm -v $MODELS:/models $IMAGE hf download Qwen/Qwen3-8B --local-dir /models/Qwen3-8B
DEVICES=$(for d in /dev/neuron*; do printf -- '--device %s ' "$d"; done)

docker run -d --name qwen3 --init --ipc=host --cap-add SYS_NICE -p 127.0.0.1:30000:30000 $DEVICES \
  -v $MODELS:/models:ro -v $CACHE:/cache \
  -e TORCH_NEURONX_NEFF_CACHE_DIR=/cache/neff -e TORCH_NEURONX_HLO_CACHE_DIR=/cache/hlo \
  -e NKI_COMPILE_CACHE_URL=/cache/nki \
  $IMAGE python3 -m sglang.launch_server --model-path /models/Qwen3-8B --tp-size 4 \
  --context-length 4096 --chunked-prefill-size 2048 --max-running-requests 64 --max-total-tokens 266240 \
  --mem-fraction-static 0.6 --host 0.0.0.0 --port 30000

# Wait until the server answers /health. The first start on a new cache compiles the model.
until curl -sf http://127.0.0.1:30000/health; do
  [ "$(docker inspect -f '{{.State.Running}}' qwen3)" = true ] || { echo "exited: docker logs qwen3"; break; }
  sleep 10; done

Replace /path/to/models and /path/to/cache with host directories. Do not mount a volume at /workspace, because the image installs torch-neuronx from that directory.

Step 3: Send a request

curl -s http://127.0.0.1:30000/v1/chat/completions -H 'Content-Type: application/json' -d '{
  "model": "/models/Qwen3-8B", "temperature": 0, "max_tokens": 64,
  "chat_template_kwargs": {"enable_thinking": false},
  "messages": [{"role": "user", "content": "What is 17 * 23? Answer with the number."}]}'

The answer is 17 * 23 = 391. To stop the server, run docker stop -t 120 qwen3 && docker rm qwen3.

By default, the plug-in samples greedily on the device, so each request sets temperature to 0 or top_k to 1. To serve sampled requests, set SGLANG_NEURON_SAMPLER=temperature or SGLANG_NEURON_SAMPLER=topk_topp when you start the server.

Quickstart scripts for each model

The examples folder in the plug-in repository has launch scripts with the validated flags for each LLM except Qwen3-MoE. Each example has the commands to start it in the container image, on one node or on a prefill and decode pair, and a GSM8K script that checks the server's accuracy. For the diffusion models, each model recipe gives the sglang generate or sglang serve command. The container guide explains each docker run option.

Disaggregated serving

A prefill server and a decode server must both run this image, because the KV transfer uses the patched NIXL. Each container also needs the EFA devices (--device /dev/infiniband/uverbs<N> for each device, with --cap-add IPC_LOCK --ulimit memlock=-1) and usually --network host. Clients reach the pair through the load balancer in the image, which needs no Neuron device:

docker run --rm --network host \
  public.ecr.aws/neuron/pytorch-inference-sglang-neuron:v0.1.0-neuron-py314-sdk-2.33.0-ubuntu26.04 \
  python3 -m sglang_neuron.disaggregation.balancer \
    --prefill http://<prefill host>:30000 --decode http://<decode host>:30000

The balancer listens on 127.0.0.1:8000. For the server flags of each model, see the model recipes in the sglang-neuron repository on GitHub.

For the supported models, see the following list.

Supported models

In this public preview, the LLMs run on trn3.48xlarge instances, and the diffusion models run on trn2.48xlarge instances. Kimi-K3 needs a Trn3 UltraServer. The LLMs are served text only. In the Details column, TP is tensor parallelism, EP is expert parallelism, CP is context parallelism, DP is data parallelism, PP is pipeline parallelism, DCP is decode context parallelism, and CFG is classifier-free guidance parallelism.

Model Type Details
Kimi-K3 LLM Checkpoint moonshotai/Kimi-K3. Needs a Trn3 UltraServer (at least two trn3.48xlarge instances connected by a switch). Disaggregated across three nodes: one prefill node (PP8 x TP16 x DCP16) and a decode that spans two nodes (8 DP-attention groups of attention TP32, EP256, DCP32) with DSPARK speculative decoding. MXFP4 experts and an FP8 KV cache. Multi-head latent attention (MLA) with Kimi Delta Attention state. The prefill uses a chunk-keyed prefix cache with HiCache, or the kimi-prompt prefix cache, which keeps the end of each finished prompt for the next turn.
GLM-5.2 (MXFP4) LLM Disaggregated only, one prefill server and one decode server. MLA with DeepSeek sparse attention. The prefill runs 4 pipeline stag...
Read more
Assets 2
Loading