Quickstart Guide

Quickstart Guide#

Welcome to NKIPy! This guide will help you get started quickly.

Installation#

For installation instructions, see the README or the detailed Installation Guide.

Quick install with uv:

git clone https://github.com/aws-neuron/nkipy.git
cd nkipy
uv sync

Basic Usage#

NKIPy kernels can run in three modes:

  1. Pure NumPy (CPU) - Run kernels directly as NumPy code for prototyping

  2. Simulation Mode - Use @simulate_jit to trace and simulate execution

  3. Trainium Hardware - Use @baremetal_jit to compile and run on Trainium

For complete code examples, see the Basic Usage section in the README.

Learning Path#

1. Start with Tutorials#

The best way to learn NKIPy is through hands-on tutorials:

  • Simple Tutorial - Learn the basics with a softmax kernel

    • Define a NKIPy kernel

    • Run as NumPy, simulation, and on Trainium hardware

    • Understand tracing and compilation

  • NKIPy to NKI - Learn how NKIPy lowers to NKI code

  • MLP Tutorial - Build a multi-layer perceptron

See all tutorials in the Tutorials Index.

2. Explore Examples#

Check out the examples/ directory for more complex use cases:

  • examples/playground/simple_nkipy_kernel.py - Simple standalone example

  • examples/models/qwen3_embedding/ - Full model implementation

3. Dive Deeper#

Once you’re comfortable with the basics:

Key Concepts#

  • Tracing: NKIPy traces your NumPy-like code to build a computation graph

  • Compilation: The traced graph is lowered and compiled to binary

  • Execution: Compiled kernels run on Trainium hardware

Getting Help#

Contents