AWS Trainium Profiling Cuts RMSNorm Time Sixfold With TorchNeuron Fusion
AWS shows how TorchNeuron profiling on Trainium cuts RMSNorm device time from 2.355 ms to 375.5 microseconds through kernel fusion with torch.compile.
Summary
On October 21, 2026, AWS engineer Josh Longenecker detailed system profiling with TorchNeuron, Trainium’s open source PyTorch backend supporting eager execution, torch.compile, standard APIs, and custom Neuron Kernel Interface operators. The walkthrough profiles RMSNorm on a 4 by 512 bfloat16 tensor using NeuronConfig with DEVICE and RUNTIME streams, NeuronCore 0, and device-time correlation. torch.profiler captures CPU and PrivateUse1 activity on a wait 2, warmup 1, active 1 schedule; torch.neuron.synchronize prevents traces exporting before device completion. NeuronProfiler writes trace.json with system traces, while Neuron Explorer links PyTorch dispatch, nrta_execute_schedule runtime events, and nc_exec_running hardware execution. Device captures produce numerically matched NEFF executables and NTFF instruction traces.
Naive RMSNorm issued six operations and took 2.355 ms of Neuron time and about 125 microseconds of CPU time. Three aten::mul calls consumed 1,164.5 microseconds, 49.5 percent, followed by aten::add at 434.6 microseconds, aten::rsqrt at 415.0 microseconds, and aten::mean at 340.8 microseconds; adding the CPU-resident epsilon also triggered a host-to-device copy. Built-in torch.rms_norm lowered to one fused kernel, taking 403.3 microseconds on Neuron and 47.6 microseconds on CPU. Compiling the naive function with TorchDynamo, AOTAutograd, and backend="neuron" fused the graph into one NEFF dispatch, taking 375.5 microseconds on Neuron and 98.4 microseconds on CPU. Both fused paths used about one-sixth of the naive device time by eliminating dispatch gaps and intermediate x_sq and mean_sq tensors. Part 2 will inspect TensorEngine, VectorEngine, and DMA activity, then use a roofline analysis to classify compute or memory limits.
Positives
- 375.5 microseconds of Neuron time made torch.compile slightly faster than the 403.3 microsecond purpose-built fused RMSNorm in this run.
- One fused dispatch reduced device time to roughly one-sixth of the naive implementation’s 2.355 ms.
- Neuron Explorer correlates PyTorch dispatch, runtime scheduling, and physical NeuronCore execution in one dependency chain.
- TorchNeuron supports eager execution, torch.compile, standard PyTorch APIs, and custom Neuron Kernel Interface operators with minimal code changes.
- NEFF and NTFF artifacts expose compiled executables and matching hardware instruction traces for deeper analysis.
Risks & concerns
- 2.355 ms of Neuron time made naive RMSNorm nearly six times slower than either fused implementation.
- Three separate aten::mul calls consumed 1,164.5 microseconds, or 49.5 percent of naive device time.
- A CPU-resident epsilon scalar caused aten::add to trigger an avoidable host-to-device transfer.
- Missing torch.neuron.synchronize can export the profile before device completion, leaving Neuron timing data empty.
- Device-level compute and memory bottlenecks remain unresolved until the planned TensorEngine, VectorEngine, DMA, and roofline analysis.