NKI Guides

This section provides hands-on tutorials for the Neuron Kernel Interface (NKI), demonstrating how to write custom kernels for AWS Trainium and Inferentia instances. These tutorials cover fundamental operations, advanced techniques, and distributed computing patterns using NKI.

Tutorials

Matrix Multiplication

Learn the fundamentals of implementing matrix multiplication in your NKI kernels.

Transpose 2D

Implement efficient 2D matrix transpose operations using NKI

Average Pooling 2D

Create custom 2D average pooling kernels for computer vision workloads

Fused Mamba

Implement fused Mamba state space model kernels

SPMD Tensor Addition

Single Program Multiple Data tensor addition

Multi-Core SPMD Addition

Advanced SPMD tensor addition across multiple NeuronCores

Architecture Guides

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.

How-To Guides

How to Insert NKI Kernels into Models

How to insert a NKI kernel as a custom operator into a PyTorch or JAX model using simple code examples.