Announcing the SGLang AWS Trainium Plugin Beta
SGLang now supports AWS Trainium through a plugin in beta.
Extract
Over the past year, we have continued to expand the hardware SGLang supports, alongside launch-day support for new open models and RL rollouts in Miles. Now, we are collaborating with AWS to bring SGLang to AWS Trainium. As frontier AI moves toward heterogeneous hardware, AWS Trainium represents a unique opportunity. Purpose-built by AWS for AI and available at scale in the cloud, Trainium is designed to deliver strong price-performance for high-volume workloads. Open-model serving and RL rollout generation are prime examples: both run continuously, and cost per token directly determines how much capacity teams can deploy. Our goal is to make SGLang a first-class option for running these workloads on Trainium. In close collaboration with AWS, and in conjunction with the TorchNeuron launch, we are announcing the beta of the SGLang Trainium plugin, built on the newly released TorchNeuron. The beta enables both inference serving and RL rollouts on Trainium. Built on TorchNeuron from the start TorchNeuron is AWS’s native PyTorch backend for Trainium. We began building SGLang’s Neuron backend on TorchNeuron during its private preview, and that decision shaped the entire effort. Because TorchNeuron is native PyTorch, SGLang’s model components and core abstractions carried over to Trainium largely intact. Models on Trainium are served through the same interface developers already use on GPUs, allowing teams to move workloads between GPUs and Trainium with minimal changes. For additional performance, NKI (the Neuron Kernel Interface) enables us to write custom Trainium kernels within the PyTorch workflow. For an open-source project spanning many hardware platforms, this gives contributors a familiar programming model, along with low-level hardware control where performance requires it. What’s in the beta The beta supports two workloads on AWS Trainium2 and Trainium3: LLM inference serving and RL rollout generation for post-training. Initial coverage includes: Models: Qwen3 (dense), Qwen3-MoE-30B-A3B, DeepSeek-V2-Lite, GLM-4, and GLM-5.2/5.3, with support for more large frontier models in progress. Features: Data Parallel / Tensor Parallel / Expert Parallel, Radix Cache, BF16 / FP8, Async scheduler. RL rollouts: SGLang on Trainium as the rollout engine for heterogeneous RL pipelines, described below. We will expand model and feature coverage over the course of the beta, with broader availability planned for later this year. Early preview: heterogeneous RL In …