VS Init

Goal

Create a new examples/model-serving/<name>/ input bundle using today's VibeServe conventions. Prefer copying and adapting the closest existing example over inventing shared infrastructure.

The user should only need to provide:

• Hugging Face model id, and revision if important.
• Hardware target, such as H100, A100, Trainium, MacBook/MLX, CPU.
• Natural-language workload / optimization goal, for example: "maximize p95 latency for 128-token chat completions at 8 req/s" or "measure prefix caching for long shared prompts."

If the workload goal is missing or vague, ask the user to describe it in plain English before writing files. Always infer the model/workload modality and a bundle slug, then ask the user to confirm both before creating files.

Workflow

1. Inspect existing examples with find examples/model-serving -maxdepth 2 -type d | sort.
2. Identify the workload's input-to-output modality. Do not assume the Hugging Face model id is sufficient. Use model id, tags/config, README/model card if available, and the user's natural-language goal as clues.
3. Propose:
   • Inferred modality/API, such as "text/chat -> text over OpenAI completions".
   • Source example to copy.
   • Bundle slug, such as qwen3-8b-h100-json or llama-3-8b-a100-chat.
4. Ask the user to confirm or correct both the modality/API and slug before creating files. Keep the question concise:
   • "I infer this is <input> -> <output> via <API>, and I would create examples/model-serving/<slug>/. Is that right?"
5. Pick the nearest source example by input-to-output modality and public API first; use hardware as a secondary tie-breaker. The accuracy checker and benchmark are usually tied more tightly to the modality/workload contract than to the accelerator.
   • Text prompt/chat -> text completion, OpenAI completions/chat API: copy from Llama-3-8B.
   • Text prompt/chat -> text completion on Trainium: copy from Llama-3-8B-trn2 only when Trainium-specific setup matters; otherwise start from the generic text example and adjust hardware notes.
   • Text/code input -> edited code output with predicted-output requests: copy from qwen3-32b-code-edit.
   • Text prompt/schema -> constrained JSON text output: copy from Llama-3.1-8B-Instruct-MLX-8bit.
   • Text prompt with long shared prefixes -> text completion with prefix-cache-sensitive benchmark: copy from olmo-hybrid-prefix-caching.
   • Audio stream or WAV input -> transcript text: copy from moonshine-streaming.
   • Text prompt -> image bytes/base64: copy from show-o2-1.5B-HQ*.
   • If no example matches the modality, ask the user which input/output contract to measure and create the smallest checker/benchmark by adapting the closest transport pattern (HTTP, WebSocket, or local VibeServeModel).
6. Create examples/model-serving/<slug>/ with the same layout:
   • OBJECTIVE.md
   • README.md
   • requirements.txt
   • config.json if the source example has one
   • reference/README.md, reference/meta.json, reference/config.json if applicable, reference/reference.py
   • accuracy_checker/README.md, accuracy_checker/checker.py
   • benchmark/README.md, benchmark/benchmark.py
7. Preserve the existing VibeServe contract:
   • Reference path is passed with --ref examples/model-serving/<slug>/reference.
   • Accuracy checker path is passed with --acc-checker examples/model-serving/<slug>/accuracy_checker.
   • Benchmark path is passed with --bench examples/model-serving/<slug>/benchmark.
   • Checkers should be executable as python checker.py.
   • Benchmarks should be executable as python benchmark.py --url http://localhost:8000 ....
8. Do not add a new shared evaluator library unless the user explicitly asks. This skill is for practical example setup, not refactoring.

Modality Inference

Treat model modality as a hypothesis, not a fact, until the user confirms it.

Useful clues:

• AutoModelForCausalLM, text-generation, chat, instruct, code: usually text -> text.
• response_format, JSON schema, grammar/constrained decoding goal: text/schema -> structured text.
• image-to-text, vision-language, vl, vllava, qwen-vl: image+text -> text, which may need a new closest-example adaptation because current examples are mostly text-only and text-to-image.
• text-to-image, diffusion, StableDiffusion, Show-o: text -> image.
• automatic-speech-recognition, speech-to-text, audio, whisper, moonshine: audio -> text.
• text-to-speech, tts: text -> audio, currently no direct model-serving source example; ask before adapting.

Ambiguous cases to clarify:

• Multimodal generative models that support more than one route, such as text -> text and image+text -> text.
• Base models where the user's workload is specialized, such as code editing, JSON generation, long-context caching, or chat serving.
• Model ids that name a framework or hardware format rather than task semantics, such as MLX, GGUF, AWQ, GPTQ, or Neuron.

Even when the modality seems obvious, ask for confirmation before writing files. If it is ambiguous, include the uncertainty in the confirmation question: "I think this is probably text -> text, but it may be image+text -> text. Which modality should this benchmark exercise?"

Slug Naming

Generate a concise lowercase slug with letters, digits, and hyphens only. Prefer:

<model-family>-<size>-<hardware-or-workload>

Examples:

• qwen3-8b-h100-chat
• llama-3-8b-trn2
• qwen3-32b-code-edit
• olmo-prefix-caching

Ask the user to confirm the slug before creating examples/model-serving/<slug>/. If the user supplies a name, normalize it to hyphen-case and confirm it unless the request explicitly says to use that exact name.

Objective

Write OBJECTIVE.md from the user's natural-language goal. Include:

• Model name and serving modality.
• Hardware target.
• Primary metric and whether higher or lower is better.
• Required API shape, usually OpenAI-compatible /v1/completions and optionally /v1/chat/completions.
• Accuracy requirement, usually "must pass the accuracy checker."
• Notes about allowed implementation approaches and disallowed shortcuts.

Keep it concrete enough for an optimizer to know what to improve. If the goal says "benchmark X", name X as the headline metric rather than relying on profiler-only timings.

Reference

For normal Hugging Face text-generation models, make reference/meta.json the source of truth:

{
  "model_id": "org/model-name",
  "revision": null,
  "task": "text-generation"
}

Use revision when the user gives one or reproducibility matters. VibeServe already knows how to use meta.json to materialize model weights in reference/model.

For generic causal LMs, adapt reference/reference.py from Llama-3-8B unless the model requires custom code. Use AutoTokenizer and AutoModelForCausalLM. Set trust_remote_code=True only when the model requires it or the user explicitly accepts it.

For nonstandard models, keep reference/reference.py as an explanatory, runnable reference implementation matching the source example's pattern.

Accuracy Checker

For a new HF causal LM, start from examples/model-serving/Llama-3-8B/accuracy_checker/checker.py.

Adapt:

• Model path default (--model-dir ../model is usually fine once copied into the bundle).
• Device and dtype defaults for the hardware target.
• Prompt suite if the workload is specialized.
• Chat-template behavior if the model is instruct/chat tuned.
• Comparator strictness only when justified.

Default text-matching behavior should be strict greedy matching against Hugging Face outputs:

• Use deterministic generation: do_sample=False, temperature 0 or unset.
• Compare generated token ids first.
• Print decoded reference/custom text and first differing token when failing.
• Exit 0 only when all required cases pass.

For workload-specific checks, add targeted cases instead of weakening the checker. Examples:

• Prefix caching: include long shared prefixes and divergent suffixes.
• Code editing: include buggy code and compare against gold fixes or similarity thresholds.
• JSON/constrained decoding: validate JSON schema and include randomized sentinel text to catch prompt-ignoring shortcuts.

Benchmark

