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'

Step 2: Execute Kernel

Run your kernel script. This compiles and executes the kernel, generating the NEFF file in the output directory.

python my_kernel.py

Important: Compute references on CPU. If the test script computes a reference result (e.g., torch.matmul for comparison), do it on CPU — not on the XLA device. Every XLA graph compiled on-device generates a separate NEFF. Running reference operations on-device creates extra NEFFs that make it hard to identify the NKI kernel's NEFF.

# CORRECT: Reference on CPU — generates only the NKI kernel NEFF
expected = torch.relu(torch.matmul(lhs.cpu(), rhs.cpu()))
result = my_nki_kernel(lhs, rhs)  # Only this generates a NEFF

# WRONG: Reference on device — generates an extra NEFF
expected = torch.relu(torch.matmul(lhs, rhs))  # Compiles to its own NEFF!
result = my_nki_kernel(lhs, rhs)                # Another NEFF

The runtime creates a subdirectory with instance and process ID naming:

./output/
└── i-0823210096b01e7ec_pid_1187583/
    └── neff_*_vnc_0.neff

Step 3: Create Dedicated Profile Folder

Create a dedicated folder for this profiling iteration. This prevents confusion when comparing multiple optimization attempts:

mkdir -p ./profiles/run_001

Organize profile iterations:

./profiles/
├── run_001/              # Baseline profiling
│   ├── profile.ntff
│   └── metrics.json
├── run_002/              # After first optimization
│   ├── profile.ntff
│   └── metrics.json
└── run_003/              # After second optimization

Step 4: Capture Profile with neuron-explorer

Locate the NKI kernel NEFF and capture execution profile:

# Find NKI kernel NEFF by function name (skips XLA-generated NEFFs)
NEFF_PATH=$(python3 scripts/identify-neffs.py ./output my_kernel_func_name)

# Capture profile trace
neuron-explorer capture \
    -n $NEFF_PATH \
    -s ./profiles/run_001/profile.ntff \
    --profile-nth-exec=2 \
    --enable-dge-notifs

Step 5: View Results with neuron-explorer (JSON)

Generate a JSON summary of the profile results:

neuron-explorer view \
    --output-format summary-json \
    -n $NEFF_PATH \
    -s ./profiles/run_001/profile.ntff

This outputs structured JSON with all metrics. Parse for specific values:

neuron-explorer view \
    --output-format summary-json \
    -n $NEFF_PATH \
    -s ./profiles/run_001/profile.ntff | jq '.latency'

Common JSON queries:

# Get latency in milliseconds
jq '.latency'

# Get all engine utilizations
jq '{tensor: .tensor_engine_active_time_percent, vector: .vector_engine_active_time_percent}'

# Get memory metrics
jq '{hbm_read: .hbm_read_bytes, hbm_write: .hbm_write_bytes}'

# Check if memory-bound or compute-bound
jq '{intensity: .mm_arithmetic_intensity, peak_ratio: .peak_flops_bandwidth_ratio}'

Save the full JSON output for later comparison:

neuron-explorer view \
    --output-format summary-json \
    -n $NEFF_PATH \
    -s ./profiles/run_001/profile.ntff > ./profiles/run_001/metrics.json

Step 6: Querying the profile and/or profile analysis (optional)

For detailed analysis of the kernel profile, use the /neuron-nki-profile-querying skill. It allows for high level performance bounds analysis, as well as zoomed in, instruction level investigation of specific inefficiencies through python on parquet.

Output Directory Structure

Understanding the generated file structure:

./output/                                    # NEURON_RT_INSPECT_OUTPUT_DIR
└── i-0823210096b01e7ec_pid_1187583/        # Instance ID + process ID
    ├── neff_307444798579300_vnc_0.neff     # One NEFF per compiled XLA graph
    ├── neff_324387526933418_vnc_0.neff     # (may include non-NKI NEFFs)
    ├── 307444798579300_vnc_0.ntff          # Matching execution traces
    ├── 324387526933418_vnc_0.ntff
    └── ntrace.pb                           # System trace metadata

