Debugging NKI on Neuron Hardware

This skill provides a workflow for debugging NKI kernel compilation and execution on Trainium/Inferentia hardware.

Quick Start

Minimal working example to test kernel compilation on device:

import os
import torch
from torch_xla.core import xla_model as xm
import nki
import nki.language as nl
import nki.isa as nisa

os.environ["NEURON_CC_FLAGS"] = "--target trn2 --lnc 1"
os.environ["NEURON_PLATFORM_TARGET_OVERRIDE"] = "trn2"

@nki.jit
def add_kernel(a_input, b_input):
    """Element-wise addition kernel."""
    a_tile = nl.ndarray(a_input.shape, dtype=a_input.dtype, buffer=nl.sbuf)
    nisa.dma_copy(dst=a_tile, src=a_input[0:a_input.shape[0], 0:a_input.shape[1]])

    b_tile = nl.ndarray(b_input.shape, dtype=b_input.dtype, buffer=nl.sbuf)
    nisa.dma_copy(dst=b_tile, src=b_input[0:b_input.shape[0], 0:b_input.shape[1]])

    c_tile = nl.ndarray(a_input.shape, dtype=a_input.dtype, buffer=nl.sbuf)
    nisa.tensor_tensor(dst=c_tile, data1=a_tile, data2=b_tile, op=nl.add)

    c_output = nl.ndarray(a_input.shape, dtype=a_input.dtype, buffer=nl.shared_hbm)
    nisa.dma_copy(dst=c_output, src=c_tile)
    return c_output

# Get XLA device and run
device = xm.xla_device()
a = torch.ones((4, 3), dtype=torch.float16).to(device=device)
b = torch.ones((4, 3), dtype=torch.float16).to(device=device)

c = add_kernel(a, b)
print(c)  # Forces XLA compilation and execution

Prerequisites

Before running kernels on device, 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 device tests:

source $NKI_VENV_PATH/bin/activate

Platform Detection

Before compilation, detect the current hardware platform:

Current platform: !neuron-ls | head -3

Platform Target Mapping

Match the --target flag and platform_target decorator argument to your detected hardware.

Standard Debugging Workflow

Step 1: Set Environment Variables

import os

# Standard debugging flags (minimal, fast compilation)
os.environ["NEURON_CC_FLAGS"] = "--target trn2 --lnc 1"
os.environ["NEURON_PLATFORM_TARGET_OVERRIDE"] = "trn2"

# Pin to a specific neuron core to avoid conflicts with concurrent sessions
os.environ["NEURON_RT_VISIBLE_CORES"] = "0"

See references/compiler-flags.md for complete flag reference.

Step 2: Apply Matching Decorator

@nki.jit  # Must match --target and NEURON_PLATFORM_TARGET_OVERRIDE
def my_kernel(input_tensor):
    ...

The platform_target environment variable MUST match the --target in NEURON_CC_FLAGS.

Step 3: Create Test Script

import os
import torch
from torch_xla.core import xla_model as xm
import nki

os.environ["NEURON_CC_FLAGS"] = "--target trn2 --lnc 1"
os.environ["NEURON_PLATFORM_TARGET_OVERRIDE"] = "trn2"

@nki.jit
def kernel(input_tensor):
    # Your kernel implementation
    ...
    return output_tensor

# XLA device execution pattern
device = xm.xla_device()
input_data = torch.randn((128, 512), dtype=torch.float32).to(device=device)

output = kernel(input_data)
print(output)  # Forces XLA compilation - triggers actual compilation

Step 4: Run and Observe

source $NKI_VENV_PATH/bin/activate
python your_test_script.py

Compilation errors appear in the console output. The print() statement forces XLA compilation, which triggers the neuronx-cc compiler.

Step 5: Validate Numerically

Compare device output against a CPU-computed reference using multiple complementary checks — no single metric catches all issues:

• atol / rtol (torch.allclose): Per-element pass/fail gate
• Maximum absolute difference: Worst-case outlier check
• Norm of the difference tensor: Detects widespread small drift
• Cosine similarity: Catches directional errors in high-dimensional outputs

Important: Compute references on CPU, not on the XLA device. Every XLA graph compiled on-device generates a separate NEFF file. Running reference operations (e.g., torch.matmul, torch.softmax) on the XLA device creates extra NEFFs, making it hard to identify which NEFF belongs to the NKI kernel during profiling.