For a new OpenAI-compatible text-generation server, start from examples/model-serving/Llama-3-8B/benchmark/benchmark.py.

Keep these reusable behaviors unless the workload says otherwise:

• --url, --endpoint, --rate, --duration, --num-requests, --max-tokens, --temperature, --prompt-len, --seed, --output-json.
• Streaming SSE request handling.
• TTFT, TPOT, total latency, request throughput, and output token throughput.
• Poisson arrivals for open-loop load.
• Structured JSON output.

Adapt the prompt pool and metric focus to the natural-language goal:

• Throughput workloads: report token throughput as headline.
• Latency workloads: report p50/p95/p99 total latency or TTFT as headline.
• Prefix-cache workloads: include repeated shared-prefix prompts.
• Long-context workloads: add synthetic or dataset-backed long prompts.
• Code-edit/predicted-output workloads: copy from qwen3-32b-code-edit and keep the prediction field behavior.

Do not make benchmark success depend on hidden implementation details. Measure end-to-end behavior through the public API unless the user specifically asks for a local/in-process benchmark.

README And Requirements

Write a short README.md with the exact paths to use:

Use:
- --ref examples/model-serving/<slug>/reference
- --acc-checker examples/model-serving/<slug>/accuracy_checker
- --bench examples/model-serving/<slug>/benchmark

Update requirements.txt from the copied example. Include only dependencies the checker/reference/benchmark need, such as torch, transformers, accelerate, httpx, datasets, jsonschema, soundfile, or websockets.

Validation

After creating or editing the example:

1. Run syntax checks on changed Python files with python3 -m py_compile ... when dependencies are not installed.
2. Run lightweight --help checks for accuracy_checker/checker.py and benchmark/benchmark.py if imports allow it.
3. Verify reference/meta.json is valid JSON.
4. Verify the final layout with find examples/model-serving/<slug> -maxdepth 3 -type f | sort.
5. Do not download large model weights or run GPU-heavy checks unless the user asks.

Handoff

End with:

• The new bundle path.
• The source example copied/adapted.
• The optimization goal captured in OBJECTIVE.md.
• The exact vibeserve command to start optimizing, using the bundle's --ref, --acc-checker, and --bench paths.

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:

Research $ARGUMENTS thoroughly:

1. Find relevant files through provided indices using Glob and Grep
2. Read and analyze document and code
3. Summarize findings with specific file references

NKI Documentation Lookup

This skill provides comprehensive access to NKI (Neuron Kernel Interface) documentation for AWS Trainium/Inferentia kernel development. Use this skill to answer questions about NKI APIs, tutorials, hardware architecture, error codes, and optimization techniques.

Query Type Routing

Route queries to the appropriate documentation based on the query pattern:

Search Strategy

1. Start with indices - Use the index files in references/indices/ to find the exact documentation needed
2. Follow links - Read the linked documentation file for full details
3. Cross-reference - For complex topics, combine information from multiple sources

Index Files

• indices/symbol-lookup.md - Alphabetical API function index with direct links
• indices/task-routing.md - Task-oriented guide ("I want to...")
• indices/hierarchical-toc.md - Complete documentation tree structure

Directory Guide

Quick API Reference

Most Common APIs

Data Movement:

• nki.isa.dma_copy - Load/store data between HBM and SBUF
• nki.isa.dma_transpose - Transpose during DMA transfer

Tensor Operations:

• nki.isa.nc_matmul - Matrix multiplication on Tensor Engine
• nki.isa.tensor_tensor - Element-wise operations
• nki.isa.tensor_scalar - Broadcast scalar operations
• nki.isa.tensor_reduce - Reduction along axes

Memory Allocation:

• nl.ndarray - Create tensor in SBUF
• nl.zeros - Create zero-initialized tensor

Loop Constructs:

• range - Standard loop iterator (recommended)
• nl.affine_range / nl.sequential_range / nl.static_range - Legacy aliases for range (all have identical effect in NKI 0.3.0+)

SPMD:

• nl.program_id - Get current program index
• nl.num_programs - Get total program count

Module Aliases

The documentation uses these standard aliases:

import nki
import nki.language as nl
import nki.isa as nisa

Activation Function Operators

Important: Symbols like nl.exp, nl.sigmoid, nl.tanh, etc. are op specifiers, not standalone callable functions. They must be passed as the op parameter to nisa.activation():

# CORRECT: Use as op parameter
result = nisa.activation(op=nl.exp, data=tensor, scale=scale_tensor, bias=bias_tensor)

# INCORRECT: These are NOT callable functions
# result = nl.exp(tensor)  # This does NOT work

Hardware Constraints

Critical limits to remember when answering questions:

Hardware Generations

Source NKI Kernel

The Python example files in references/downloads/ use the latest NKI API patterns and can be used as reference implementations:

• average_pool2d_nki_kernels.py — 2D average pooling kernel
• matrix_multiplication_nki_kernels.py — Tiled matrix multiplication kernel
• transpose2d_nki_kernels.py — 2D transpose kernel
• mamba_nki_kernels.py — Fused Mamba state-space model kernel
• test_nki_isa_local_gather.py — Local gather ISA test

Note: The fused_mamba tutorial (references/programming/tutorials/fused_mamba.md) still uses Beta 1 syntax (nl.load/nl.store). The optimization concepts (loop reordering, data reuse, tiling for temporal locality) are valid, but do NOT copy its code patterns — use nisa.dma_copy from /neuron-nki-writing for any code generation.

Common Query Examples

API Lookup

Query: "What is nisa.nc_matmul?" → Read indices/symbol-lookup.md → Find link → Read programming/api/api-nki-isa-tensor.md#nki-isa-nc_matmul

Error Code

Query: "What does NCC_EVRF001 mean?" → Read debugging/error-codes-index.md → Read debugging/error-codes/EVRF001.md

Tutorial

Query: "How do I implement matrix multiplication?" → Read indices/task-routing.md → Read programming/tutorials/matrix_multiplication.md

Architecture

Query: "What's different in Trainium3?" → Read architecture/trainium3_arch.md

Optimization

Query: "How do I profile my kernel?" → Read optimization/use-neuron-profile.md

Response Guidelines

When answering NKI documentation queries:

1. Cite sources - Reference the specific documentation file used
2. Include signatures - For API queries, include function signatures and key parameters
3. Note hardware requirements - Mention if APIs are generation-specific (v2/v3/v4)
4. Link related docs - Point to related tutorials or guides for context
5. Warn about deprecated patterns - If documentation shows old patterns, note this

Error Code Format

NKI compiler errors follow the pattern NCC_<category><number>:

• EOOM* - Out of memory errors
• EVRF* - Verification/validation errors
• EUOC* - Unsupported operation errors
• EBVF* - Backend verification failures
• ESPP* - Shape/type errors

The full index is in debugging/error-codes-index.md with individual files in debugging/error-codes/.

Related Skills

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

serving-systems

This skill bundles the curated reference material for LLM and multimodal serving-system development as a flat library of topic-focused notes under references/. Open the specific reference whose topic matches the task; do not preload everything.

How to use this skill

1. Read this file once to learn what's covered.
2. For the active task, identify the one or two topics that match it (use the index below).
3. Open references/<topic>.md directly with your file-read tool. Each is self-contained.
4. Some topics have follow-up files named references/<topic>-<sub>.md (e.g. cuda-graph-runner.md); the main <topic>.md says when to read them.

Default-on optimizations (NVIDIA serving)

