AWS工程师Josh Longenecker演示Trainium推理剖析：融合内核将设备时间降至朴素实现的约六分之一

许多机器学习从业者对让模型跑得快有直觉，但更少人知道如何分析模型为什么慢。Amazon工程师Josh Longenecker在Hugging Face博客发布了一篇Trainium推理性能剖析教程，从配置profiler开始，到读取key_averages()表格，再到用Neuron Explorer追踪单条操作从框架调度到硬件执行的完整生命周期，给出了一套可复现的系统方法。

文章的核心对比围绕RMSNorm这一小操作展开，使用三种实现：逐行PyTorch张量运算（乘法、均值、加法、倒数平方根、再乘）、PyTorch内置torch.rms_norm()（被Neuron降低为单个融合内核）、以及用torch.compile编译朴素实现。三组脚本使用相同输入、相同profiler配置和相同schedule，结果可直接比较。

在朴素实现中，key_averages()表格显示三个aten::mul调用占据Neuron设备时间的最大份额（49.5%），总设备时间为2.355 ms。Neuron Explorer的追踪揭示了原因：每个ATen操作独立跨越调度边界，遵循各自的框架→运行时→硬件路径，硬件执行之间存在间隙。

切换到torch.rms_norm()后，表格只显示一条aten::_fused_rms_norm记录，设备时间降至403.3 μs。追踪中只有一条调度路径和一条硬件事件，组件操作之间没有间隙。用torch.compile编译朴素函数得到类似结果：一条Torch-Compiled Region记录，375.5 μs。编译器从原始Python函数中发现融合，将操作合并、优化内存访问并编译为单个NEFF工件，作为一次硬件调度执行。

Longenecker总结：在这几组运行中，两种融合方式都只占朴素实现设备时间的约六分之一。融合消除了硬件执行之间的间隙，并避免了在分别调度的操作之间物化中间值（如x_sq和mean_sq）。这套流程可直接迁移到更大模型：在表格中找到热点，在追踪中验证，再判断融合操作或torch.compile能否降低成本。

需要指出的是，所有数字来自作者单次运行，未报告重复次数或方差；RMSNorm是小操作，结论能否外推到更大模型需Part 2的TensorEngine/VectorEngine/DMA分解和roofline分析进一步验证。