{ "@context": "https://schema.org", "@id": "https://aggr.aymericbeaumet.com/items/huggingface-co/2026-10-21-profiling-torchneuron-part-1-system-level-profiling-on-aws/#webpage", "@type": [ "WebPage", "ArchiveComponent" ], "articleBody": "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.\n\nThis 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](https://github.com/aws-neuron/torch-neuronx), 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)](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/nki/index.html) to enable custom operators and performance optimization for Trainium hardware. By the end, you should know:\n\n- How to configure `torch.profiler` with the Neuron backend and capture a clean trace.\n- How to read the `key_averages()` table to identify which operations consume Neuron device time and compare their Self CPU time.\n- How to follow one operation through Neuron Explorer, from framework dispatch through runtime scheduling to hardware execution on a NeuronCore.\n- How these signals change across three implementations of the same operation.\n\n* * *\n\n## Neuron as a PyTorch Backend\n\nBefore you profile a workload, it helps to understand how [Neuron](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/index.html) integrates with PyTorch.\n\nThe 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](https://docs.pytorch.org/tutorials/advanced/privateuseone.html).\n\n* * *\n\n## Set Up the Profiler\n\nYou use two objects to profile a Neuron workload:\n\n- `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`.\n- `NeuronProfiler` takes a `NeuronConfig` and provides an `export_trace` callback. The callback writes the framework trace, `trace.json`, alongside the Neuron system traces.\n\n```python\nnc = NeuronConfig(\n modes=[ProfileMode.DEVICE, ProfileMode.RUNTIME],\n profile_output_dir=OUTPUT_DIR,\n capture_enabled_for_nc='0',\n correlate_device_time=True,\n)\nexporter = NeuronProfiler(nc)\n```\n\nThe `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](http://awsdocs-neuron.readthedocs-hosted.com/en/latest/tools/neuron-explorer/how-to-profile-workload.html#neuronconfig-profile-modes).\n\n* * *\n\n## Tooling: Neuron Explorer\n\nNeuron 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.\n\nAfter 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](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/tools/neuron-explorer/get-started.html#step-5-upload-your-profile) for the UI and command-line instructions.\n\n* * *\n\n## Profile RMSNorm\n\nFor this first analysis, you profile RMSNorm, a small operation with a clear sequence of arithmetic and reduction operations.\n\nYou profile three implementations of the same operation:\n\n1. Naive RMSNorm with individual PyTorch operations: multiply, mean, add, reciprocal square root, and multiply.\n2. `torch.rms_norm()`, the PyTorch built-in lowered to a single fused Neuron kernel.\n3. Naive RMSNorm with `torch.compile`, using the same operations as the first implementation.\n\nThe scripts use the same inputs, profiler configuration, and schedule so the results are directly comparable.\n\nAll scripts to reproduce this work are available in the `torch-neuronx` [examples directory](https://github.com/aws-neuron/torch-neuronx/tree/main/examples/hf_profiling_blog_pt1).\n\n* * *\n\n## Profile 1: Naive PyTorch Operations\n\nStart by writing RMSNorm with individual PyTorch operations. Each line maps to one or more ATen operations and can create intermediate tensors.\n\n```python\ndevice = torch.device('neuron')\n\nx = torch.randn(4, 512, dtype=torch.bfloat16, device=device)\nweight = torch.ones(512, dtype=torch.bfloat16, device=device)\n\ndef rms_norm_fn(t):\n x_sq = t * t # square\n mean_sq = x_sq.mean(dim=-1, keepdim=True) # mean of squares\n rms = torch.rsqrt(mean_sq + 1e-6) # 1 / sqrt(mean_sq + eps)\n return t * rms * weight # normalize + scale\n```\n\nCall the profiler with two arguments required for the Neuron backend:\n\n1. Include `ProfilerActivity.PrivateUse1` to collect activity from the Neuron device.\n2. Pass `NeuronConfig` as `experimental_config` to enable the Neuron-specific Kineto features.\n\nThe schedule uses `wait=2`, `warmup=1`, and `active=1` across four `prof.step()` calls. This records one warm iteration without including initial setup overhead.\n\nCall `torch.neuron.synchronize()` inside the profiler loop. Without this synchronization, the profiler can export before the device finishes and leave the Neuron time column empty.\n\n```python\nwith torch.profiler.profile(\n activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1],\n experimental_config=nc,\n record_shapes=False,\n on_trace_ready=exporter.export_trace,\n schedule=torch.profiler.schedule(wait=2, warmup=1, active=1),\n) as prof:\n for step in range(4):\n with torch.no_grad():\n _ = rms_norm_fn(x)\n torch.neuron.synchronize()\n prof.step()\n```\n\nAfter profiling finishes, inspect the `key_averages()` table. Sort by `self_privateuse1_time_total` to place the operations with the highest device cost first. The table shown here is filtered to the RMSNorm operations that report Neuron device time. The full `key_averages()` output also includes CPU-side events such as allocations and host-to-device transfers.\n\n```python\nprint(prof.key_averages().table(sort_by=\"self_privateuse1_time_total\", row_limit=4))\n```\n\n[![The key averages table shows that three aten::mul calls account for the largest share of Neuron device time.](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/mlyyxfL6i5GfbOgUsyIuL.png)](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/mlyyxfL6i5GfbOgUsyIuL.png)\n\n*The key averages table shows that three aten::mul calls account for the largest share of Neuron device time.*\n\nThe table shows which ATen operations ran, how often they ran, and how much time they spent on the CPU and Neuron device. `aten::mul` has the highest total device time because the function performs three multiplications.\n\n| Name | \\# Calls | Self CPU | Self NEURON | % NEURON |\n| ----------- | -------- | -------- | ----------- | -------- |\n| aten::mul | 3 | 60.9 μs | 1,164.5 μs | 49.5% |\n| aten::add | 1 | 26.6 μs | 434.6 μs | 18.5% |\n| aten::rsqrt | 1 | 15.6 μs | 415.0 μs | 17.6% |\n| aten::mean | 1 | 21.1 μs | 340.8 μs | 14.5% |\n\n*Self NEURON total: 2.355 ms*\n\nEach row maps directly to the function. `aten::mul` runs for `x * x`, `x * rms`, and the final multiplication by `weight`. `aten::mean` reduces the last dimension. `aten::add` adds the epsilon value and triggers a host-to-device copy because the scalar starts on the CPU. `aten::rsqrt` computes the reciprocal square root.\n\n[![Neuron Explorer follows an aten::mul operation from framework dispatch through runtime scheduling to execution on the Neuron hardware lanes.](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/GUGp8s5G4EV7rb09YBDbf.png)](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/GUGp8s5G4EV7rb09YBDbf.png)\n\n*Neuron Explorer follows an aten::mul operation from framework dispatch through runtime scheduling to execution on the Neuron hardware lanes.*\n\n### Read the Trace in Explorer\n\nIn Explorer, time flows from left to right and the event lifecycle flows from top to bottom. Each ATen operation follows the same pattern: creation in the framework lane, scheduling in the runtime lane as `nrta_execute_schedule`, and device execution in the hardware lane as `nc_exec_running`.\n\nThe system trace view organizes events into lanes. In the Default grouping mode, labels follow the pattern instance/PID/subsystem/identifier:\n\n| Lane | Label suffix | What it shows |\n| ------------------ | ---------------------- | ------------------------------------------------------------------------- |\n| Framework dispatch | `framework/TID:` | Operator dispatch events (PrivateUse1 runtime), grouped by thread ID. |\n| Framework device | `framework/Stream:` | Device interaction events (PrivateUse1 driver), grouped by stream number. |\n| Neuron Runtime | `neuron_rt/NC:` | Neuron Runtime API calls, grouped by logical NeuronCore. |\n| Neuron Hardware | `neuron_hw/NC:` | Hardware execution events, grouped by physical NeuronCore. |\n\nIn eager mode, each operation runs independently. Use the dependency chain viewer to follow `aten::mul.Tensor` from the framework layer through runtime scheduling to device execution.\n\n[![The dependency chain viewer highlights the lifecycle of one operation across the framework, runtime, and hardware lanes.](https://huggingface.co/datasets/jlongen-aws/blog-assets/resolve/main/hf_profiling_blog_pt1/dependency_chain.gif)](https://huggingface.co/datasets/jlongen-aws/blog-assets/resolve/main/hf_profiling_blog_pt1/dependency_chain.gif)\n\n*The dependency chain viewer highlights the lifecycle of one operation across the framework, runtime, and hardware lanes.*\n\n* * *\n\n## Profile 2: torch.rms\\_norm\n\nReplace the individual operations with PyTorch's built-in `rms_norm` function. Neuron lowers the operation to one fused kernel:\n\n```python\ndef rms_norm_fn(t):\n return torch.rms_norm(t, [512], weight)\n```\n\n[![The fused RMSNorm appears as one device operation in the key averages table.](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/uYhQmUtuUp1LYPOu8f9bf.png)](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/uYhQmUtuUp1LYPOu8f9bf.png)\n\n*The fused RMSNorm appears as one device operation in the key averages table.*\n\n| Name | \\# Calls | Self CPU | Self NEURON | % NEURON |\n| ------------------------ | -------- | -------- | ----------- | -------- |\n| aten::\\_fused\\_rms\\_norm | 1 | 47.6 μs | 403.3 μs | 100% |\n\n*Self NEURON total: 403.3 μs*\n\n[![The fused RMSNorm follows one framework-to-hardware path instead of dispatching each arithmetic operation separately.](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/OfM6vRT9iPBG42s8iWK1V.png)](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/OfM6vRT9iPBG42s8iWK1V.png)\n\n*The fused RMSNorm follows one framework-to-hardware path instead of dispatching each arithmetic operation separately.*\n\n### Compare the Lowering Path\n\nIn the naive implementation, each operation crosses the dispatch boundary independently and follows its own framework-to-runtime-to-hardware path. In the fused implementation, one Python call covers the complete computation. The trace contains one dispatch and one hardware event, with no gaps between component operations.\n\n* * *\n\n## Profile 3: Naive PyTorch with torch.compile\n\nFor the final example, compile the original naive RMSNorm function with `torch.compile`.\n\nWith `@torch.compile(backend=\"neuron\")`, TorchDynamo captures the Python bytecode as an FX graph and AOTAutograd lowers it to ATen operations. The Neuron backend receives the complete graph instead of dispatching each operation separately. It fuses the operations, optimizes memory access, and compiles the graph into a single NEFF artifact. The artifact runs as one hardware dispatch.\n\n```python\n@torch.compile(backend='neuron', fullgraph=True)\ndef naive_compiled(t):\n x_sq = t * t\n mean_sq = x_sq.mean(dim=-1, keepdim=True)\n rms = torch.rsqrt(mean_sq + 1e-6)\n return t * rms * weight\n```\n\n[![The compiled function appears as one Torch-Compiled Region in the key averages table.](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/TSJ2RWScmKPKmSroHCZaI.png)](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/TSJ2RWScmKPKmSroHCZaI.png)\n\n*The compiled function appears as one Torch-Compiled Region in the key averages table.*\n\n| Name | \\# Calls | Self CPU | Self NEURON | % NEURON |\n| -------------------------- | -------- | -------- | ----------- | -------- |\n| Torch-Compiled Region: 0/0 | 1 | 98.4 μs | 375.5 μs | 100% |\n\n*Self NEURON total: 375.5 μs*\n\n[![The compiled graph follows one path through runtime scheduling and hardware execution.](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/yWIFqisMxCtJHLYb2Rcsh.png)](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/yWIFqisMxCtJHLYb2Rcsh.png)\n\n*The compiled graph follows one path through runtime scheduling and hardware execution.*\n\nThe table shows one device execution, as in Profile 2, but the code takes a different path. Instead of calling a purpose-built fused operation, the compiler discovers the fusion from the original Python function. In this run, the Neuron time is comparable to `torch.rms_norm`, at 375 μs compared with 403 μs.\n\n* * *\n\n## Side-by-Side Summary\n\n| Variant | NEURON time | Self CPU time | Operator calls |\n| --------------------- | ----------- | ------------- | -------------- |\n| Naive ops | 2,355 µs | \\~125 µs | 6 |\n| torch.rms\\_norm | 403 µs | \\~48 µs | 1 |\n| Naive + torch.compile | 376 µs | \\~98 µs | 1 |\n\nBoth `torch.rms_norm` and `torch.compile` fuse RMSNorm into one kernel dispatch. In these runs, either approach takes about one-sixth of the naive implementation's device time. Fusion removes the gaps between hardware executions and avoids materializing intermediate values such as `x_sq` and `mean_sq` between separately dispatched operations.\n\n* * *\n\n## Conclusion\n\n`key_averages()` shows which operations consume device time, while its full output also reveals CPU-side setup and transfers. Explorer provides the corresponding visual trace. You can follow one operation from framework dispatch through runtime scheduling to hardware execution and see the gaps that fusion removes.\n\nYou can use the same workflow on larger models: find a hotspot in the table, verify it in the trace, and then determine whether a fused operation or `torch.compile` can reduce its cost.\n\nIn Part 2, we will move into the device-level profile. You open the NEFF, inspect the TensorEngine, VectorEngine, and direct memory access (DMA) breakdown, and place the kernel on the roofline to determine whether it is compute-bound or memory-bound.\n\nSpecial thanks to Aritra Roy Gosthipaty for the inspiration to begin this series, and to my peers Pinak Panigrahi, Samir Araujo, Akhil Azhikodan, Esha Lakhotia and the rest for reviewing!\n", "authors": [], "captured_at": "2026-10-21T19:14:42.600331207+00:00", "category": "ai", "content_html": "

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.

\n

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:

\n
    \n
  • How to configure torch.profiler with the Neuron backend and capture a clean trace.
    • \n
  • How to read the key_averages() table to identify which operations consume Neuron device time and compare their Self CPU time.
    • \n
  • How to follow one operation through Neuron Explorer, from framework dispatch through runtime scheduling to hardware execution on a NeuronCore.
    • \n
  • How these signals change across three implementations of the same operation.
    • \n
\n
\n

Neuron as a PyTorch Backend

\n

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

\n

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.

\n
\n

Set Up the Profiler

\n

You use two objects to profile a Neuron workload:

\n
    \n
  • 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.
    • \n
  • NeuronProfiler takes a NeuronConfig and provides an export_trace callback. The callback writes the framework trace, trace.json, alongside the Neuron system traces.
    • \n
\n
nc = NeuronConfig(\n    modes=[ProfileMode.DEVICE, ProfileMode.RUNTIME],\n    profile_output_dir=OUTPUT_DIR,\n    capture_enabled_for_nc='0',\n    correlate_device_time=True,\n)\nexporter = NeuronProfiler(nc)\n
\n

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.

\n
\n

Tooling: Neuron Explorer

\n

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.

\n

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.

\n
\n

Profile RMSNorm

\n

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

\n

You profile three implementations of the same operation:

\n
    \n
  1. Naive RMSNorm with individual PyTorch operations: multiply, mean, add, reciprocal square root, and multiply.
    2. \n
  2. torch.rms_norm(), the PyTorch built-in lowered to a single fused Neuron kernel.
    3. \n
  3. Naive RMSNorm with torch.compile, using the same operations as the first implementation.
    4. \n
\n

The scripts use the same inputs, profiler configuration, and schedule so the results are directly comparable.

\n

All scripts to reproduce this work are available in the torch-neuronx examples directory.

\n
\n

Profile 1: Naive PyTorch Operations

\n

Start by writing RMSNorm with individual PyTorch operations. Each line maps to one or more ATen operations and can create intermediate tensors.

\n
device = torch.device('neuron')\n\nx = torch.randn(4, 512, dtype=torch.bfloat16, device=device)\nweight = torch.ones(512, dtype=torch.bfloat16, device=device)\n\ndef rms_norm_fn(t):\n    x_sq = t * t                                    # square\n    mean_sq = x_sq.mean(dim=-1, keepdim=True)       # mean of squares\n    rms = torch.rsqrt(mean_sq + 1e-6)               # 1 / sqrt(mean_sq + eps)\n    return t * rms * weight                         # normalize + scale\n
\n

Call the profiler with two arguments required for the Neuron backend:

\n
    \n
  1. Include ProfilerActivity.PrivateUse1 to collect activity from the Neuron device.
    2. \n
  2. Pass NeuronConfig as experimental_config to enable the Neuron-specific Kineto features.
    3. \n
\n

The schedule uses wait=2, warmup=1, and active=1 across four prof.step() calls. This records one warm iteration without including initial setup overhead.

\n

Call torch.neuron.synchronize() inside the profiler loop. Without this synchronization, the profiler can export before the device finishes and leave the Neuron time column empty.

\n
with torch.profiler.profile(\n    activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1],\n    experimental_config=nc,\n    record_shapes=False,\n    on_trace_ready=exporter.export_trace,\n    schedule=torch.profiler.schedule(wait=2, warmup=1, active=1),\n) as prof:\n    for step in range(4):\n        with torch.no_grad():\n            _ = rms_norm_fn(x)\n        torch.neuron.synchronize()\n        prof.step()\n
\n