Three techniques every production serving system on NVIDIA ships with — confirm all three are in place before pursuing workload-specific optimizations:

1. Continuous batching — see references/algorithms/continuous-batching.md. Skip only when the workload is single-batch by contract.
2. Fused attention kernel — FlashInfer, FlashAttention, or SDPA. For the picker (workload → backend, plus the per-feature matrix), open references/backends/attention-backend-comparison.md; for deep usage, references/backends/flashinfer.md, references/backends/flashattention.md, or references/backends/sdpa.md. On NVIDIA, never skip a fused kernel.
3. CUDA graphs — see references/backends/cuda-graph.md. Required for low-latency single-batch and high-throughput batched serving.

On AWS Trainium (Neuron) instead

Trainium is not a GPU — the floor is different. Start from references/frameworks/neuron-pytorch.md and references/hardware/aws-trainium.md. The defaults there: static-shape graphs (bucket prompt/decode lengths; neuronx-cc recompiles on new shapes), a persistent compile cache, BF16, on-device sampling, and continuous batching over static buffers. The decisive decode optimization is a device-resident, in-place KV cache — use NxD's KVCacheManager + ModelBuilder aliasing (references/frameworks/nxd-kv-cache.md); a from-scratch torch_neuronx.trace decode that passes the KV cache through the graph boundary every token is host-bound and slow. Flash attention is available here too — via the NKI kernel nki_flash_attn_func (references/frameworks/neuron-flash-attention.md); the FlashAttention/FlashInfer libraries are NVIDIA-only but the algorithm ships on Neuron, and it cuts the attention activation peak that drives the HBM OOM. For custom NeuronCore kernels use the bundled neuron-nki-* skills (the Trainium analog of FlashInfer/Triton kernels).

Reference index

Each entry below is one file under references/. The bracketed phrase shows what triggers it.

Model architectures

• references/models/image-generation.md — Image generation serving — diffusion (U-Net, DiT) and flow-matching.

• references/models/omni-multimodal.md — Omni-modal serving — multi-modality in AND out.

• references/models/speech-generation.md — Speech generation serving — TTS and speech-to-speech.

• references/models/speech-language.md — Speech-language serving — audio understanding (ASR, speech translation, audio-text chat).

• references/models/ssm-hybrid.md — State-space and hybrid SSM+attention serving — Mamba/Mamba2, Jamba, Zamba/Zamba2, Nemotron-H (dense and hybrid), Jet-Nemotron, Falcon-Mamba.

• references/models/text-dense.md — The foundational architecture most modern LLMs build on.

• references/models/text-moe.md — Mixture-of-Experts text decoders for serving — coarse-grained MoE (Mixtral 8x7B, 8x22B) to fine-grained MoE with shared experts (DeepSeek V2/V3/R1, Qwen3-MoE 30B-A3B / 235B-A22B, Llama-4 Scout/Maverick) and DeepSeek's MLA+MoE+MTP stack.

• references/models/video-generation.md — Video generation serving — diffusion-based with 3D attention, temporal+spatial denoising, large activations.

• references/models/vision-language.md — Vision-language serving — LLaVA / LLaVA-NeXT / LLaVA-OneVision (fixed and dynamic tiling), Qwen-VL / 2-VL / 2.5-VL / 3-VL (native-resolution ViT + M-RoPE), InternVL, mllama (Llama-3.2 Vision, cross-attention), Molmo, DeepSeek-VL.

Serving algorithms

• references/algorithms/async-scheduling.md — Hide CPU scheduler / model-runner / Python overhead behind the GPU forward.

• references/algorithms/attention-variants.md — Attention variants in modern serving — three orthogonal axes: head sharing (MHA / MQA / GQA / MLA), masking pattern (causal / bidirectional / sliding-window / cross-attention / 3D / tree), complexity class (quadratic vs SSM / linear / RetNet / RWKV / hybrid).

• references/algorithms/batched-sampling.md — Efficient batched sampling on GPU for LLM serving — temperature, top-p, top-k, min-p, repetition / presence / frequency penalties, typical, rejection sampling for combined filters — without per-request CPU-GPU sync.

• references/algorithms/chunked-prefill.md — Chunked prefill scheduling — split long prompts into smaller token chunks interleaved with decode requests in a single forward pass, preventing a long prefill from stalling decode latency.

• references/algorithms/continuous-batching.md — Implement continuous batching for an LLM inference server.

• references/algorithms/disaggregated-serving.md — Disaggregated prefill / decode (P/D) serving — separate worker pools for prefill and decode stages, with KV cache transferred between them.

• references/algorithms/heterogeneous-kv-cache.md — Memory management and prefix caching for hybrid models (full-attn + sliding-window, attention + SSM/Mamba, attention + linear).

• references/algorithms/moe-routing-dispatch.md — MoE routing and dispatch for serving — top-k gating (softmax, group-limited, auxiliary-loss-free), token-to-expert dispatch/combine (padding, permutation, DeepEP all-to-all), grouped-GEMM expert FFN, Marlin-MoE / CUTLASS-MoE kernels, expert parallelism (EP), expert load balancing (EPLB), shared experts.

• references/algorithms/paged-attention.md — Paged KV cache design for LLM serving — block-based non-contiguous KV storage with a page table per request, allowing dynamic growth without fragmentation.

• references/algorithms/parallelism.md — Parallelism strategies for LLM serving — TP, PP, EP, DP, SP, and combos (TP+EP, DP-attention + EP-MoE, TP+SP, TP+PP, context parallel).

• references/algorithms/quantization-schemes.md — Quantization schemes for LLM serving — FP8 (E4M3/E5M2, per-tensor vs per-channel vs block), INT4 weight-only (AWQ, GPTQ, Marlin, GGUF, petit), INT8, FP4 (MXFP4, NVFP4), FP quant mixed-precision (qutlass / modelopt), weight-only vs.

• references/algorithms/radix-prefix-caching.md — Prefix / radix KV cache reuse — share KV cache pages across requests with common prefixes via a radix tree, LRU eviction.

• references/algorithms/speculative-decoding.md — Speculative decoding for LLM serving — draft proposals verified in one target pass.

• references/algorithms/structured-output.md — Structured output at serving time — grammar-guided decoding (XGrammar, Outlines, llguidance, lm-format-enforcer), JSON mode, regex-constrained output, CFG / pushdown grammars, tool / function calling, logits biasing.

Kernel-library backends

• references/backends/attention-backend-comparison.md — Picking among FlashInfer, FlashAttention, and SDPA — workload-to-backend table, per-feature matrix, plan/run cost on single-batch latency, hardware support, migration paths, common pitfalls.

• references/backends/cuda-graph.md — CUDA graph capture/replay for serving — eliminate per-kernel CPU launch overhead.

• references/backends/flashattention.md — Integrate FlashAttention into a serving engine with explicit KV-cache management.

• references/backends/flashinfer.md — FlashInfer library usage in serving engines.

• references/backends/sdpa.md — Use PyTorch scaled dot product attention (SDPA) as a serving attention backend.

• references/backends/triton-kernels.md — Consuming existing Triton kernels in serving engines — invocation patterns, autotune caching, compile-time constants, warmup, and interaction with CUDA graphs and torch.compile.

Frameworks (PyTorch / Triton / MLX / Neuron)

