AWS Neuron Documentation#

AWS Neuron is the software development kit for deep learning and generative AI on AWS Inferentia and AWS Trainium instances. Neuron supports multiple development paths: serving large language models with vLLM, training and inference with PyTorch (native) and JAX, authoring custom kernels with NKI, and direct use of the Neuron Graph Compiler and Runtime.

Current release: Neuron 2.33.0

Released October 19, 2026. Select this card for the details!

Train and deploy models on Neuron#

Neuron now supports native PyTorch on Trn2 and Trn3 platforms with TorchNeuron. Run your existing PyTorch code in eager mode or with torch.compile, and change the device from cuda to neuron. TorchNeuron is currently in Beta.

TorchNeuron: Native PyTorch on Trainium (Beta)

Use standard PyTorch APIs on Trainium: eager mode, torch.compile with backend="neuron", and the FSDP, DTensor, DDP, and Tensor Parallelism distributed APIs. Most PyTorch workflows run with few or no changes.

Train models on Neuron

What Neuron provides for model training: native PyTorch (TorchNeuron), TorchTitan Neuron, and Hugging Face.

Deploy models on Neuron

What Neuron provides for inference: vLLM Neuron and vLLM Omni Neuron for model serving.

Monitor and optimize your model#

Use Neuron Observer to monitor the health of your Trainium cluster, and use Neuron Explorer to profile your model and find ways to make it faster.

Neuron Observer (Beta)

Monitor the health of Trainium clusters on Amazon EKS. Neuron Observer detects hardware faults and finds the affected node and device. Supported on Trn2 and Trn3.

Neuron Explorer

Profile and debug your model and kernels on Trainium and Inferentia. Use the CLI, the UI, or the VS Code integration, and get AI-driven optimization recommendations.

Develop with agent support#

Neuron Agentic Development

Use open-source AI agents and skills in coding environments such as Kiro and Claude Code to develop on Trainium. Write, debug, and profile NKI kernels, or port Hugging Face transformer models to run on Trainium, with natural-language prompts.

Who Neuron is for#

ML engineers deploying production models — Deploy prepared Neuron Deep Learning AMIs (DLAMIs) and Deep Learning Containers (DLCs) on Amazon EC2 Trainium and Inferentia instances. Start with the DLAMI setup guide or the DLC quickstart.
- Serving LLMs — Use vLLM Neuron to serve open-source LLMs with minimal code changes. Start with the online serving quickstart or offline serving quickstart.
ML researchers and model developers — Use native PyTorch on Trainium with eager mode, torch.compile, and standard distributed APIs. Start with TorchNeuron.
Performance engineers optimizing kernels — Use NKI to write custom kernels with direct NeuronCore access, or pick from the NKI Library’s pre-optimized kernels. Start with the NKI quickstart and NKI Library.

Start here#

Pick the task that matches what you want to do.

Get started with a Neuron DLAMI and PyTorch

Launch a Trainium or Inferentia EC2 instance with a pre-configured Neuron Deep Learning AMI (DLAMI) and PyTorch. The DLAMI bundles the Neuron SDK, framework virtual environments (PyTorch, JAX, vLLM), and the system tools — no manual install required. See Install PyTorch via Deep Learning AMI and get started!

Serve a large language model

Run LLM inference on Trainium and Inferentia with vLLM Neuron. Supports OpenAI-compatible APIs, continuous batching, and speculative decoding. See the offline or online serving quickstart.

Serve a diffusion or video model

Run diffusion and multimodal generation models (such as WAN 2.2 video generation) on Trainium 2 and Trainium 3 with vLLM Omni Neuron (Beta). See the offline or online serving quickstart.

Train a model with PyTorch

Use native PyTorch on Trainium (TorchNeuron) with eager mode, torch.compile, and the standard distributed APIs (FSDP, DTensor, DDP). Existing PyTorch code runs with minimal changes; primarily swap cuda for neuron on your tensors.

Write custom NKI kernels

Program NeuronCores directly with NKI when you need finer control than framework-level compilation provides. NKI offers tile-level programming with Python and NumPy-like syntax, and ships with a library of pre-optimized kernels (attention, MoE, and others).

Neuron SDK organization#

The Neuron SDK includes these major components.

Component	What it does and when to use it
TorchNeuron	Native PyTorch backend for Trainium, with eager mode, `torch.compile`, and standard distributed APIs. Use it to train or run PyTorch models with few code changes.
TorchTitan Neuron	PyTorch-native framework for large-scale model pretraining. Use it to pretrain LLMs on Trainium clusters.
vLLM Neuron	vLLM plugin for LLM serving with OpenAI-compatible APIs and continuous batching. Use it to serve open-source LLMs in production.
vLLM Omni Neuron	vLLM plugin for diffusion and multimodal generation models. Use it to serve image and video generation models on Trn2 and Trn3.
JAX NeuronX	JAX integration for Neuron (Beta). Use it to run JAX training and inference code on Trainium.
Neuron Kernel Interface (NKI)	Python programming interface for custom kernels with direct NeuronCore access. Use it when framework-level compilation does not give you the performance you need.
NKI Library	Open-source, pre-optimized kernels for attention, MLP, MoE, normalization, and more. Use these kernels in your models, or read them as reference implementations.
Neuron Graph Compiler	Compiles model graphs and NKI kernels (`neuronx-cc`) into Neuron Executable File Format (NEFF) files. Use the compiler options to tune precision and performance.
Neuron Runtime	Loads NEFF files and runs them on NeuronCores. Manages devices, memory, and collective communications. Use its settings to configure and debug execution.
Neuron Explorer	Profiling and debugging suite with AI-driven recommendations, available from the CLI, the UI, or VS Code. Use it to find and fix performance problems.
Neuron Observer	Health monitoring for Trainium clusters on Amazon EKS. Use it to detect hardware faults and find the affected node and device.
Neuron system tools	Command-line tools to monitor and inspect Neuron devices on an instance. Use them to check device usage and status.
Neuron Agentic Development	Skills and context for AI coding agents that work with Neuron. Use them to speed up model and kernel development.
Deploy on AWS	Pre-configured Deep Learning AMIs (DLAMIs) and Deep Learning Containers (DLCs) for Amazon EC2, EKS, ECS, SageMaker, and ParallelCluster. Use them to set up an environment without a manual install.
Legacy inference frameworks · Legacy training frameworks	Earlier components, such as PyTorch NeuronX (`torch-neuronx`) and the NeuronX Distributed (NxD) libraries. Use them only to maintain existing workloads.

Learn more#

What is AWS Neuron? — Background on Inferentia, Trainium, and the Neuron SDK.
Release notes — Component-by-component release notes for every Neuron SDK version.
Open source and contribute — Public GitHub repositories, contribution guidelines, and source for TorchNeuron, NKI Library, NKI Samples, vLLM Neuron, and Neuron Agentic Development.
News and blogs — Feature announcements, technical deep dives, and customer stories.
FAQ and troubleshooting — Common questions and solutions for Neuron SDK issues.
Archived documentation — Reference material for MXNet Neuron, TensorFlow Neuron, torch-neuron (Inf1), and other legacy components.

AWS Neuron Documentation

Contents