This document is relevant for: Inf2, Trn1, Trn2
NeuronDevice Architecture Guide for NKI#
NKI currently supports the following NeuronDevice generations:
Trainium/Inferentia2, available on AWS
trn1,trn1nandinf2instancesTrainium2, available on AWS
trn2instances 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 basics of any NeuronDevice generation, while Trainium2 Architecture Guide walks through architecture enhancements compared to the previous generation in details. Therefore, we suggest new NKI developers start with Trainium/Inferentia2 Architecture Guide before exploring newer NeuronDevice architecture.
This document is relevant for: Inf2, Trn1, Trn2