• references/frameworks/mlx.md — MLX framework for LLM serving on Apple Silicon — unified-memory array model, lazy evaluation and mx.eval, mx.compile, mx.fast kernels, mlx-lm reference serving path, native INT4/INT8 quantization via mx.quantize, custom Metal kernels via mx.fast.metal_kernel.

• references/frameworks/pytorch.md — PyTorch idioms for LLM serving — nn.Module + weight loading, torch.compile (Dynamo / inductor / dynamic shapes / reduce-overhead), state_dict remapping, custom op registration with torch.library, NCCL setup, HF transformers integration, inference_mode.

• references/frameworks/neuron-pytorch.md — PyTorch on AWS Trainium (Neuron) — torch-neuronx vs TorchNeuron Native, neuronx-cc compilation, static-shape bucketing, persistent compile cache, BF16, static KV cache, host-sync pitfalls, when to drop to NKI.

• references/frameworks/nxd-inference.md — NxD Inference (NeuronX Distributed Inference) — AWS's Trainium inference library (NeuronConfig / ModelBuilder / generate); full-turnkey or as building blocks inside a bespoke server.

• references/frameworks/nxd-kv-cache.md — NxD KVCacheManager — the device-resident, in-place KV cache for Trainium (resident nn.Parameter buffers + ModelBuilder input/output aliasing). The key decode optimization; what raw torch_neuronx.trace can't do.

• references/frameworks/neuron-flash-attention.md — flash attention on Trainium via the NKI kernel nki_flash_attn_func. Flash attention is NOT NVIDIA-only — the algorithm ships on Neuron; use it instead of naive softmax(QKᵀ)V to cut the activation peak (and the HBM OOM that caps batch size).

• references/frameworks/triton.md — Triton as a framework-level decision — when a custom Triton kernel pays off in serving vs reusing FlashInfer / FlashAttention / CUTLASS / liger / sgl-kernel.

Hardware specifics

• references/hardware/amd-mi300.md — AMD Instinct MI300-family hardware specs for serving — MI300X, MI325X, MI350X.

• references/hardware/apple-silicon.md — Apple Silicon (M-series SoC) hardware specs for serving — M1 / M2 / M3 / M4 (and Pro / Max / Ultra variants).

• references/hardware/aws-trainium.md — AWS Trainium2 (Trn2) specs for serving — NeuronCore-v3, SBUF/PSUM/HBM hierarchy, compute engines, logical-NeuronCore (LNC) config, what trn2.3xlarge exposes.

• references/hardware/nvidia.md — NVIDIA data-center / workstation GPU specs across 5 generations: Blackwell (B200, RTX PRO 6000), Hopper (H200, H100), Ada (L40S, L4), Ampere (A100 40/80, A10), Turing (T4).

Engine source maps

• references/engines/sglang.md — SGLang serving engine source-code lookup — SRT runtime, scheduler/tokenizer managers, attention backends (FlashInfer, FlashInfer-MLA, CUTLASS-MLA, FA, FlashMLA, Triton), MoE routing/dispatch + EPLB, radix cache + HiCache, quantization, speculative decoding (EAGLE), disaggregated P/D, TP/PP/EP, CUDA graphs, sgl-kernel custom kernels.

• references/engines/trtllm.md — TensorRT-LLM serving engine source-code lookup.

• references/engines/vllm.md — vLLM serving engine source-code lookup.

API / benchmark / profiler tooling

• references/tooling/accuracy-checker.md — Create an accuracy verification script that compares a custom generation implementation against HuggingFace model.generate() as the reference.

• references/tooling/fastapi-serving.md — Create a production-ready FastAPI inference server for HuggingFace transformer models.

• references/tooling/io-handling.md — Request-time I/O handling for LLM/multimodal serving — tokenization and chat templates, tool-call prompt formatting, image preprocessing (resize / tile / normalize / patchify), video frame sampling, audio feature extraction (log-mel), detokenization and UTF-8-safe streaming, tool-call parsing per model family, structured-output extraction.

• references/tooling/lora-serving.md — Multi-adapter LoRA serving — single base model dispatching different LoRA adapters per request at serving throughput.

• references/tooling/openai-api.md — OpenAI-compatible HTTP per modality — text (/v1/completions, /v1/chat/completions), image (/v1/images/generations), TTS (/v1/audio/speech), STT (/v1/audio/transcriptions), video (/v1/videos + async polling), realtime audio (WS /v1/realtime).

• references/tooling/profiler.md — GPU performance profilers for LLM/multimodal serving: PyTorch Profiler (Python-level op aggregation, Chrome trace), Nsight Systems (nsys, system-timeline / launch-gap / NCCL diagnosis), Nsight Compute (ncu, kernel metrics / occupancy / roofline).

• references/tooling/serving-benchmark.md — Benchmark an LLM serving endpoint — TTFT, TPOT, ITL, end-to-end latency, throughput, p50/p95/p99 across concurrency and ISL/OSL sweeps.

Out of scope

Kernel implementation (writing CUDA / Triton / CUTLASS). For that, use the separate agent-gpu-skills collection. Exception — NKI: writing NeuronCore kernels for AWS Trainium is in scope here, via the bundled neuron-nki-* skills (neuron-nki-writing, -docs, -debugging, -profiling, -profile-querying); there is no separate Trainium kernel collection.

Reference repos

The repos/ directory (excluded from materialization to agents) holds full source trees of vLLM, SGLang, and TensorRT-LLM as git submodules. Engine-source-map references in this skill cite paths like $SERVE_REPOS/<engine>/...; export SERVE_REPOS=$(git rev-parse --show-toplevel)/skills/serving-systems/repos or substitute inline.

Profile Querying

Run SQL queries against NKI kernel profile data using neuron-explorer view. This ingests NEFF+NTFF into parquet and exposes a DuckDB-backed API server on localhost. No deployment, no remote service — just the CLI and curl.

For more advanced analysis, use python on parquet to compute performance bounds and investigate precise inefficiencies within arbitrary execution intervals.

What you need: A compiled NEFF file and a captured NTFF trace file. These come from /neuron-nki-profiling or from running a kernel with the right env vars and neuron-explorer capture.

Quick Start

# Ingest and start API server (no web UI)
neuron-explorer view \
  -n ./kernel.neff \
  -s ./profile.ntff \
  --data-path ~/.local/share/neuron-profile \
  --display-name my-kernel \
  --disable-ui &

# Wait for server
sleep 10

# Query
curl -s -X POST http://localhost:3002/api/v1/db/my-kernel/_search \
  -H 'Content-Type: application/json' \
  -d '{"type":"databaseExplorerQuery","tableName":"Summary","query":"SELECT total_time, mfu_estimated_percent, tensor_engine_active_time_percent, dma_active_time_percent FROM Summary"}'

That's it. Ingest, serve, query.

Prerequisites

• neuron-explorer installed (comes with AL2023 DLAMI or aws-neuronx-tools)
• NEFF file (compiled kernel binary) + NTFF file (execution trace)

Check availability:

which neuron-explorer && neuron-explorer --version

If not found, check /opt/aws/neuron/bin/neuron-explorer.

Step-by-Step Workflow

Step 0: Check Profile Quality (Re-profile if Needed)

Note: This step is specific to NKI kernel development. If you are querying a profile that was generated outside of an NKI workflow, skip to Step 1.

Disclaimer: Query results are only as good as the profile. If the NEFF/NTFF were captured without the right env vars, key tables (DmaPacket, DmaPacketAggregated) may be empty and source-level attribution will be missing.

Check whether the profile has the data you need:

