AWS Neuron 2.32 is here! MXFP8 training and new NKI features power up Trn2

Hi, I'm Shii-chan! Today I found an update about AWS's machine learning chip stack. It might look niche at first glance, but there's some solid substance in here, so let's dig in!

AWS What's New aws.amazon.com

What was announced?

According to AWS What's New, AWS Neuron — the SDK for Trainium and Inferentia — has reached version 2.32.0. NKI (Neuron Kernel Interface), the language used to write kernels, moves to 0.6.0, and 13 new kernels have been added for MoE (Mixture of Experts) training and sparse attention. Trn2 and Trn3 also get new collective communication capabilities.

The story so far

Until now, NKI's collective instructions (like all-gather) assumed every rank held a tensor of the same size. That made it hard to write kernels for cases where the amount of data per rank varies, such as after MoE routing.

On top of that, on-device top-K processing and data-dependent iteration (loops whose length depends on the data itself) were also awkward to express in NKI.

What changes

NKI 0.6.0 addresses these gaps together.

• An on-device top-K instruction lets you compute top-K entirely on the device
• A variable-length all-gather now supports ranks holding differently sized tensors
• New loop constructs support data-dependent iteration

This means kernels for models like MoE, where load isn't evenly distributed across ranks, can now be written more naturally in NKI. For engineers training large models on Trainium or Inferentia, work that previously required custom workarounds is now supported at the library level.

Dive Deep

The lineup of the 13 new NKI Library kernels reflects some very current model architectures.

• A context-encoding kernel for DeepSeek-V3.2's sparse multi-head latent attention
• An MXFP8 attention kernel for the decode step
• A matched pair of MXFP8 forward and backward kernels that let blockwise MoE layers train end to end in MXFP8

An additional 22 kernels also ship with PyTorch reference implementations. Being able to train blockwise MoE layers entirely in the low-precision MXFP8 format end to end is great news if you're trying to cut memory and compute costs for large MoE model training.

On the Neuron Runtime side, Trn2 and Trn3 now support variable-size all-gather, reduce-scatter, and all-to-all collectives. Each rank can send or receive a different number of elements, which makes it possible to build efficient communication for distributed training workloads with uneven load.

A few other updates round out the release:

• The vLLM Neuron plugin moves to vLLM 0.24.0 and now ships in all Neuron Deep Learning AMIs and Deep Learning Containers
• A new Neuron Agentic Development skill helps port transformer models to the vLLM Neuron plugin
• The Neuron Compiler adds explicit control over 64-bit integer compilation
• Neuron Explorer's System Trace Viewer now shows per-core host CPU utilization

AWS Neuron is available in every AWS Region where Amazon EC2 Trn1, Trn2, Trn3, Inf2, and Inf1 instances are available.

Wrap-up

• AWS Neuron 2.32.0 has shipped, with NKI updated to 0.6.0
• NKI gains an on-device top-K instruction, variable-length all-gather, and data-dependent loop constructs
• 13 new kernels target MoE training and sparse attention, including end-to-end MXFP8 training for blockwise MoE and support for DeepSeek-V3.2
• Trn2/Trn3's Neuron Runtime now supports variable-size all-gather, reduce-scatter, and all-to-all
• The vLLM Neuron plugin moves to vLLM 0.24.0, and a new Neuron Agentic Development skill has arrived

If you're training your own MoE models on Trainium or Inferentia, this update is especially worth your attention!