TorchNeuron Profiling on AWS Trainium: System-Level Analysis and Fusion Benefits

• System-level profiling methodology for TorchNeuron on AWS Trainium uses torch.profiler and Neuron Explorer to trace execution from framework dispatch to hardware.
• Operator fusion via torch.rms_norm or torch.compile reduces RMSNorm device time to approximately 1/6 of the naive implementation (375 µs vs 2,355 µs).
• torch.compile with the Neuron backend achieves 375.5 µs latency, comparable to the fused torch.rms_norm kernel at 403.3 µs, by eliminating intermediate value materialization.
• Naive PyTorch operations result in 6 operator calls and 2,355 µs device time, with aten::mul accounting for 49.5% of the self NEURON time.

Profiling Tools and Configuration

• NeuronConfig passed to torch.profiler controls data stream capture and NeuronCore selection via experimental_config.
• Neuron Explorer visualizes system traces, device timelines, and engine breakdowns using .neff and .ntff artifacts.
• correlate_device_time=True enables comparison of CPU and Neuron device times in the key_averages() table.

Performance Analysis

• Fusion removes gaps between hardware executions and prevents physical generation of intermediate values like x_sq and mean_sq.
• The workflow identifies hotspots in key_averages(), verifies them in traces, and applies fused operations or compilation.
• Part 2 of the series will cover device-level profiling, including TensorEngine, VectorEngine, and DMA breakdowns.