# After ingesting (Step 2), check for DMA packet data
curl -s -X POST http://localhost:3002/api/v1/db/${PROFILE_NAME}/_search \
  -H 'Content-Type: application/json' \
  -d '{"type":"databaseExplorerQuery","tableName":"DmaPacket","query":"SELECT COUNT(*) as cnt FROM DmaPacket"}'

If cnt is 0 or the table is missing, the profile was captured without DGE notifications. If bir_debug_info_source_location is NULL on all Instruction rows, the NEFF was compiled without debug info.

To re-profile for best results, set these env vars in the kernel script before any neuron imports, then re-run and re-capture:

import os
os.environ["XLA_IR_DEBUG"] = "1"
os.environ["XLA_HLO_DEBUG"] = "1"
os.environ["NEURON_FRAMEWORK_DEBUG"] = "1"
os.environ["NEURON_RT_VISIBLE_CORES"] = ... # Restrict available cores when running experiments in parallel.
os.environ["NEURON_RT_INSPECT_ENABLE"] = "1"
os.environ["NEURON_RT_INSPECT_DEVICE_PROFILE"] = "1"
os.environ["NEURON_RT_INSPECT_SYSTEM_PROFILE"] = "0"
os.environ["NEURON_RT_INSPECT_OUTPUT_DIR"] = ... # This is for the NEFF generation if needed. NTFF will go to the -s capture path in the next command. 

Then re-capture with DGE notifications enabled:

NEFF_PATH=$(find ./output -name "*.neff" | head -1)
NEURON_RT_ENABLE_DGE_NOTIFICATIONS=1 neuron-explorer capture \
  -n "$NEFF_PATH" \
  -s profile.ntff \
  --profile-nth-exec=2

With --profile-nth-exec=2, the output file is profile_exec_2.ntff (not profile.ntff), written to the directory specified by the -s flag.

If the existing profile has the data you need, skip this step entirely.

Another thing to look out for is running torch functions on device like randomnly generating inputs. This will be fused into the kernel execution and obfuscate it's profile. Move those commands off device if you want to isolate kernel execution.

Step 1: Ingest and Start Server

If you want to run SQL queries against the Neuron Explorer DuckDB engine, use the view command with --disable-ui to start the server.

Set variables:

NEFF_PATH=<resolved neff path>
NTFF_PATH=<resolved ntff path>
PROFILE_NAME=<descriptive name, e.g. "my-matmul">
NE_DATA_PATH=~/.local/share/neuron-profile

Check if the neuron-explorer server is already running:

curl -s http://localhost:3002/api/v1/health

If the server is already running or if you are running python directly on the parquet, use --ingest-only in the following command instead of --disable-ui.

neuron-explorer view \
  -n "$NEFF_PATH" \
  -s "$NTFF_PATH" \
  --data-path "$NE_DATA_PATH" \
  --display-name "$PROFILE_NAME" \
  --disable-ui \
  > /tmp/neuron-explorer-${PROFILE_NAME}.log 2>&1 &
NE_PID=$!
echo "neuron-explorer started (PID: $NE_PID), waiting for API..."

The command may fail on an conflicting port from the existing server but the ingestion may have still succeeded. If so, check for Processing for ... is complete before the error message or rerun with --ingest-only.

Wait for API:

for i in $(seq 1 60); do
  if curl -s http://localhost:3002/api/v1/health 2>/dev/null | grep -q healthy; then
    echo "API server ready"
    break
  fi
  sleep 1
done

Step 2: Read the Schema

Before writing any queries, check the table docs in references/schema/ for interpretive guidance on the most commonly used tables. For any table not covered there, query its schema directly from neuron-explorer using the following commands:

curl -s -X POST http://localhost:3002/api/v1/db/${PROFILE_NAME}/_search \
  -H 'Content-Type: application/json' \
  -d '{"type":"tableSchema","tableName":"Instruction"}' | python3 -m json.tool

List all available tables:

curl -s -X POST http://localhost:3002/api/v1/db/${PROFILE_NAME}/_search \
  -H 'Content-Type: application/json' \
  -d '{"type": "listDbFiles"}' | python3 -m json.tool

Step 3a: Execute SQL Queries

Use databaseExplorerQuery for arbitrary SQL (SELECT only).

Summary metrics — which engine is the bottleneck?

curl -s -X POST http://localhost:3002/api/v1/db/${PROFILE_NAME}/_search \
  -H 'Content-Type: application/json' \
  -d '{"type":"databaseExplorerQuery","tableName":"Summary","query":"SELECT total_time, mfu_estimated_percent, tensor_engine_active_time_percent, vector_engine_active_time_percent, dma_active_time_percent, hbm_read_bytes, hbm_write_bytes FROM Summary"}' | python3 -m json.tool

Instruction breakdown — what is each engine doing and waiting on?

curl -s -X POST http://localhost:3002/api/v1/db/${PROFILE_NAME}/_search \
  -H 'Content-Type: application/json' \
  -d '{"type":"databaseExplorerQuery","tableName":"Instruction","query":"SELECT engine, opcode, COUNT(*) as cnt, SUM(duration_ns) as total_dur_ns, SUM(evt_wait_time_ns) as total_evt_wait_ns FROM Instruction GROUP BY engine, opcode ORDER BY total_dur_ns DESC"}' | python3 -m json.tool

NKI source line hotspots — which lines of the kernel are slowest?

curl -s -X POST http://localhost:3002/api/v1/db/${PROFILE_NAME}/_search \
  -H 'Content-Type: application/json' \
  -d '{"type":"databaseExplorerQuery","tableName":"Instruction","query":"SELECT bir_debug_info_source_location, engine, opcode, COUNT(*) as cnt, SUM(duration_ns) as total_dur_ns FROM Instruction WHERE bir_debug_info_source_location IS NOT NULL GROUP BY bir_debug_info_source_location, engine, opcode ORDER BY total_dur_ns DESC LIMIT 10"}' | python3 -m json.tool

Step 3b: Python DuckDB on Parquet

For SQL queries without the API server, use DuckDB's Python bindings directly on the parquet files. After ingestion (Step 1), the data lives at:

<data-path>/profiles/global/<display-name>@latest/<Table>.parquet

import duckdb

NE_DATA_PATH = "~/.local/share/neuron-profile"
PROFILE = "my-kernel"
PARQUET_DIR = f"{NE_DATA_PATH}/profiles/global/{PROFILE}@latest"

con = duckdb.connect()

# Load tables directly from parquet
con.execute(f"CREATE VIEW Instruction AS SELECT * FROM '{PARQUET_DIR}/Instruction.parquet'")
con.execute(f"CREATE VIEW DmaPacket AS SELECT * FROM '{PARQUET_DIR}/DmaPacket.parquet'")

# Example: measure LDWEIGHTS/MATMUL temporal overlap
result = con.execute("""
    SELECT
        CASE WHEN lw.end_ts <= mm.start_ts THEN 'lw_before'
             WHEN lw.start_ts >= mm.end_ts THEN 'lw_after'
             ELSE 'overlap' END as rel,
        COUNT(*) as cnt
    FROM Instruction mm
    JOIN Instruction lw ON mm.bir_id = lw.bir_id
    WHERE mm.opcode = 'MATMUL' AND lw.opcode = 'LDWEIGHTS'
      AND mm.tensor_instruction_type = 'REGULAR'
      AND lw.tensor_instruction_type = 'REGULAR'
    GROUP BY rel
""").fetchdf()
print(result)

