NKI Library Design Specifications#

The NKI Library provides pre-built kernels you can review and modify in your own kernel development with the AWS Neuron SDK and NKI. In this section, learn how the NKI Library kernels are designed and optimized so you can apply the same techniques to your own custom NKI kernels.

Kernel

Description

Source Code

RMSNorm-Quant kernel specification

Performs optional RMS normalization followed by quantization to fp8.

Source code