亚马逊发布TorchNeuron系统级性能分析教程，用RMSNorm演示三种实现差异
Hugging Face博客文章介绍如何在AWS Trainium上用torch.profiler与Neuron Explorer定位推理耗时，同一RMSNorm的三种写法中，朴素逐算子版本设备时间约2.36毫秒，融合算子与torch.compile版本分别约403微秒和376微秒。
AI解读：这篇来自亚马逊、发布在Hugging Face博客的文章，是“Profiling TorchNeuron”系列第一部分，面向没有性能分析经验的读者，讲如何用torch.profiler和Neuron Explorer分析AWS Trainium上的推理性能。
文章以一个RMSNorm算子为例，比较三种实现：用多个独立PyTorch算子手写、调用内置torch.rms_norm()、以及对手写版本使用torch.compile。三者输入、profiler配置和调度相同。
文章给出的单次运行数据显示，朴素版本Self NEURON总时间为2.355毫秒、6次算子调用；融合的torch.rms_norm为403.3微秒、1次调用；torch.compile版本为375.5微秒、1次调用。作者称两种融合路径在这组运行中约为朴素版本设备时间的六分之一。
文章说明，ProfilerActivity.PrivateUse1用于采集Neuron设备活动，NeuronConfig作为experimental_config传入；若不在循环内调用torch.neuron.synchronize()，设备可能未完成就导出，导致Neuron时间列为空。
这些数字来自文章展示的单次示例运行，不是对Trainium或TorchNeuron整体性能的基准评测；文章预告第二部分将做设备级分析，查看NEFF中的TensorEngine、VectorEngine和DMA分解，并用roofline判断计算受限还是内存受限。
亚马逊在Hugging Face博客发布“Profiling TorchNeuron”系列第一部分，介绍如何在AWS Trainium上用torch.profiler采集推理性能数据，并借助Neuron Explorer把耗时对应到框架调度、运行时和硬件执行各层。文章明确面向没有性能分析经验的读者。
文章称TorchNeuron是Trainium的开源原生PyTorch后端，通过PyTorch的PrivateUse1扩展点注册，支持eager执行、torch.compile和标准PyTorch API，并可对接Neuron Kernel Interface（NKI）编写自定义算子。
用NeuronConfig与NeuronProfiler采集设备数据
文章给出的配置中，NeuronConfig通过modes选择要采集的数据流（示例为DEVICE和RUNTIME），通过profile_output_dir指定输出目录，capture_enabled_for_nc='0'指定要跟踪的NeuronCore，correlate_device_time=True把设备时间写入key_averages()表，便于比较CPU与Neuron设备时间。NeuronProfiler接收该配置并提供export_trace回调。
调用torch.profiler.profile时需在activities中加入ProfilerActivity.PrivateUse1以采集Neuron设备活动，并把NeuronConfig作为experimental_config传入。示例调度为wait=2、warmup=1、active=1，共四次prof.step()，只记录一次预热后的迭代。文章特别提示，须在profiler循环内调用torch.neuron.synchronize()，否则可能在设备完成前导出，使Neuron时间列为空。
- 采集后在key_averages()中按self_privateuse1_time_total排序，可把设备耗时最高的算子排在前面。
- 完整表格还包含CPU侧的内存分配和主机到设备传输等事件，文章示例只筛选了有Neuron设备时间的RMSNorm相关行。
- Neuron Explorer是基于浏览器的查看器，可经SSH隧道或VS Code扩展访问；系统或组合profile上传整个输出目录，仅设备profile则分别上传配对的.neff与.ntff文件。
三种RMSNorm写法的示例数据
第一个实现用逐算子写法计算RMSNorm：平方、求均值、加epsilon、取倒数平方根、再乘权重。key_averages()显示aten::mul调用3次、设备时间1164.5微秒，占49.5%；aten::add为434.6微秒，aten::rsqrt为415.0微秒，aten::mean为340.8微秒，Self NEURON合计2.355毫秒。文章指出aten::add之所以触发主机到设备拷贝，是因为标量epsilon初始位于CPU。
第二个实现调用torch.rms_norm()，被降为单个融合Neuron内核，表中显示aten::_fused_rms_norm调用1次、设备时间403.3微秒。第三个实现对手写函数加@torch.compile(backend='neuron', fullgraph=True)，TorchDynamo把Python字节码捕获为FX图，AOTAutograd降为ATen算子，再由Neuron后端融合并编译为单个NEFF，表中显示为一条Torch-Compiled Region，设备时间375.5微秒。
文章并列给出三者对比：朴素版本Self NEURON约2355微秒、Self CPU约125微秒、6次算子调用；torch.rms_norm约403微秒、约48微秒、1次；torch.compile约376微秒、约98微秒、1次。作者称融合去掉了硬件执行之间的间隙，也避免了分别调度时中间张量（如x_sq和mean_sq）的物化。
- 文章强调这些脚本使用相同输入、profiler配置和调度，因此结果可直接比较。
- 在Explorer中，一个ATen算子依次出现在框架、运行时调度（nrta_execute_schedule）与硬件执行（nc_exec_running）泳道，依赖链视图可跟踪单次算子的完整生命周期。
- 作者预告第二部分将打开NEFF，查看TensorEngine、VectorEngine和DMA的分解，并用roofline判断内核是计算受限还是内存受限。