Step 3c: Pandas on Parquet

For analyses that require Python computation — interval merges, custom metrics, numpy operations, or cross-table joins with arbitrary logic — load the parquet files directly with pandas.

import pandas as pd, numpy as np, os

NE = os.path.expanduser("~/.local/share/neuron-profile/profiles/global")
profile = "my-kernel"
d = f"{NE}/{profile}@latest"

# Load tables
summary  = pd.read_parquet(f"{d}/Summary.parquet").iloc[0]
inst     = pd.read_parquet(f"{d}/Instruction.parquet")
active   = pd.read_parquet(f"{d}/ActiveTime.parquet")
metadata = pd.read_parquet(f"{d}/Metadata.parquet").iloc[0]
dma_pkts = pd.read_parquet(f"{d}/DmaPacket.parquet")
dma_agg  = pd.read_parquet(f"{d}/DmaPacketAggregated.parquet")
tensors  = pd.read_parquet(f"{d}/TensorInfo.parquet")
flow     = pd.read_parquet(f"{d}/Flow.parquet")

This is the approach used by the performance bounds computation and all investigations in the Profile Analysis workflow.

Step 4: Interpret Results

Only claim what the data shows. Profile data is precise but narrow — it tells you what happened, not always why.

• Don't diagnose from single metrics. A query result is a measurement, not a conclusion. Low utilization, high wait times, or large byte counts need context from other tables before they mean anything.
• Don't assume field names mean what they sound like. Some fields are unpopulated or misleading for NKI kernels. Check references/schema/ before building conclusions on a field you haven't validated.
• Don't compare engines without interval merging. Instructions overlap within an engine (pipelining) and across engines (parallelism). Raw sums from the Instruction table overstate wall-clock time. Use ActiveTime for wall-clock comparisons.
• Don't skip the data quality check. If DmaPacketAggregated is missing or bir_debug_info_source_location is mostly NULL, the query results are incomplete — re-profile before interpreting.

Step 8: Cleanup

kill $NE_PID 2>/dev/null

Profile Analysis

If you are asked for analysis of the profile, follow this workflow. All logic — bound definitions, gap interpretation, and investigation selection — lives in performance-bounds.md.

1. Calculate bounds

Follow the "The bounds" section of performance-bounds.md to compute all three families (memory, compute, pipeline). These require Python on parquet (Step 3c).

2. Identify the dominant gaps

Follow "Reading the gaps" in performance-bounds.md. Compute each consecutive-pair gap within the memory and compute families, plus the pipeline gap. Report all gaps and their sizes relative to total_time.

3. Run investigations

Follow "From bounds to investigations" in performance-bounds.md. Use the bottleneck engine and gap sizes to select which investigation groups to run. Each investigation has a Step 1 (detect and quantify) and Step 2 (localize to NKI source lines). Run all relevant investigations — a kernel typically has multiple active inefficiencies.

4. Report

Present a single summary:

• Bounds table: all bounds with values and the gap between each pair. Also report each engine's total time pointing out the largest one(s) as the bottleneck(s). If neither DMA nor Tensor Engine is the bottleneck, explain which engine is the bottleneck and that supporting it is still WIP.
• Per-investigation findings: gap size, source lines responsible, and their contributions. Include investigations that found nothing so the analysis is visibly complete.

Order the presented inefficiencies and investigation findings according to it's relevance to the bottlenecks and the measured gaps.

4. Follow up (After an optimization step/attempt)

After an optimization step or attempt, investigate the new profile to identify exactly what improved or regressed. Follow the full process and present a side by side report of all of the bounds and engine times as well as the new investigation findings. Highlight changes but do not over-interpret, only relay what the evidence shows. Static code analysis is faulty, you will be tempted to over-intepret the causes and effects, DON'T (unless EXPLICITELY) asked to.

Worked Examples

For end-to-end examples of profile-guided optimization, see:

Multi-Kernel Querying

All profiles sharing the same --data-path are served by one server. Each profile is queried by its --display-name.

NE_DATA_PATH=~/.local/share/neuron-profile

neuron-explorer view -n $NEFF_A -s $NTFF_A --data-path "$NE_DATA_PATH" --display-name kernel-a --disable-ui &
neuron-explorer view -n $NEFF_B -s $NTFF_B --data-path "$NE_DATA_PATH" --display-name kernel-b --disable-ui &

# Query either through same server
curl localhost:3002/api/v1/db/kernel-a/_search ...
curl localhost:3002/api/v1/db/kernel-b/_search ...

For batch ingestion without a server, use --ingest-only instead of --disable-ui. It writes parquet and exits. Any future server on the same data-path discovers the ingested profiles.

Parquet lands at <data-path>/profiles/global/<display-name>@latest/.

Port Conflicts

If port 3002 is already in use, ingestion still succeeds — parquet is written to disk before the server attempts to bind.

lsof -i :3002 | head -5

• If it's neuron-explorer on the same data-path: reuse it — it discovers newly ingested profiles automatically.
• If it's something else: use --api-server-port 4002 (or any free port).

Important Notes

• Use neuron-explorer not neuron-profile for all capture and view commands.
• DGE notifications are required for DMA packet-level tables (DmaPacket, DmaPacketAggregated). Set NEURON_RT_ENABLE_DGE_NOTIFICATIONS=1 in the environment before neuron-explorer capture. Do NOT rely on the CLI flag — use the env var directly.
• Always pass --data-path explicitly.
• The API server binds to localhost only.
• Only SELECT queries are supported via the API.
• --display-name becomes the profile identifier in API URLs.
• --disable-ui skips the web UI (port 3001) but starts the API server (port 3002).
• --ingest-only writes parquet and exits — no server at all.

Related Skills

Writing NKI Kernels

This skill guides writing and modifying NKI (Neuron Kernel Interface) kernels — from new kernel creation (PyTorch/NumPy/natural language translation) to editing existing kernels (adding shape/dtype support, refactoring tiling, implementing new features). Focus on correctness using documented APIs.

Critical: NKI Language Constraints

BEFORE writing any NKI code, read references/nki-language-constraint.md for the complete list of required and forbidden API patterns covering Beta 1 → Beta 2, Beta 2 → NKI 0.3.0, and NKI 0.3.0 → NKI 0.4.0 migration rules. Violating ANY rule is a compilation failure.

Quick Start

Minimal working kernel structure:

import nki
import nki.isa as nisa
import nki.language as nl

@nki.jit
def my_kernel(input_hbm: nl.ndarray) -> nl.ndarray:
    """One-line description of kernel operation."""
    # 1. Allocate SBUF tile
    tile = nl.ndarray(input_hbm.shape, dtype=input_hbm.dtype, buffer=nl.sbuf)

    # 2. Load from HBM to SBUF
    nisa.dma_copy(dst=tile, src=input_hbm[0:input_hbm.shape[0], 0:input_hbm.shape[1]])

    # 3. Compute (example: element-wise exp)
    result = nl.ndarray(tile.shape, dtype=tile.dtype, buffer=nl.sbuf)
    nisa.activation(dst=result, data=tile, op=nl.exp)

    # 4. Allocate and store to HBM
    output = nl.ndarray(input_hbm.shape, dtype=input_hbm.dtype, buffer=nl.shared_hbm)
    nisa.dma_copy(dst=output, src=result)

    return output

Complexity Assessment

