neuron-nki-debugging

name	neuron-nki-debugging
description	This skill guides debugging NKI compilation errors on Neuron hardware. Use when encountering "compiler error on device", "debug NKI kernel", "test kernel on trn2/trn3", "neuronx-cc compilation failed", "validate kernel on hardware", "run kernel on trainium", or asking "how to debug NKI compilation errors on device".
argument-hint	[kernel file]

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

Hardware	Instance	Target Flag	Generation
Trainium 1	trn1	--target trn1	gen2
Trainium 1n	trn1n	--target trn1n	gen2
Inferentia 2	inf2	--target inf2	gen2
Trainium 2	trn2	--target trn2	gen3
Trainium 3	trn3	--target trn3	gen4

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"

Flag	Purpose
--target	Hardware platform (trn1, trn2, trn3, inf2)
--lnc 1	Single NeuronCore (simplifies debugging)
NEURON_RT_VISIBLE_CORES	Pin to specific core(s) — prevents contention when multiple agents run concurrently

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

Error Pattern	Category	Reference
NCC_EVRF*	Verification error	See references/ncc-verification-errors.md
NCC_EOOM*	Out of memory	See references/ncc-memory-resource-errors.md
NCC_E* (other)	Type/operation error	See references/ncc-type-operation-errors.md

Quick Reference

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

Common Error Quick Fixes

Error Code	Category	Quick Fix
NCC_EVRF001	Unsupported operator	Use alternative operator from neuronx-cc list-operators
NCC_EOOM001	Memory exceeded	Reduce batch size, use tensor/pipeline parallelism
NCC_EVRF007	Instruction limit	Apply model parallelism
NCC_EVRF005	Unsupported FP8 type	Convert to float16/bfloat16 or use gen3+ hardware
NCC_EARG001	LNC configuration	Use supported LNC count for target hardware
NCC_EVRF024	Output tensor > 4GB	Reduce tensor size or use tensor parallelism

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:

Setting	Source	Description
nki_venv_path	.claude/nki-dev-suite.local.md or NKI_VENV_PATH	Python venv with neuronx packages

Related skills:

Skill	Use When
/neuron-nki-profiling	Profile kernel performance
/neuron-nki-docs	Look up API documentation and error codes