This document is relevant for: Trn2, Trn3
NKI Library Kernel Utilities Reference#
The NKI Library provides utility modules to simplify common patterns in NKI kernel development. These utilities help manage memory allocation, dimension tiling, and data broadcasting.
Source code for these utilities can be found at: aws-neuron/nki-library
Memory Management#
Stack-based SBUF memory allocator with scope management and multi-buffering support. |
Tensor Operations#
Broadcasts a single partition across the partition dimension using hardware shuffle. |
Iteration Helpers#
TiledRange |
Divides dimensions into tiles with automatic remainder handling. |
This document is relevant for: Trn2, Trn3