After profiling finishes, inspect the key_averages() table. Sort by self_privateuse1_time_total to place the operations with the highest device cost first. The table shown here is filtered to the RMSNorm operations that report Neuron device time. The full key_averages() output also includes CPU-side events such as allocations and host-to-device transfers.

\n
print(prof.key_averages().table(sort_by="self_privateuse1_time_total", row_limit=4))\n
\n

\"The

\n

The key averages table shows that three aten::mul calls account for the largest share of Neuron device time.

\n

The table shows which ATen operations ran, how often they ran, and how much time they spent on the CPU and Neuron device. aten::mul has the highest total device time because the function performs three multiplications.

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
Name# CallsSelf CPUSelf NEURON% NEURON
aten::mul360.9 μs1,164.5 μs49.5%
aten::add126.6 μs434.6 μs18.5%
aten::rsqrt115.6 μs415.0 μs17.6%
aten::mean121.1 μs340.8 μs14.5%
\n

Self NEURON total: 2.355 ms

\n

Each row maps directly to the function. aten::mul runs for x * x, x * rms, and the final multiplication by weight. aten::mean reduces the last dimension. aten::add adds the epsilon value and triggers a host-to-device copy because the scalar starts on the CPU. aten::rsqrt computes the reciprocal square root.

\n

\"Neuron

