Skill v1.0.0

currentAutomated scan100/100

aws-neuron/neuron-agentic-development/neuron-nki-debugging

──Details

PublishedOctober 29, 2026 at 09:32 AM

Content Hashsha256:15e74b64cdfd1e7e...

Git SHA4782c0ba245f

──Files

Files (1 file, 10.5 KB)

SKILL.md10.5 KBactive

SKILL.md · 306 lines · 10.5 KB

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]" metadata: version: "1.0.0"

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:

python

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:

Check environment: echo $NKI_VENV_PATH
If empty, read .claude/nki-dev-suite.local.md and extract nki_venv_path from YAML frontmatter
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:

bash

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

python

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

python

@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

python

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

bash

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.

python

# 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)

python

# 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:

python

# 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

python

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

All versions