./profiles/                                  # Organized profile iterations
├── run_001/
│   ├── profile.ntff                        # Execution trace
│   └── metrics.json                        # Summary metrics
├── run_002/
│   └── ...

The instance/pid subdirectory naming (i-xxx_pid_xxx) is automatic and includes the EC2 instance ID and process ID for traceability.

NEFF Identification

When the output directory contains multiple NEFFs (from multiple on-device operations), use the included identify-neffs.py script to identify which NEFF belongs to which kernel:

# List all NEFFs with identification
python3 scripts/identify-neffs.py ./output/i-*_pid_*/
# Output:
#   [NKI:matmul_relu] ./output/.../neff_389250674131083_vnc_0.neff
#     inputs: ['lhs_T', 'rhs', 'tmp.4']  outputs: ['output.51']
#     ntff: ./output/.../389250674131083_vnc_0.ntff
#   [XLA:broadcast,dot,maximum] ./output/.../neff_418708643727628_vnc_0.neff

# Find a specific kernel by name (useful with multiple NKI kernels)
NEFF_PATH=$(python3 scripts/identify-neffs.py ./output/i-*_pid_*/ matmul_relu)

How it works: Each NEFF embeds its compile workdir path (/tmp/.../neuroncc_compile_workdir/<uuid>/). The script reads the HLO module in that workdir. NKI kernels appear as custom-call ops with AwsNeuronCustomNativeKernel and carry a base64-encoded JSON blob containing func_name, input_names, and output_names.

Limitation: Depends on compile workdirs in /tmp/ still existing. Run promptly after kernel execution.

Key Metrics Quick Reference

Comparing Optimization Iterations

When optimizing a kernel, compare metrics across iterations:

# Baseline measurement
neuron-explorer view --output-format summary-json \
    -n $NEFF -s ./profiles/baseline/profile.ntff > ./profiles/baseline/metrics.json

# After optimization
neuron-explorer view --output-format summary-json \
    -n $NEFF -s ./profiles/optimized/profile.ntff > ./profiles/optimized/metrics.json

# Compare latencies
echo "Baseline: $(jq .latency ./profiles/baseline/metrics.json)"
echo "Optimized: $(jq .latency ./profiles/optimized/metrics.json)"

Optimization tracking table:

Keep notes on what changed between iterations to correlate optimizations with metric improvements.

Complete Example

See examples/basic-profiling-workflow.py for a complete end-to-end profiling script demonstrating all steps: environment setup, kernel execution, NEFF identification, profile capture, and JSON metric extraction.

Configuration

Required settings:

Environment variables (set in kernel script):

Related Skills

Troubleshooting

Multiple NEFFs generated (can't tell which is the NKI kernel):

• Primary fix: Compute reference operations (e.g., torch.matmul) on CPU, not on the XLA device. Each on-device XLA graph generates its own NEFF.
• Identify NEFFs: Use python3 scripts/identify-neffs.py ./output to list all NEFFs with their type (NKI vs XLA) and kernel names. See the NEFF Identification section for details.
• Match NEFF to NTFF: Each NEFF neff_<ID>_vnc_0.neff has a matching trace <ID>_vnc_0.ntff in the same directory.

No NEFF file generated:

• Verify NEURON_RT_INSPECT_ENABLE=1 is set before imports
• Check NEURON_RT_INSPECT_OUTPUT_DIR path exists and is writable
• Ensure kernel actually executed (print forces XLA compilation)
• Confirm you are on Trainium/Inferentia hardware: neuron-ls

neuron-explorer capture fails:

• Verify running on Trainium/Inferentia hardware
• Check NEFF file path is correct with ls -la <path>
• Ensure neuronx packages are installed in venv
• Check sufficient disk space for NTFF file

Empty or minimal profile data:

• Use --profile-nth-exec=2 to skip warmup execution
• Add --enable-dge-notifs for detailed DMA analysis
• Verify kernel ran successfully before profiling
• Check NTFF file size is non-trivial: ls -lh profile.ntff

Latency varies between runs:

• Use --profile-nth-exec=2 or higher to skip warmup
• Ensure system is not under other load
• Run multiple iterations and average results
• Check for thermal throttling in profile output