neuron-nki-profiling

Profiling NKI Kernels

This skill provides a complete workflow for profiling NKI kernel execution on Trainium/Inferentia hardware using Neuron profiling tools.

Quick Start

Minimal workflow to profile a kernel:

# 1. Set environment variables in Python before kernel execution
os.environ['NEURON_RT_INSPECT_ENABLE'] = '1'
os.environ['NEURON_RT_INSPECT_DEVICE_PROFILE'] = '1'
os.environ['NEURON_RT_INSPECT_OUTPUT_DIR'] = './output'

# 2. Run kernel to generate NEFF
python my_kernel.py

# 3. Find the NKI kernel NEFF (skip XLA-generated NEFFs)
NEFF_PATH=$(python3 scripts/identify-neffs.py ./output my_kernel_func_name)

# 4. Capture profile with neuron-explorer
neuron-explorer capture -n $NEFF_PATH -s profile.ntff --profile-nth-exec=2 --enable-dge-notifs

# 5. View results with neuron-explorer
neuron-explorer view --output-format summary-json -n $NEFF_PATH -s profile.ntff

The workflow generates two key artifacts:

• NEFF file: Compiled kernel binary, generated during execution
• NTFF file: Execution trace captured by neuron-explorer

Prerequisites

Before profiling kernels, resolve the NKI virtual environment path:

1. Check environment: echo $NKI_VENV_PATH
2. If empty, read .claude/nki-dev-suite.local.md and extract nki_venv_path from YAML frontmatter
3. If still not found, report: "NKI_VENV_PATH not configured. Set the environment variable or create .claude/nki-dev-suite.local.md with nki_venv_path in frontmatter."

Activate before running any profiling commands:

source $NKI_VENV_PATH/bin/activate

Hardware requirement: Profiling requires execution on actual Trainium/Inferentia hardware.

Complete Profiling Workflow

Step 1: Set Environment Variables

Add these environment variables in your Python script before kernel execution:

import os

# Enable runtime inspection and device profiling
os.environ['NEURON_RT_INSPECT_ENABLE'] = '1'
os.environ['NEURON_RT_INSPECT_DEVICE_PROFILE'] = '1'
os.environ['NEURON_RT_INSPECT_OUTPUT_DIR'] = './output'

# Compiler flags for target hardware
os.environ['NEURON_CC_FLAGS'] = '--target trn2 --lnc 1' # use lnc=2 if explicitely told to.  

# Pin to a specific neuron core(s) to avoid conflicts with concurrent sessions
os.environ['NEURON_RT_VISIBLE_CORES'] = '0' # '0,1', '0-1'