NKIPy Documentation

Welcome to the NKIPy documentation! NKIPy provides a NumPy-like tensor-level programming layer on top of NKI (Neuron Kernel Interface) for AWS Trainium.

For an overview of the project, installation instructions, and basic usage examples, see the README.

Key Features

• NumPy-like API: Write kernels using familiar NumPy syntax with Python control flow

• HLO Lowering: Tensor operations are traced and lowered to HLO for compilation

• Neuron Compiler Integration: Direct integration with neuronx-cc for generating NKI or executables

Runtime

For kernel execution, NKIPy uses Spike, a lightweight Pythonic runtime layer for AWS Neuron. See the Spike README for details on the runtime architecture and API.

Getting Started

• Installation
  • Prerequisites
  • Installation with uv (Recommended)
  • Alternative: Installation with pip
  • Verifying Installation
  • Troubleshooting
• Quickstart Guide
  • Installation
  • Basic Usage
  • Learning Path
  • Key Concepts
  • Getting Help

Learning NKIPy

• Tutorials
  • 1. Simple NKIPy Tutorial
  • 2. Generate NKI Kernels from NKIPy Kernels
  • 3. MLP NKIPy Tutorial

User Guide

• Tensor Indexing and Slicing
  • Supported Operations
  • Unsupported Operations
  • Limitations
  • Examples
• NKIPy Tracing Architecture
  • Overview
  • How Tracing Works
  • Supported Operations
  • Special Features

Developing NKIPy

• Extending NKIPy Operations
  • Overview
  • Adding a New Operation
• Testing & Code Quality
  • Running Tests
  • Test Categories
  • Linting & Formatting
  • Type Checking
  • Pre-commit Checklist
  • Pre-commit Hooks (Automated)
• Known Issues
  • Frontend
• Building Documentation
  • Prerequisites
  • Building HTML Documentation
  • What Happens During Build
  • Other Build Targets
  • Documentation Structure
  • Adding New Pages
  • Troubleshooting
• API Reference
  • Operations Reference
  • NumPy API Compatibility
  • Distributed APIs
  • Core APIs
  • Runtime APIs