Artificial Intelligence1 min read

Deploy frontier AI models on AWS Trainium with SGLang-Neuron

By · Published by Everything Blog

In short

The tool runs on native PyTorch via an out-of-tree architecture, separating the core server operations from custom chip optimizations. Under this model, the base framework handles prefix caching and request management while the specialized plugin drives hardware-specific attention, graph compilation, and mixture-of-experts logic. Additionally, the system enables prefill and decode disaggregation by transferring calculated key-value caches between servers using NIXL over Elastic Fabric Adapter. Users can deploy workloads directly onto Trainium3 Amazon Elastic Compute Cloud instances using the dedicated container.

Key points

  • Developers can now execute frontier model inference and reinforcement learning rollouts u…: Developers can now execute frontier model inference and reinforcement learning rollouts using standard SGLang APIs on Trainium.
  • The modular out-of-tree plugin design enables hardware enhancements to evolve without alt…: The modular out-of-tree plugin design enables hardware enhancements to evolve without altering the core framework logic.
  • Model serving is immediately accessible on Trainium3 EC2 instances through the pytorch-in…: Model serving is immediately accessible on Trainium3 EC2 instances through the pytorch-inference-sglang-neuronx container.

Sources checked by our analysis pipeline: 1

Artificial Intelligence

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,…

Artificial Intelligence