\n

Neuron Explorer follows an aten::mul operation from framework dispatch through runtime scheduling to execution on the Neuron hardware lanes.

\n

Read the Trace in Explorer

\n

In Explorer, time flows from left to right and the event lifecycle flows from top to bottom. Each ATen operation follows the same pattern: creation in the framework lane, scheduling in the runtime lane as nrta_execute_schedule, and device execution in the hardware lane as nc_exec_running.

\n

The system trace view organizes events into lanes. In the Default grouping mode, labels follow the pattern instance/PID/subsystem/identifier:

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
LaneLabel suffixWhat it shows
Framework dispatchframework/TID:<id>Operator dispatch events (PrivateUse1 runtime), grouped by thread ID.
Framework deviceframework/Stream:<n>Device interaction events (PrivateUse1 driver), grouped by stream number.
Neuron Runtimeneuron_rt/NC:<n>Neuron Runtime API calls, grouped by logical NeuronCore.
Neuron Hardwareneuron_hw/NC<n>:<id>Hardware execution events, grouped by physical NeuronCore.
\n

In eager mode, each operation runs independently. Use the dependency chain viewer to follow aten::mul.Tensor from the framework layer through runtime scheduling to device execution.

\n

\"The

\n

The dependency chain viewer highlights the lifecycle of one operation across the framework, runtime, and hardware lanes.