Before reading references, classify the task to avoid unnecessary overhead:

Simple (element-wise op, single reduction, activation, layernorm, add/multiply):

• Use the Quick Start template and Step 4 API table directly
• Skip utility library references entirely
• Start writing code immediately — consult references only when stuck
• Target: working kernel within 5 minutes

Medium (matmul, softmax, multi-step fusion, transpose with tiling):

• Read references/common-patterns.md, references/api-translation.md and references/memory-patterns.md
• Skip utility library references unless tiling is complex
• Target: working kernel within 15 minutes

Complex (multi-head attention, transformer blocks, state-space models, MoE):

• Full reference loading appropriate
• Read utility selection guide and relevant patterns
• Target: working kernel within 30 minutes

Start writing code as soon as possible. Reference reading should supplement coding, not precede it. Write the kernel structure first, then consult references for specific API details as needed.

Translation Workflow

Step 1: Identify Operations

Map PyTorch/NumPy operations to NKI equivalents using references/api-translation.md.

Step 2: Design Tiling Strategy

NKI operates on tiles with hardware constraints:

For tensors exceeding limits, use explicit tiling with TiledRange for remainder-safe iteration (see Utility Selection Guide below).

Step 3: Implement Memory Access

Consult the Utility Selection Guide to choose the right utilities for the kernel's access patterns, then follow references/memory-patterns.md and references/transpose-and-layout.md:

• Contiguous DMA: For aligned, sequential access (most efficient)
• Strided DMA: Use TensorView.slice(step=N) for gather/scatter patterns (see transpose-and-layout.md)
• Tiling loops: Use TiledRange for dimensions requiring tiling with remainder handling
• Transpose operations: Use nisa.nc_transpose() for P↔F transpose, TensorView for layout manipulation
• Partition broadcast: Use stream_shuffle_broadcast when a bias/scale in partition 0 must reach all PEs
• Buffer management: Use SbufManager when the kernel has 4+ SBUF allocations or shared sub-functions

Step 4: Add Compute Operations

Use ISA functions with explicit dst parameter:

# Element-wise
nisa.activation(dst=result, data=input, op=nl.exp)
nisa.tensor_tensor(dst=result, data1=a, data2=b, op=nl.add)
nisa.tensor_scalar(dst=result, data=input, op0=nl.multiply, operand0=2.0)

# Reductions
nisa.tensor_reduce(dst=result, data=input, op=nl.add, axis=1)

# Matrix multiply
nisa.nc_matmul(dst=psum_result, stationary=a, moving=b)

Step 5: Validate Complex Translations

Important: Always compute reference results on CPU, not on the XLA device. Every on-device XLA graph generates a separate NEFF file, making it hard to identify the NKI kernel's NEFF during profiling. Use tensor.cpu() before computing PyTorch/NumPy references.

For simple kernels (single operation, few tiles), comparing the final output against a PyTorch/NumPy reference is usually sufficient. For complex kernels with multiple computation stages, validate incrementally:

1. Identify logical stages in the source operation. For example, a fused attention kernel has: QK matmul → scale → mask → softmax → AV matmul.
2. Translate and validate one stage at a time. Write the first stage, store its output to HBM, and compare against the corresponding PyTorch intermediate. Only proceed to the next stage once the current one matches.
3. Compose validated stages. Once each stage is verified independently, connect them (keeping intermediates in SBUF instead of round-tripping through HBM) and validate the final output.

This catches translation errors early — a mismatch in the final output of a 5-stage kernel is much harder to diagnose than a mismatch after stage 2.

Use multiple complementary checks (atol/rtol, max absolute difference, tensor norm of the difference, cosine similarity) rather than relying on a single metric.

Hardware Constraints Quick Reference

Loop Types

Common Patterns

For detailed code examples, anti-patterns, and production patterns (cumsum, rmsnorm_quant), see references/common-patterns.md.

Element-wise Operations

• Reshape to 2D, tile P dimension (≤128), use nisa.activation() / nisa.tensor_tensor()

Matrix Multiply — Key Rules

• Allocate PSUM: nl.ndarray(..., buffer=nl.psum) — uninitialized is correct
• K-dimension loop: always nl.affine_range(), never nl.sequential_range (serializes execution)
• Multiple nisa.nc_matmul() writes to same PSUM buffer triggers hardware accumulation
• Never write PSUM to HBM between accumulation steps
• Operands: stationary [K≤128, M≤128], moving [K≤128, N≤512/4096], result [M, N] in PSUM
• Copy PSUM→SBUF via nisa.tensor_copy() before further ops
• See examples/simple_matmul.py for complete examples

Fused ScalarE Operations

• nisa.activation(op=nl.exp, data=x, scale=s) → exp(x * s) in one instruction
• Available: scale, bias, or both before any activation function

Sequential Operations & Associative Scan

• Use nisa.tensor_tensor_scan instead of explicit sequential loops
• Pattern: out[i] = op0(data[i], out[i-1]) op1 data1[i]
• Multi-tile: pass final state as initial= to next tile's scan
• See examples/associative_scan.py for complete pattern

Skill References

References are tiered to minimize overhead on simple tasks. Load only what you need based on the Complexity Assessment above.

Always load (core references):

• references/nki-language-constraint.md - MANDATORY: Required and forbidden API patterns for NKI 0.4.0, reference kernel template
• references/common-patterns.md - Full code examples: matmul PSUM accumulation, fused ScalarE, associative scan, production patterns
• references/api-translation.md - PyTorch/NumPy to NKI operation mapping
• references/kernel-template.md - Standard kernel template with self-contained utilities
• references/indexing-patterns.md - Complete indexing guide: memory-type rules (HBM/SBUF/PSUM), operation constraints (matmul/transpose/reduce), dynamic indexing with DGE modes

Load when tiling or DMA patterns are needed (medium+ complexity):

• references/memory-patterns.md - DMA and tiling patterns with code examples
• references/nkilib/core/tiled-range.md - TiledRange: dimension tiling with remainder handling
• references/nkilib/core/kernel-helpers.md - Math helpers, SPMD, dtype utilities

Load when layout manipulation is needed:

• references/transpose-and-layout.md - Transpose and layout transformation guide: nc_transpose, TensorView, array patterns, strided DMA, decision trees for layout operations
• references/nkilib/core/tensor-view.md - TensorView: zero-copy tensor manipulation

Load when advanced patterns are needed (complex kernels only):

• references/performance-basics.md - Optimization patterns (fusion, double buffering)
• references/nkilib/core/allocator.md - SbufManager: stack/heap SBUF allocation
• references/nkilib/core/tile-info.md - TiledDimInfo: tile tracking with subtile support
• references/nkilib/ops/ - Copy/broadcast operations (stream-shuffle, tp-broadcast)
• references/nkilib/types/ - Enum types and logging utilities
• references/nkilib/patterns/ - Reusable kernel patterns (quantization, normalization, layout conversion, MoE)

Bundled Source Code

Full source for nkilib/core utilities and subkernels:

• references/nkilib/core/utils/ - Utility source (TensorView, TiledRange, kernel_helpers, SbufManager, etc.)
• references/nkilib/core/subkernels/ - Reusable sub-operations (RMSNorm, LayerNorm, normalization utils)

Configuration

Default: inline nkilib utility source directly into the user's kernel file from the bundled source above. If nkilib is installed in the user's environment, use from nkilib.core.utils.X import Y imports instead.

Utility Selection Guide

Always Use

