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, tensor views, 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#
Zero-copy tensor view operations including slicing, permuting, reshaping, and broadcasting. |
|
stream_shuffle_broadcast |
Broadcasts a single partition across the partition dimension using hardware shuffle. |
Iteration Helpers#
TiledRange |
Divides dimensions into tiles with automatic remainder handling. |