\n
\n

Profile 2: torch.rms_norm

\n

Replace the individual operations with PyTorch's built-in rms_norm function. Neuron lowers the operation to one fused kernel:

\n
def rms_norm_fn(t):\n    return torch.rms_norm(t, [512], weight)\n
\n

\"The

\n

The fused RMSNorm appears as one device operation in the key averages table.

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
Name# CallsSelf CPUSelf NEURON% NEURON
aten::_fused_rms_norm147.6 μs403.3 μs100%
\n

Self NEURON total: 403.3 μs

\n

\"The

\n

The fused RMSNorm follows one framework-to-hardware path instead of dispatching each arithmetic operation separately.

\n

Compare the Lowering Path

\n

In the naive implementation, each operation crosses the dispatch boundary independently and follows its own framework-to-runtime-to-hardware path. In the fused implementation, one Python call covers the complete computation. The trace contains one dispatch and one hardware event, with no gaps between component operations.

\n
\n

Profile 3: Naive PyTorch with torch.compile

\n

For the final example, compile the original naive RMSNorm function with torch.compile.

\n

With @torch.compile(backend="neuron"), TorchDynamo captures the Python bytecode as an FX graph and AOTAutograd lowers it to ATen operations. The Neuron backend receives the complete graph instead of dispatching each operation separately. It fuses the operations, optimizes memory access, and compiles the graph into a single NEFF artifact. The artifact runs as one hardware dispatch.

