This document is relevant for: Trn2, Trn3

NKI Tutorials#

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.

The full source code of the following tutorials can be also viewed on the nki-samples repository on GitHub.

Basic Operations#

Matrix Multiplication

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

2D Transpose

Implement efficient 2D matrix transpose operations using NKI

Average Pooling 2D

Create custom 2D average pooling kernels for computer vision workloads

Advanced Kernels#

Fused Mamba

Implement fused Mamba state space model kernels

Kernel Optimization

Learn the recommended workflow for optimizing NKI kernels using profiling and performance analysis

This document is relevant for: Trn2, Trn3