NKI and Neuron Architecture

NKI currently supports the following NeuronDevice generations:

• Trainium/Inferentia2, available on AWS trn1, trn1n and inf2 instances

• Trainium2, available on AWS trn2 instances and UltraServers

• Trainium3, available on AWS trn3 instances and UltraServers

The documents below provide an architecture deep dive of each NeuronDevice generation, with a focus on areas that NKI developers can directly control through kernel implementation.

• Trainium/Inferentia2 Architecture Guide serves as a foundational architecture guide for understanding the basics of any NeuronDevice generation.

• Trainium2 Architecture Guide walks through the architecture enhancements when compared to the previous generation.

• Trainium3 Architecture Guide covers the enhancements for the next-generation Trainium ML accelerators.

Neuron recommends new NKI developers start with Trainium/Inferentia2 Architecture Guide before exploring newer NeuronDevice architecture.

Trainium/Inferentia2 Architecture Guide

Foundational architecture guide for understanding NeuronDevice basics.

Trainium2 Architecture Guide

Architecture enhancements and improvements in the Trainium2 generation.

Trainium3 Architecture Guide

Latest architecture features and capabilities in Trainium3 devices.