# CORRECT: Reference computed on CPU — only the NKI kernel generates a NEFF
cpu_input = input_data.cpu()
reference_output = reference_implementation(cpu_input)
device_output = output.cpu()

# Use dtype-appropriate tolerances
assert torch.allclose(device_output, reference_output, rtol=1e-5, atol=1e-8)

# WRONG: Reference computed on device — generates an extra NEFF
reference_output = reference_implementation(input_data)  # Compiles to separate NEFF!

For complex kernels where the final output is wrong: decompose the kernel into logical stages and examine intermediate tensors at each boundary. Store intermediates to HBM temporarily, compare each against the matching reference stage, and binary-search for the stage that introduces the error. Once the failing stage is identified, test it with minimal input shapes (e.g., a single 128x128 tile) to isolate whether the issue is in the core logic or in tiling/boundary handling. Remove the debug stores once the issue is resolved.

Compiler Artifacts Mode

For advanced debugging that preserves compiler outputs for inspection, use when you need to understand detailed compilation behavior.

When to use: "compiler artifacts", "compiler flags", "inspect compiler log"

See references/compiler-artifacts.md for:

• Compiler debug flag configuration (--verbose, --target, --lnc)
• Finding the compiler temp folder
• Understanding generated artifacts (*.neff, log-neuron-cc.txt)

Error Resolution

Error Categories

Quick Reference

See references/compiler-error-codes.md for the complete index of all 28 NCC_* error codes.

Common Error Quick Fixes

Profiling (Optional)

To capture execution traces for profiling:

# Add before running kernel
os.environ['NEURON_RT_INSPECT_ENABLE'] = '1'
os.environ['NEURON_RT_INSPECT_DEVICE_PROFILE'] = '1'
os.environ['NEURON_RT_INSPECT_OUTPUT_DIR'] = './output'

This captures NEFF (compiled binary) and NTFF (execution trace) files in the output directory.

Complete Example

import os
import torch
from torch_xla.core import xla_model as xm
import nki
import nki.language as nl
import nki.isa as nisa

# Standard debugging configuration
os.environ["NEURON_CC_FLAGS"] = "--target trn2 --lnc 1"

# Optional: Enable profiling
os.environ['NEURON_RT_INSPECT_ENABLE'] = '1'
os.environ['NEURON_RT_INSPECT_DEVICE_PROFILE'] = '1'
os.environ['NEURON_RT_INSPECT_OUTPUT_DIR'] = './output'
os.environ["NEURON_PLATFORM_TARGET_OVERRIDE"] = "trn2"

@nki.jit
def softmax_kernel(input_tensor):
    """Simple softmax along last dimension."""
    # Load input tile
    tile = nl.ndarray(input_tensor.shape, dtype=input_tensor.dtype, buffer=nl.sbuf)
    nisa.dma_copy(dst=tile, src=input_tensor)

    # Compute softmax
    exp_tile = nl.ndarray(input_tensor.shape, dtype=input_tensor.dtype, buffer=nl.sbuf)
    nisa.activation(dst=exp_tile, data=tile, op=nl.exp)

    sum_tile = nl.ndarray((input_tensor.shape[0], 1), dtype=input_tensor.dtype, buffer=nl.sbuf)
    nisa.tensor_reduce(dst=sum_tile, data=exp_tile, op=nl.add, axis=(1,))

    recip_sum = nl.ndarray((input_tensor.shape[0], 1), dtype=input_tensor.dtype, buffer=nl.sbuf)
    nisa.reciprocal(dst=recip_sum, data=sum_tile)

    result = nl.ndarray(input_tensor.shape, dtype=input_tensor.dtype, buffer=nl.sbuf)
    nisa.tensor_scalar(dst=result, data=exp_tile, op0=nl.multiply, operand0=recip_sum)

    # Store output
    output = nl.ndarray(input_tensor.shape, dtype=input_tensor.dtype, buffer=nl.shared_hbm)
    nisa.dma_copy(dst=output, src=result)
    return output

# Test execution
device = xm.xla_device()
x = torch.randn((64, 128), dtype=torch.float32).to(device=device)

y = softmax_kernel(x)
print(y)  # Triggers compilation

# Validate against PyTorch reference
reference = torch.softmax(x.cpu(), dim=-1)
assert torch.allclose(y.cpu(), reference, rtol=1e-4, atol=1e-6)
print("Validation passed!")

Configuration

Required settings:

Related skills: