Inference on AWS Neuron#

AWS Neuron runs inference on AWS Trainium and Inferentia for large language models (LLMs), image and video generation, and custom model architectures. This page helps you pick a serving path — the vLLM Neuron plugin, vLLM Omni Neuron, or native PyTorch. Neuron supports both low-latency online serving and high-throughput batch inference. NxD Inference and NxD Core are legacy frameworks, announced end of support starting the Neuron 2.34 release; see Legacy inference frameworks below.

Does this page fit your task?

Use this page if you are choosing how to serve a model on AWS Trainium or Inferentia and want the right path without knowing which library implements it. Scenario: you have a trained model (an open-weight LLM, a fine-tuned checkpoint, or a diffusion model) and want to deploy it for online or offline inference on Neuron. Not this page if you are training a model — see Train models on Neuron instead.

Choose your inference path#

The recommended serving path for new inference work on Neuron is the vLLM Neuron plugin. For diffusion and multimodal generation, use vLLM Omni Neuron. You can also build on native PyTorch directly. Pick the path that matches your model and hardware.

vLLM Omni Neuron (Beta — diffusion and multimodal)#

For diffusion and multimodal generation models such as video generation, use the vLLM Omni Neuron plugin. It is an out-of-tree Neuron backend for vLLM Omni that reuses vLLM Omni’s upstream serving components and adds Neuron-specific execution, caching, and sharding. The initial Beta release serves the Wan2.2-T2V-A14B (text-to-video) and Wan2.2-I2V-A14B (image-to-video) models on Trn2 and Trn3, with offline generation through the Omni.generate API and OpenAI-compatible online serving at the /v1/videos endpoint. Start from the online or offline serving quickstart. vLLM Omni Neuron is also published as a Neuron Deep Learning Container; see the vLLM Omni Neuron setup guide.

Native PyTorch#

All inference on Trainium runs on native PyTorch (TorchNeuron) underneath, and you can build on it directly — running native PyTorch models with torch.compile, developing a custom inference framework, or adapting another model-serving framework for Trainium. A serving plugin is the recommended path when it fits your model, but it is not required. See Native PyTorch inference on Neuron for how inference works on Trainium and whether to use the torch-neuronx or libtorch-neuronx-lite package. If your workflow is built around the Hugging Face ecosystem, you can also serve Neuron-compiled models with Optimum Neuron, maintained by Hugging Face.

Custom models and performance optimization#

To bring up a custom diffusion or multimodal generation model on Trn2/Trn3, the vLLM Omni Neuron plugin (Beta) ships Neuron-optimized NKI kernels and reference implementations you can adapt. See its model onboarding guide.

To implement a custom model directly on native PyTorch, or to build a custom inference stack, see Native PyTorch inference on Neuron. The legacy NxD Inference and NxD Core frameworks are also available for custom NxDI-API serving and for building custom models on the NxD Core primitives, but they are announced end of support starting Neuron 2.34. See the NKI note below for where custom kernels fit.

Learn the concepts#

overview

vLLM Neuron plugin

The recommended serving path for new inference work — continuous batching, speculative decoding, quantization, and OpenAI-compatible online serving.

deep dive

vLLM Neuron features guide

The serving features vLLM Neuron supports and how to configure them for your workload.

overview

Native PyTorch inference

Run inference with native PyTorch on Trainium, and build a custom inference stack with the torch-neuronx or libtorch-neuronx-lite package.

Start serving#

quickstart

Online serving quickstart

Stand up an OpenAI-compatible vLLM Neuron server and send your first request.

quickstart

Offline serving quickstart

Run offline batch inference with vLLM Neuron for high-throughput workloads.

tutorial

Llama 3 recipe

End-to-end recipe for serving Llama 3 on Neuron with vLLM Neuron.

tutorial

EAGLE3 speculative decoding

Accelerate Llama 3.1 serving with EAGLE3 speculative decoding on Trn2/Trn3.

Note

Where NKI fits in inference

The serving backends use NKI kernels internally — the NKI Library ships tuned attention, MoE, and quantization kernels that the vLLM Neuron plugin and NxD Inference call for supported architectures. Write a custom NKI kernel when you need to close a remaining performance gap.

Legacy inference frameworks#

NxD Inference (NxDI, the neuronx-distributed-inference library) and NxD Core are legacy and are announced end of support on Trn2 and Trn3 starting the Neuron 2.34 release. For new inference work, use the recommended paths above.

The documentation for both lives under Legacy inference frameworks.