\n
@torch.compile(backend='neuron', fullgraph=True)\ndef naive_compiled(t):\n    x_sq = t * t\n    mean_sq = x_sq.mean(dim=-1, keepdim=True)\n    rms = torch.rsqrt(mean_sq + 1e-6)\n    return t * rms * weight\n
\n

\"The

\n

The compiled function appears as one Torch-Compiled Region in the key averages table.

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
Name# CallsSelf CPUSelf NEURON% NEURON
Torch-Compiled Region: 0/0198.4 μs375.5 μs100%
\n

Self NEURON total: 375.5 μs

\n

\"The

\n

The compiled graph follows one path through runtime scheduling and hardware execution.

\n

The table shows one device execution, as in Profile 2, but the code takes a different path. Instead of calling a purpose-built fused operation, the compiler discovers the fusion from the original Python function. In this run, the Neuron time is comparable to torch.rms_norm, at 375 μs compared with 403 μs.

\n
\n

Side-by-Side Summary

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
VariantNEURON timeSelf CPU timeOperator calls
Naive ops2,355 µs~125 µs6
torch.rms_norm403 µs~48 µs1
Naive + torch.compile376 µs~98 µs1
\n

Both torch.rms_norm and torch.compile fuse RMSNorm into one kernel dispatch. In these runs, either approach takes about one-sixth of the naive implementation's device time. Fusion removes the gaps between hardware executions and avoids materializing intermediate values such as x_sq and mean_sq between separately dispatched operations.

\n
\n

Conclusion

\n

key_averages() shows which operations consume device time, while its full output also reveals CPU-side setup and transfers. Explorer provides the corresponding visual trace. You can follow one operation from framework dispatch through runtime scheduling to hardware execution and see the gaps that fusion removes.

\n

You can use the same workflow on larger models: find a hotspot in the table, verify it in the trace, and then determine whether a fused operation or torch.compile can reduce its cost.

\n

In Part 2, we will move into the device-level profile. You open the NEFF, inspect the TensorEngine, VectorEngine, and direct memory access (DMA) breakdown, and place the kernel on the roofline to determine whether it is compute-bound or memory-bound.

\n

Special thanks to Aritra Roy Gosthipaty for the inspiration to begin this series, and to my peers Pinak Panigrahi, Samir Araujo, Akhil Azhikodan, Esha Lakhotia and the rest for reviewing!

