Profiling TorchNeuron (Part 1): System-level profiling on AWS Trainium.

Many ML practitioners have some intuition about how to make models run fast, but fewer know how to analyze why a model is slow. In Part 1 of this series, you will learn how to systematically analyze inference performance and identify where the time is spent.

This post starts from scratch and assumes no prior profiling experience. The walkthrough uses AWS Trainium, a purpose-built AI chip for deep learning training and inference workloads, and TorchNeuron , its open source native PyTorch backend. TorchNeuron supports standard PyTorch workflows on Trainium, including eager execution, torch.compile, and standard PyTorch APIs, with minimal code changes. Additionally, TorchNeuron integrates with Neuron Kernel Interface (NKI) to enable custom operators and performance optimization for Trainium hardware. By the end, you should know:

- How to configure torch.profiler with the Neuron backend and capture a clean trace.

- How to read the key_averages() table to identify which operations consume Neuron device time and compare their Self CPU time.

- How to follow one operation through Neuron Explorer, from framework dispatch through runtime scheduling to hardware execution on a NeuronCore.

- How these signals change across three implementations of the same operation.

Neuron as a PyTorch Backend

Before you profile a workload, it helps to understand how Neuron integrates with PyTorch.

The Neuron backend uses PyTorch's PrivateUse1 extension point. This mechanism lets an accelerator integrate with the framework without adding a dedicated dispatch key such as cuda. The torch_neuronx package works with upstream PyTorch to register Neuron as a native backend for PyTorch code. For details about this extension mechanism, see Facilitating New Backend Integration by PrivateUse1 .

Set Up the Profiler

You use two objects to profile a Neuron workload:

- NeuronConfig controls which Neuron data streams and NeuronCores to capture and where to write the output. You pass it to torch.profiler.profile as experimental_config.

- NeuronProfiler takes a NeuronConfig and provides an export_trace callback. The callback writes the framework trace, trace.json, alongside the Neuron system traces.

nc = NeuronConfig( modes=[ProfileMode.DEVICE, ProfileMode.RUNTIME], profile_output_dir=OUTPUT_DIR, capture_enabled_for_nc= '0' , correlate_device_time= True , ) exporter = NeuronProfiler(nc)

The modes parameter selects the data streams to capture, such as device instructions and runtime traces. The capture_enabled_for_nc parameter selects the NeuronCores to trace. Setting correlate_device_time=True propagates device time from the trace into the key_averages() table, where you can compare CPU and Neuron device time. For the complete parameter reference and available ProfileMode values, see the NeuronConfig Profile Modes documentation .

Tooling: Neuron Explorer

Neuron Explorer is a browser-based viewer for the system trace, device timeline, and per-engine breakdown. You can access it through an SSH tunnel or the Visual Studio Code extension. In this post, you use Explorer to confirm what key_averages() reports numerically and inspect what happens between dispatches.

After profiling, the output directory contains the system trace and, when you capture device data, matched .neff and .ntff files. The Neuron Executable File Format (NEFF) contains the compiled Neuron executable, while the matching NTFF contains its hardware instruction trace. The files share a numeric hash that identifies them as a pair. To open the profile in Neuron Explorer, upload the complete output directory for a system or combined profile. For a device-only profile, upload the matching .neff and .ntff files individually. See Step 5: Upload your profile for the UI and command-line instructions.

Profile RMSNorm

For this first analysis, you profile RMSNorm, a small operation with a clear sequence of arithmetic and reduction operations.

You profile three implementations of the same operation:

- Naive RMSNorm with individual PyTorch operations: multiply, mean, add, reciprocal square root, and multiply.

- torch.rms_norm(), the PyTorch built-in lowered to a single fused Neuron kernel.

- Naive RMSNorm with torch.compile, using the same operations as the first implementation.