Neuron Nki Profiling

# 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

source $NKI_VENV_PATH/bin/activate

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'

python my_kernel.py

# 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

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

mkdir -p ./profiles/run_001

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

# 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

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

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

# 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}'

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

./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/
│   └── ...

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