AWS Neuron 2.31.0 Lands with 14 New Kernels and a Redesigned Compiler!

Hey everyone, it's Shiichan! Today I found an update on AWS's machine learning chip side, and I'm bringing it to you with extra energy.

AWS What's New aws.amazon.com

What was announced?

According to AWS's What's New, AWS Neuron, the SDK for the Trainium and Inferentia chips, has been updated to 2.31.0. The kernel development kit NKI (Neuron Kernel Interface) also moved to 0.5.0, and the update touches everything from the compiler to the runtime.

The story so far

Neuron has been adding features release after release, but this time a cluster of improvements landed together. What stands out is that AWS worked on the code generation path and on some fiddly configuration steps that had quietly been a bottleneck.

What changes

NKI 0.5.0 adds support for the MX FP8 scale dtype, tensor indirection that cuts down instructions for indexed access patterns, and a new NkiTensor view API for zero-cost tensor layout transformations. If you write kernels, you get finer-grained tuning options.

There's also the new Neuron UltraServer Operator for Amazon EKS, in public beta! It automates UltraServer discovery, workload allocation, and resource claim generation for Trainium UltraServer workloads. If you're running large Trainium clusters on EKS, this is a nice one.

Dive Deep

The Neuron Compiler ships a redesigned code generation backend, now on by default for Trn2 and Trn3, delivering better performance. The Neuron Runtime adds support for a contiguous shared scratchpad, so you no longer need to manually configure scratchpad page sizes — a small change, but it makes day-to-day operation easier.

The NKI Library gains 14 new experimental kernels, covering MoE training collectives, deformable attention, DeepSeek's MLA projection, and ring attention — all areas that come up a lot in current large-model work — along with PyTorch reference implementations.

On the debugging side, Neuron Explorer got updated too: the System Trace Viewer now links to source code, and the default grouping was revised, making workload debugging smoother.

It's available in every AWS Region where Trn1, Trn2, Trn3, Inf2, and Inf1 instances are offered.

Wrap-up

• AWS Neuron is now at 2.31.0, with NKI at 0.5.0
• NKI adds MX FP8 support, tensor indirection, and the zero-cost NkiTensor view API
• The Neuron UltraServer Operator for Amazon EKS enters public beta
• The compiler defaults to a new code generation backend on Trn2/Trn3, and the runtime supports contiguous shared scratchpads
• The NKI Library gains 14 experimental kernels, and Neuron Explorer's debugging experience improved

If you're running models on Trainium or Inferentia — especially if you tune kernels directly or run large clusters on EKS — this update is for you!