\n", "content_markdown": "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.\n\nThis 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](https://github.com/aws-neuron/torch-neuronx), 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)](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/nki/index.html) to enable custom operators and performance optimization for Trainium hardware. By the end, you should know:\n\n- How to configure `torch.profiler` with the Neuron backend and capture a clean trace.\n- How to read the `key_averages()` table to identify which operations consume Neuron device time and compare their Self CPU time.\n- How to follow one operation through Neuron Explorer, from framework dispatch through runtime scheduling to hardware execution on a NeuronCore.\n- How these signals change across three implementations of the same operation.\n\n* * *\n\n## Neuron as a PyTorch Backend\n\nBefore you profile a workload, it helps to understand how [Neuron](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/index.html) integrates with PyTorch.\n\nThe 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](https://docs.pytorch.org/tutorials/advanced/privateuseone.html).\n\n* * *\n\n## Set Up the Profiler\n\nYou use two objects to profile a Neuron workload:\n\n- `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`.\n- `NeuronProfiler` takes a `NeuronConfig` and provides an `export_trace` callback. The callback writes the framework trace, `trace.json`, alongside the Neuron system traces.\n\n```python\nnc = NeuronConfig(\n modes=[ProfileMode.DEVICE, ProfileMode.RUNTIME],\n profile_output_dir=OUTPUT_DIR,\n capture_enabled_for_nc='0',\n correlate_device_time=True,\n)\nexporter = NeuronProfiler(nc)\n```\n\nThe `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](http://awsdocs-neuron.readthedocs-hosted.com/en/latest/tools/neuron-explorer/how-to-profile-workload.html#neuronconfig-profile-modes).\n\n* * *\n\n## Tooling: Neuron Explorer\n\nNeuron 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.\n\nAfter 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](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/tools/neuron-explorer/get-started.html#step-5-upload-your-profile) for the UI and command-line instructions.\n\n* * *\n\n## Profile RMSNorm\n\nFor this first analysis, you profile RMSNorm, a small operation with a clear sequence of arithmetic and reduction operations.\n\nYou profile three implementations of the same operation:\n\n1. Naive RMSNorm with individual PyTorch operations: multiply, mean, add, reciprocal square root, and multiply.\n2. `torch.rms_norm()`, the PyTorch built-in lowered to a single fused Neuron kernel.\n3. Naive RMSNorm with `torch.compile`, using the same operations as the first implementation.\n\nThe scripts use the same inputs, profiler configuration, and schedule so the results are directly comparable.\n\nAll scripts to reproduce this work are available in the `torch-neuronx` [examples directory](https://github.com/aws-neuron/torch-neuronx/tree/main/examples/hf_profiling_blog_pt1).\n\n* * *\n\n## Profile 1: Naive PyTorch Operations\n\nStart by writing RMSNorm with individual PyTorch operations. Each line maps to one or more ATen operations and can create intermediate tensors.\n\n```python\ndevice = torch.device('neuron')\n\nx = torch.randn(4, 512, dtype=torch.bfloat16, device=device)\nweight = torch.ones(512, dtype=torch.bfloat16, device=device)\n\ndef rms_norm_fn(t):\n x_sq = t * t # square\n mean_sq = x_sq.mean(dim=-1, keepdim=True) # mean of squares\n rms = torch.rsqrt(mean_sq + 1e-6) # 1 / sqrt(mean_sq + eps)\n return t * rms * weight # normalize + scale\n```\n\nCall the profiler with two arguments required for the Neuron backend:\n\n1. Include `ProfilerActivity.PrivateUse1` to collect activity from the Neuron device.\n2. Pass `NeuronConfig` as `experimental_config` to enable the Neuron-specific Kineto features.\n\nThe schedule uses `wait=2`, `warmup=1`, and `active=1` across four `prof.step()` calls. This records one warm iteration without including initial setup overhead.\n\nCall `torch.neuron.synchronize()` inside the profiler loop. Without this synchronization, the profiler can export before the device finishes and leave the Neuron time column empty.\n\n```python\nwith torch.profiler.profile(\n activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1],\n experimental_config=nc,\n record_shapes=False,\n on_trace_ready=exporter.export_trace,\n schedule=torch.profiler.schedule(wait=2, warmup=1, active=1),\n) as prof:\n for step in range(4):\n with torch.no_grad():\n _ = rms_norm_fn(x)\n torch.neuron.synchronize()\n prof.step()\n```\n\nAfter profiling finishes, inspect the `key_averages()` table. Sort by `self_privateuse1_time_total` to place the operations with the highest device cost first. The table shown here is filtered to the RMSNorm operations that report Neuron device time. The full `key_averages()` output also includes CPU-side events such as allocations and host-to-device transfers.\n\n```python\nprint(prof.key_averages().table(sort_by=\"self_privateuse1_time_total\", row_limit=4))\n```\n\n[![The key averages table shows that three aten::mul calls account for the largest share of Neuron device time.](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/mlyyxfL6i5GfbOgUsyIuL.png)](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/mlyyxfL6i5GfbOgUsyIuL.png)\n\n*The key averages table shows that three aten::mul calls account for the largest share of Neuron device time.*\n\nThe table shows which ATen operations ran, how often they ran, and how much time they spent on the CPU and Neuron device. `aten::mul` has the highest total device time because the function performs three multiplications.\n\n| Name | \\# Calls | Self CPU | Self NEURON | % NEURON |\n| ----------- | -------- | -------- | ----------- | -------- |\n| aten::mul | 3 | 60.9 μs | 1,164.5 μs | 49.5% |\n| aten::add | 1 | 26.6 μs | 434.6 μs | 18.5% |\n| aten::rsqrt | 1 | 15.6 μs | 415.0 μs | 17.6% |\n| aten::mean | 1 | 21.1 μs | 340.8 μs | 14.5% |\n\n*Self NEURON total: 2.355 ms*\n\nEach row maps directly to the function. `aten::mul` runs for `x * x`, `x * rms`, and the final multiplication by `weight`. `aten::mean` reduces the last dimension. `aten::add` adds the epsilon value and triggers a host-to-device copy because the scalar starts on the CPU. `aten::rsqrt` computes the reciprocal square root.\n\n[![Neuron Explorer follows an aten::mul operation from framework dispatch through runtime scheduling to execution on the Neuron hardware lanes.](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/GUGp8s5G4EV7rb09YBDbf.png)](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/GUGp8s5G4EV7rb09YBDbf.png)\n\n*Neuron Explorer follows an aten::mul operation from framework dispatch through runtime scheduling to execution on the Neuron hardware lanes.*\n\n### Read the Trace in Explorer\n\nIn Explorer, time flows from left to right and the event lifecycle flows from top to bottom. Each ATen operation follows the same pattern: creation in the framework lane, scheduling in the runtime lane as `nrta_execute_schedule`, and device execution in the hardware lane as `nc_exec_running`.\n\nThe system trace view organizes events into lanes. In the Default grouping mode, labels follow the pattern instance/PID/subsystem/identifier:\n\n| Lane | Label suffix | What it shows |\n| ------------------ | ---------------------- | ------------------------------------------------------------------------- |\n| Framework dispatch | `framework/TID:` | Operator dispatch events (PrivateUse1 runtime), grouped by thread ID. |\n| Framework device | `framework/Stream:` | Device interaction events (PrivateUse1 driver), grouped by stream number. |\n| Neuron Runtime | `neuron_rt/NC:` | Neuron Runtime API calls, grouped by logical NeuronCore. |\n| Neuron Hardware | `neuron_hw/NC:` | Hardware execution events, grouped by physical NeuronCore. |\n\nIn eager mode, each operation runs independently. Use the dependency chain viewer to follow `aten::mul.Tensor` from the framework layer through runtime scheduling to device execution.\n\n[![The dependency chain viewer highlights the lifecycle of one operation across the framework, runtime, and hardware lanes.](https://huggingface.co/datasets/jlongen-aws/blog-assets/resolve/main/hf_profiling_blog_pt1/dependency_chain.gif)](https://huggingface.co/datasets/jlongen-aws/blog-assets/resolve/main/hf_profiling_blog_pt1/dependency_chain.gif)\n\n*The dependency chain viewer highlights the lifecycle of one operation across the framework, runtime, and hardware lanes.*\n\n* * *\n\n## Profile 2: torch.rms\\_norm\n\nReplace the individual operations with PyTorch's built-in `rms_norm` function. Neuron lowers the operation to one fused kernel:\n\n```python\ndef rms_norm_fn(t):\n return torch.rms_norm(t, [512], weight)\n```\n\n[![The fused RMSNorm appears as one device operation in the key averages table.](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/uYhQmUtuUp1LYPOu8f9bf.png)](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/uYhQmUtuUp1LYPOu8f9bf.png)\n\n*The fused RMSNorm appears as one device operation in the key averages table.*\n\n| Name | \\# Calls | Self CPU | Self NEURON | % NEURON |\n| ------------------------ | -------- | -------- | ----------- | -------- |\n| aten::\\_fused\\_rms\\_norm | 1 | 47.6 μs | 403.3 μs | 100% |\n\n*Self NEURON total: 403.3 μs*\n\n[![The fused RMSNorm follows one framework-to-hardware path instead of dispatching each arithmetic operation separately.](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/OfM6vRT9iPBG42s8iWK1V.png)](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/OfM6vRT9iPBG42s8iWK1V.png)\n\n*The fused RMSNorm follows one framework-to-hardware path instead of dispatching each arithmetic operation separately.*\n\n### Compare the Lowering Path\n\nIn the naive implementation, each operation crosses the dispatch boundary independently and follows its own framework-to-runtime-to-hardware path. In the fused implementation, one Python call covers the complete computation. The trace contains one dispatch and one hardware event, with no gaps between component operations.\n\n* * *\n\n## Profile 3: Naive PyTorch with torch.compile\n\nFor the final example, compile the original naive RMSNorm function with `torch.compile`.\n\nWith `@torch.compile(backend=\"neuron\")`, TorchDynamo captures the Python bytecode as an FX graph and AOTAutograd lowers it to ATen operations. The Neuron backend receives the complete graph instead of dispatching each operation separately. It fuses the operations, optimizes memory access, and compiles the graph into a single NEFF artifact. The artifact runs as one hardware dispatch.\n\n```python\n@torch.compile(backend='neuron', fullgraph=True)\ndef naive_compiled(t):\n x_sq = t * t\n mean_sq = x_sq.mean(dim=-1, keepdim=True)\n rms = torch.rsqrt(mean_sq + 1e-6)\n return t * rms * weight\n```\n\n[![The compiled function appears as one Torch-Compiled Region in the key averages table.](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/TSJ2RWScmKPKmSroHCZaI.png)](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/TSJ2RWScmKPKmSroHCZaI.png)\n\n*The compiled function appears as one Torch-Compiled Region in the key averages table.*\n\n| Name | \\# Calls | Self CPU | Self NEURON | % NEURON |\n| -------------------------- | -------- | -------- | ----------- | -------- |\n| Torch-Compiled Region: 0/0 | 1 | 98.4 μs | 375.5 μs | 100% |\n\n*Self NEURON total: 375.5 μs*\n\n[![The compiled graph follows one path through runtime scheduling and hardware execution.](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/yWIFqisMxCtJHLYb2Rcsh.png)](https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/yWIFqisMxCtJHLYb2Rcsh.png)\n\n*The compiled graph follows one path through runtime scheduling and hardware execution.*\n\nThe table shows one device execution, as in Profile 2, but the code takes a different path. Instead of calling a purpose-built fused operation, the compiler discovers the fusion from the original Python function. In this run, the Neuron time is comparable to `torch.rms_norm`, at 375 μs compared with 403 μs.\n\n* * *\n\n## Side-by-Side Summary\n\n| Variant | NEURON time | Self CPU time | Operator calls |\n| --------------------- | ----------- | ------------- | -------------- |\n| Naive ops | 2,355 µs | \\~125 µs | 6 |\n| torch.rms\\_norm | 403 µs | \\~48 µs | 1 |\n| Naive + torch.compile | 376 µs | \\~98 µs | 1 |\n\nBoth `torch.rms_norm` and `torch.compile` fuse RMSNorm into one kernel dispatch. In these runs, either approach takes about one-sixth of the naive implementation's device time. Fusion removes the gaps between hardware executions and avoids materializing intermediate values such as `x_sq` and `mean_sq` between separately dispatched operations.\n\n* * *\n\n## Conclusion\n\n`key_averages()` shows which operations consume device time, while its full output also reveals CPU-side setup and transfers. Explorer provides the corresponding visual trace. You can follow one operation from framework dispatch through runtime scheduling to hardware execution and see the gaps that fusion removes.\n\nYou can use the same workflow on larger models: find a hotspot in the table, verify it in the trace, and then determine whether a fused operation or `torch.compile` can reduce its cost.\n\nIn Part 2, we will move into the device-level profile. You open the NEFF, inspect the TensorEngine, VectorEngine, and direct memory access (DMA) breakdown, and place the kernel on the roofline to determine whether it is compute-bound or memory-bound.\n\nSpecial thanks to Aritra Roy Gosthipaty for the inspiration to begin this series, and to my peers Pinak Panigrahi, Samir Araujo, Akhil Azhikodan, Esha Lakhotia and the rest for reviewing!\n", "dateCreated": "2026-10-21T19:14:42.600331207+00:00", "date_created": "2026-10-21T16:04:43+00:00", "date_modified": "2026-10-21T16:04:43+00:00", "date_published": "2026-10-21T16:04:43+00:00", "description": "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…", "external_url": "https://huggingface.co/blog/amazon/profiling-torchneuron-part-1", "git": { "commit": "33e7967b937ef6329f2d9194d87eba92105b7ced", "path": "items/huggingface-co/2026/10/2026-10-21-profiling-torchneuron-part-1-system-level-profiling-on-aws.md", "permalink": "https://github.com/aymericbeaumet/aggr-instance/blob/33e7967b937ef6329f2d9194d87eba92105b7ced/items/huggingface-co/2026/10/2026-10-21-profiling-torchneuron-part-1-system-level-profiling-on-aws.md" }, "id": "urn:aggr:item:6f9d7593c57a852034c39c47a70c7c743427beb7", "isBasedOn": { "@id": "https://huggingface.co/blog/amazon/profiling-torchneuron-part-1" }, "mainEntity": { "@id": "https://huggingface.co/blog/amazon/profiling-torchneuron-part-1", "@type": "CreativeWork", "archivedAt": { "@id": "https://aggr.aymericbeaumet.com/items/huggingface-co/2026-10-21-profiling-torchneuron-part-1-system-level-profiling-on-aws/#webpage" }, "dateModified": "2026-10-21T16:04:43+00:00", "datePublished": "2026-10-21T16:04:43+00:00", "headline": "Profiling TorchNeuron (Part 1): System-level profiling on AWS Trainium.", "image": { "@type": "ImageObject", "caption": null, "height": 90, "url": "https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/hEodPVvQZI24jETLbPooP.png", "width": 160 }, "timeRequired": "PT8M", "url": "https://huggingface.co/blog/amazon/profiling-torchneuron-part-1", "wordCount": 1794 }, "name": "Archived snapshot: Profiling TorchNeuron (Part 1): System-level profiling on AWS Trainium.", "preview": { "alt": null, "color": null, "height": 90, "placeholder": { "data_url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=", "hash": "" }, "url": "https://cdn-uploads.huggingface.co/production/uploads/6ac67ea29579ea34f1f7e470/hEodPVvQZI24jETLbPooP.png", "width": 160 }, "reading_minutes": 8, "replicated_at": null, "source": "huggingface-co", "source_name": "Hugging Face - Blog", "summary": null, "tags": [], "temporalCoverage": "2026-10-21T19:14:42.600331207+00:00", "title": "Profiling TorchNeuron (Part 1): System-level profiling on AWS Trainium.", "url": "https://aggr.aymericbeaumet.com/items/huggingface-co/2026-10-21-profiling-torchneuron-part-1-system-level-profiling-on-aws/", "word_count": 1794 }