Use When Pattern Matches

Specialized: TiledDimInfo (subtile metadata), tp_broadcast (P→F broadcast, very rare).

Decision Flowchart

Kernel needs tiling?
├─ Yes → Use TiledRange for each tiled dimension
│        Use div_ceil() for tile count computations
├─ Nested tiles (subtiles within tiles)?
│  └─ Yes → TiledRange supports nesting: TiledRange(outer_tile, subtile_size)
│           For CTE-style metadata tracking → also consider TiledDimInfo
└─ No → Plain nl.affine_range()

Kernel accesses tensors non-contiguously?
├─ Strided/interleaved → TensorView.slice(step=N)
├─ Broadcasting → TensorView.broadcast(dim, size)
├─ Reshape without copy → TensorView.reshape_dim() / flatten_dims()
├─ Dynamic expert selection → TensorView.select(dim, scalar_offset)
└─ Simple contiguous → plain tensor[start:end, :]

Kernel needs to broadcast a scalar/vector to all partitions?
├─ 1D value in partition 0 → stream_shuffle_broadcast(src, dst)
└─ Column vector (P-dim) to row (F-dim) → tp_broadcast(src, dst, ...)

Kernel allocates many SBUF buffers?
├─ 4+ buffers, or sub-functions share SBUF → SbufManager
└─ 1-3 simple buffers → plain nl.ndarray(..., buffer=nl.sbuf)

Coding Conventions

Follow these conventions unless the user's instructions or existing project style indicate otherwise.

• Prefer kernel_assert() over Python assert - Produces structured error messages ([NCC_INKI016] Kernel validation exception: ...) that clearly identify errors originating from NKI kernels, which is helpful when kernels run inside larger frameworks. The kernel template (references/kernel-template.md) provides an inline definition; alternatively, import from nkilib if installed.
• Include docstrings with Args/Returns/Notes sections for non-trivial kernels
• Validate inputs with shape checks before the kernel body
• Use descriptive variable names (e.g., partition_idx not p)

Kernel Efficiency Guidelines

These are basic efficiency practices to follow when writing any kernel. They do not require advanced allocation or pipelining — just sensible layout, tiling, and data flow choices.

1. Tensor Layout Flexibility (Conditional)

If the input/output tensor layout would make the kernel significantly harder to write (e.g., requiring many strided DMAs or complex reshaping), ask the user:

"The current tensor layout [describe issue] would require [strided DMA / reshaping]. Would you allow changing the layout to [proposed layout]? This would simplify the kernel and improve performance."

Layout changes that typically help:

• Putting the reduction dimension last (contiguous in memory)
• Aligning dimensions to 128 (partition) and 512 (PSUM free)
• Transposing to avoid strided DMA patterns

2. Large Contiguous Free Dimension in DMA (≥2KB)

DMA efficiency depends on the free dimension (second dimension onwards) being large and contiguous in memory. Target ≥2KB contiguous free dimension to saturate memory bandwidth.

Key concept: In a tile [P, F], the partition dimension P (first dim) is distributed across hardware partitions. The free dimension F (second dim onwards) should be large and contiguous.

What "contiguous" means: The free dimension elements are adjacent in HBM memory with stride=1.

Production example from mlp_tkg_gate_up_projection.py:169-181:

# Weight layout: [H, I] where I is the contiguous free dimension
# Load weight tile [HTile=2048, I] where I is large and contiguous
# HTile = 2048 for non-quantized (reshapes to [128, 16, I] in SBUF)

nisa.dma_copy(
    dst=weight_tiles[weight_idx][0:H0, 0:h1_tiles, 0:I],  # SBUF: [128, h1_tiles, I]
    src=unsharded_weight.ap(
        pattern=[
            [H1 * unsharded_weight.shape[1], H0],  # Partition dim
            [unsharded_weight.shape[1], h1_tiles], # Batch dim
            [1, I],                                # Free dim - contiguous (stride=1)
        ],
        offset=h_offset * dims.I + weight_shard_offset,
    ),
)
# Here I (intermediate dimension) is the large contiguous free dimension

Why this matters: Strided DMA (free dim not contiguous) has significant overhead. If your tensor layout requires strided access, consider asking the user to change the layout (see Tensor Layout Flexibility above).

3. Keep Intermediates in SBUF

Avoid unnecessary HBM round-trips by keeping intermediate results in SBUF between operations.

Common pattern: MatMul → Element-wise → HBM

# MatMul result in PSUM
psum_result = nl.ndarray((P, F), dtype=nl.float32, buffer=nl.psum)
nisa.nc_matmul(dst=psum_result, stationary=a, moving=b)

# Copy to SBUF for element-wise ops (PSUM → SBUF, no HBM)
sbuf_result = nl.ndarray((P, F), dtype=nl.float32, buffer=nl.sbuf)
nisa.tensor_copy(dst=sbuf_result, src=psum_result)

# Element-wise activation in SBUF (still no HBM)
nisa.activation(dst=sbuf_result, data=sbuf_result, op=nl.gelu)

# Only final result written to HBM
nisa.dma_copy(dst=output_hbm, src=sbuf_result)

Anti-pattern to avoid:

# BAD: Writing matmul result to HBM, then reading back for activation
nisa.dma_copy(dst=hbm_temp, src=psum_result)      # Unnecessary write
nisa.dma_copy(dst=sbuf_for_act, src=hbm_temp)     # Unnecessary read
nisa.activation(dst=sbuf_for_act, data=sbuf_for_act, op=nl.gelu)

4. Maximize Hardware Parallelism (P=128)

Always try to use the full partition dimension (128) for hardware parallelism.

Production example from mlp_tkg_constants.py:156:

# Hardware partition dimension constraint - always use 128
_pmax = nl.tile_size.pmax  # Max partition dimension in SBUF = 128

# Derived dimensions maximize partition usage
H0 = _pmax  # 128 (partition dimension)
H1 = H // H0  # Remaining hidden dimension

# All SBUF tiles use full partition dimension
tile = nl.ndarray((H0, free_dim), dtype=dtype, buffer=nl.sbuf)  # [128, ...]

5. Minimum Tile Sizes

Production MatMul example from mlp_tkg_gate_up_projection.py:188-204:

# Standard matmul tile: stationary [128, T], moving [128, 512]
for i_tiles in TiledRange(I, dims._psum_fmax):  # _psum_fmax = 512
    nisa.nc_matmul(
        dst=result_psums[i_tiles.index][
            nl.ds(dims.column_tiling_dim * column_idx, T),
            0 : i_tiles.size,  # Up to 512
        ],
        stationary=hidden.ap(
            pattern=[[T * H1, H0], [H1, T]],
            offset=h_offset + column_tile_offset,
        ),
        moving=weight_tiles[weight_idx][
            0:H0,  # 128 partition dim
            column_tile_offset,
            nl.ds(i_tiles.start_offset, i_tiles.size),  # Up to 512
        ],
    )

Vector/Scalar tile sizing from mlp_tkg_constants.py:194-206:

# column_tiling_dim sets the free dimension for vector/scalar ops
# (e.g., activation functions, element-wise ops after matmul)
# Minimum 64 ensures efficient hardware utilization
if T <= 32:
    column_tiling_dim = 32  # Small T: use 32
elif T <= 64:
    column_tiling_dim = 64  # Medium T: use 64 (minimum for efficiency)
else:
    column_tiling_dim = 128  # Large T: use 128

Related Skills