This document is relevant for: Inf2, Trn1, Trn2, Trn3

nrtpy model API#

This page documents NrtpyModel and BenchmarkResult — the classes you use to load a compiled NEFF file, run it on a NeuronCore, and benchmark its execution.

NrtpyModel#

class nrtpy.NrtpyModel#

A wrapper class for executing compiled kernels from NEFF files.

input_tensors_info: dict[str, TensorMetadata]#

Mapping of input tensor name to metadata (shape, size). Use this to discover the expected input names and sizes for a loaded NEFF:

model = NrtpyModel.load_from_neff("model.neff")
for name, info in model.input_tensors_info.items():
    print(f"{name}: shape={info.shape}, size={info.size}")
output_tensors_info: dict[str, TensorMetadata]#

Mapping of output tensor name to metadata (shape, size).

classmethod load_from_neff(neff_path, name=None, core_id=0, cc_enabled=False, rank_id=0, world_size=1)#

Load a NEFF file and return an NrtpyModel instance.

Parameters:

  • neff_path (pathlib.Path or str) – Path to the NEFF file to load.

  • name (str or None) – Optional name for the model. If None, the NEFF filename stem is used.

  • core_id (int) – Target NeuronCore for the model (default 0).

  • cc_enabled (bool) – Enable collective communication.

  • rank_id (int) – Rank of this model within the collective group.

  • world_size (int) – Total number of ranks in the collective group.

Returns:

An NrtpyModel wrapping the loaded model.

Return type:

nrtpy.NrtpyModel

__call__(inputs, outputs=None, save_trace=False, ntff_name=None)#

Execute the model. Invoke by calling the model instance directly: model(inputs={...}).

Parameters:

  • inputs (dict[str, nrtpy.NrtpyTensor]) – Mapping of input tensor name to NrtpyTensor. Keys must match the NEFF input names.

  • outputs (dict[str, nrtpy.NrtpyTensor] or None) – Mapping of output tensor name to NrtpyTensor. Keys must match the NEFF output names. If None, output tensors are allocated automatically and returned.

  • save_trace (bool) – Whether to save an execution trace (.ntff file).

  • ntff_name (str or None) – Optional path for the trace file. If None, the trace is written next to the NEFF file with an .ntff suffix.

Returns:

The auto-allocated output tensors when outputs is None; otherwise None.

Return type:

dict[str, nrtpy.NrtpyTensor] or None

benchmark(inputs, outputs=None, warmup_iter=5, benchmark_iter=5, mode='device')#

Benchmark model execution and return timing statistics.

Parameters:

  • inputs (dict[str, nrtpy.NrtpyTensor]) – Mapping of input tensor name to NrtpyTensor.

  • outputs (dict[str, nrtpy.NrtpyTensor] or None) – Mapping of output tensor name to NrtpyTensor. If None, output tensors are allocated automatically.

  • warmup_iter (int) – Number of warmup iterations before timing.

  • benchmark_iter (int) – Number of timed iterations. Must be at least 1.

  • mode (str) –

    Timing mode:

    • "device" — NeuronCore execution time via device-side tracing. Most accurate for kernel timing.

    • "host" — Host wall-clock time including host-device overhead. No tracing overhead.

Returns:

The benchmark statistics.

Return type:

nrtpy.BenchmarkResult

Raises:

  • ValueError – If mode is not "device" or "host", or if benchmark_iter is less than 1.

  • RuntimeError – In "device" mode, if no device execution events are captured.

BenchmarkResult#

class nrtpy.BenchmarkResult#

Result of a model benchmark run, returned by benchmark(). All timing values are in milliseconds.

mean_ms: float#

Mean execution time across timed iterations.

min_ms: float#

Minimum execution time.

max_ms: float#

Maximum execution time.

std_dev_ms: float#

Standard deviation of execution time.

iterations: int#

Number of timed iterations reflected in the statistics.

warmup_iterations: int#

Number of warmup iterations performed before timing.

durations_ms: list[float]#

Per-iteration execution times.

Examples#

Basic execution#

import numpy as np
from nrtpy import NrtpyModel, NrtpyTensor

# Load a compiled NEFF model
model = NrtpyModel.load_from_neff("path/to/model.neff")

# Build an input tensor from a NumPy array
input_data = np.random.randn(1, 10).astype(np.float32)
input_tensor = NrtpyTensor.from_numpy(input_data, name="input")

# Execute — output tensors are allocated automatically and returned
outputs = model(inputs={"input": input_tensor})
result = outputs["output"].numpy()

Benchmarking#

# Device-side timing (most accurate for kernel latency)
stats = model.benchmark(
    inputs={"input": input_tensor},
    warmup_iter=5,
    benchmark_iter=100,
)
print(f"Mean: {stats.mean_ms:.2f} ms")
print(f"Min: {stats.min_ms:.2f} ms, Max: {stats.max_ms:.2f} ms")
print(f"Std dev: {stats.std_dev_ms:.4f} ms")

# Host-side timing (includes host-device overhead)
host_stats = model.benchmark(
    inputs={"input": input_tensor},
    warmup_iter=5,
    benchmark_iter=100,
    mode="host",
)