AWS Engineer Josh Longenecker Demonstrates Trainium Inference Profiling: Fused Kernels Reduce Device Time to Approximately One-Sixth of Naive Implementation
Josh Longenecker (Amazon) uses three groups of identically configured RMSNorm comparisons to demonstrate how to locate scheduling gaps on AWS Trainium using torch.profiler and Neuron Explorer, and quantifies the effect of fused kernels and torch.compile in reducing device time from 2.355 ms to approximately 400 μs.
Many machine learning practitioners have an intuition for making models run fast, but fewer know how to analyze why a model is slow. Amazon engineer Josh Longenecker published a tutorial on profiling Trainium inference performance on the Hugging Face blog, providing a reproducible systematic approach that starts with configuring the profiler, reading the key_averages() table, and using Neuron Explorer to trace the complete lifecycle of a single operation from framework dispatch to hardware execution.
The core comparison in the article centers on RMSNorm, a small operation, using three implementations: row-by-row PyTorch tensor operations (multiplication, mean, addition, reciprocal square root, then multiplication), the PyTorch built-in torch.rms_norm() (lowered by Neuron into a single fused kernel), and a naive implementation compiled with torch.compile. The three sets of scripts use identical inputs, profiler configurations, and schedules, allowing for direct comparison of results.
In the naive implementation, the key_averages() table shows that three aten::mul calls account for the largest share of Neuron device time (49.5%), with a total device time of 2.355 ms. Tracing in Neuron Explorer reveals the cause: each ATen operation independently crosses the dispatch boundary, following its own framework→runtime→hardware path, leaving gaps between hardware executions.
Switching to torch.rms_norm() results in the table showing only one record for aten::_fused_rms_norm, with device time dropping to 403.3 μs. The trace shows only one scheduling path and one hardware event, with no gaps between component operations. Compiling the naive function with torch.compile yields similar results: one Torch-Compiled Region record at 375.5 μs. The compiler discovers fusion from the original Python function, merging operations, optimizing memory access, and compiling them into a single NEFF artifact executed as one hardware dispatch.