This document is relevant for: Trn2, Trn3

vLLM Neuron Plugin (Beta) Documentation#

The vLLM Neuron plugin brings the full vLLM serving stack to AWS Trainium accelerators. It supports continuous batching, speculative decoding (EAGLE3), disaggregated inference, structured outputs, multimodal models, and more — all accessible through the standard vllm serve command and OpenAI-compatible API.

For a high-level overview of inference on Neuron and help choosing the right inference solution, see Inference on Neuron. The source code for the vLLM Neuron plugin is hosted in the vLLM Neuron GitHub repository.

Get started#

Setup guide

Install and configure vLLM Neuron on a Trainium instance.

Online serving quickstart

Launch an OpenAI-compatible API server and send your first chat request.

Offline serving quickstart

High-throughput batch inference with the vllm.LLM Python API.

Migration from NxD Inference

Migrate existing NxDI deployments to the vLLM Neuron plugin.

Deploy & serve#

Features guide

Configure and tune all serving features — bucketing, quantization, DI, speculation, and more.

Configuration reference

All Neuron-specific options in additional_config and environment variables.

Profiling workloads

Capture Neuron Runtime profiles via built-in profiler endpoints.

Model recipes

Supported models and their feature tables.

Model Recipes#

Deploy Llama 3

Model recipe for the Llama 3 family (1B, 8B, 70B) on Trn2/Trn3.

Deploy GPT-OSS

Model recipe for GPT-OSS 20B and 120B (MoE) on Trn2/Trn3.

Deploy Qwen3-VL

Model recipe for Qwen3-VL 32B (multimodal) on Trn2/Trn3.

Deploy Qwen3-Embedding 8B

Model recipe for Qwen3-Embedding 8B (pooling / embeddings) on Trn2/Trn3.

Tutorials#

EAGLE3 speculative decoding (Llama 3.1)

Run Llama 3.1 8B with an EAGLE3 draft model for higher throughput.

EAGLE3 speculative decoding (GPT-OSS)

Run GPT-OSS-120B with an EAGLE3 draft model for higher throughput.

GPT-OSS deployment tutorial

End-to-end deployment of GPT-OSS on Trn2/Trn3.

Qwen3-VL multimodal tutorial

Deploy Qwen3-VL 32B for multimodal inference.

Disaggregated inference: 1P1D and xPyD

Configure disaggregated inference topologies.

Disaggregated encoder: 1E1PD and xEyPD

Configure encoder-disaggregated (EPD) multimodal topologies.

Prefix caching benchmark

Measure TTFT improvement from prefix caching with GPT-OSS.

Deploy Qwen3-Embedding-8B

Serve embeddings via /v1/embeddings with a pooling model.

Model development#

Onboard a new model

Implement and register a new architecture with vLLM.

Onboard a vision-language model

Add a vision encoder tower on top of the text-decoder flow.

Optimizing a vision-language model

Roofline, sharding, and profiling to optimize a VLM.

CPU development workflow

Develop and test without Neuron hardware.

NKI CPU simulator

Validate NKI kernel correctness on CPU.

Debugging model code

Use pdb and print statements to inspect model execution.

Debugging accuracy issues

Methodology for isolating where accuracy drift is introduced.

Accuracy debugger tools

Run the automated debugger pipeline and interpret results.

Concepts & architecture#

Parallelism#

Topic	Description
Attention DP	Sharding Q/O weights across DP groups
Component DP sharding	Per-component independent sharding
Data parallelism	Data parallelism overview
Decode Context Parallelism	KV cache sequence sharding for long contexts
Expert parallelism	Expert parallelism for MoE
Tensor parallelism	Tensor parallelism overview
Vision encoder parallelism	Independent TP/DP for vision encoders

Speculation#

Topic	Description
Speculative decoding	EAGLE3 internals on Neuron

Multimodal#

Topic	Description
Block packing attention	FFD block packing for multi-image attention efficiency
On-Device Encoder Cache	Block-based on-device cache for vision encoder outputs
M-RoPE	Spatial position embeddings for VLMs

vLLM integration#

Topic	Description
vLLM integration design reference	Plugin registration, scheduler, KV cache
KV cache integration	KV cache integration points with vLLM
Additional config	Additional configuration options
Async scheduling and execution	Async scheduling and execution design
Context-length bucketing	Reducing decode HBM reads
Memory management	Available memory determination
KV cache quantization	FP8 KV cache mechanics
Metrics	Production metrics design
Neuron profiling	Profiling integration
Neuron scheduler	Holdback queue and admission control
Block KV batching and padding	Block KV cache architecture
Prefix caching	Prefill segmentation and KV reuse
Pooling models	Pooling model execution on Neuron
Prompt embeddings	Prompt embedding support
Structured outputs and tool calling	JSON, regex, grammar-constrained generation
Disaggregated inference	DI architecture, NIXL transport, hybrid TP, DCP

Framework#

Topic	Description
Async execution	Async execution double buffering
Checkpoint loading	Weight loading with parallelism
Model bringup	Model bringup workflow
Model factory	Model registry and factory pattern

Compilation#

Topic	Description
Compilation cache	Compilation cache (hit/miss, remote store)
CPU compilation	Ahead-of-time CPU compilation (NEFF extraction)
FX passes architecture	FX passes architecture
Aliasing output rewrite pass	Aliasing output rewrite pass
Device rewriting FX pass	Device rewriting FX pass
Inplace to out-of-place pass	Inplace to out-of-place rewrite

Accuracy#

Topic	Description
Accuracy debugging design	Accuracy debugging framework
KV cache analysis	KV cache analysis
Logit validation	Logit validation
Module test guidelines	Module test guidelines
Tensor capture	Tensor capture
Tensor compare	Tensor compare
Tensor replacement	Tensor replacement