nki.jit Decorator Reference#

This topic provides reference for using NKI JIT decorators and kernel execution APIs with the AWS Neuron SDK. Use this page to look up the available decorators for compiling and running NKI kernels on NeuronDevices.

Overview#

NKI provides decorators and APIs for compiling and executing kernels on NeuronDevices. The primary entry point is nki.jit, which automatically detects the ML framework in use and compiles the kernel accordingly.

Compilation and Execution#

jit

This decorator compiles a top-level NKI function to run on NeuronDevices.