PyTorch torch.compile性能突破：LayerNorm与RMSNorm内核优化，GPU性能提升至SOTA水平

关键词：torch.compile、归一化算子、LayerNorm、RMSNorm、GPU性能优化

LayerNorm与RMSNorm是深度学习模型中的基础归一化算子，用于对输入数据进行标准化处理。它们是确保大模型训练平稳收敛、提升推理效率的关键模块。在高性能GPU平台上，其内核性能直接决定了整体训练吞吐量。

目前，业界顶尖的归一化内核多依赖于手工深度优化。相比之下，PyTorch原生的torch.compile在该类算子上长期存在性能差距，主要表现为内存带宽利用率低、归约逻辑与分块策略不合理等问题，难以充分发挥H100、B200等新一代GPU的硬件潜力。

本文以手工优化的Quack内核为性能基准，针对torch.compile的Inductor后端展开改进。通过优化内部归约分块、线程束配置等核心参数，提出混合阶数归约方案，并结合拆分归约、软件流水线与自动调优策略，我们大幅提升了前向与反向传播内核的效率。最终使编译生成的归一化算子达到行业顶尖水平，同时保留了自动融合周边算子的优势，为大模型编译优化提供了实用方案。

我们在英伟达H100与B200 GPU上对LayerNorm/RMSNorm的torch.compile性能进行了评估与优化，使其在逐内核层面达到接近行业顶尖（SOTA）的性能，同时还能通过自动融合能力实现进一步的速度提升。

本文目录

• 一、前向传播
  • 1.1 层归一化（LayerNorm）
  • 1.2 均方根归一化（RMSNorm）
  • 1.3 Quack基准内核
  • 1.4 torch.compile的生成逻辑
  • 1.5 基准测试结果
• 二、反向传播
  • 2.1 融合归约
  • 2.2 Inductor生成的融合反向传播内核
  • 2.3 拆分尺寸自动调优
  • 2.4 软件流水线优化
  • 2.5 基准测试结果
• 结论

一、前向传播

1.1 层归一化（LayerNorm）

层归一化最早在论文《Layer Normalization》中被提出。它通过计算输入张量的均值与方差进行归一化，并引入可学习的权重（γ）和偏置（β）参数进行缩放与平移。

1.2 均方根归一化（RMSNorm）

均方根归一化是LayerNorm的一种改进方案，在论文《Root Mean Square Layer Normalization》中被提出。该方法不再进行中心化处理，而是使用输入值的平方和的均方根（RMS）进行归一化。它同样使用可学习的权重（γ）进行缩放，但移除了偏置项。

LayerNorm与RMSNorm的前向传播逻辑相似，通常在连续维度上执行归约操作，并伴随若干逐元素运算。RMSNorm的计算量通常更低，因为它减少了浮点运算且无需处理偏置项。鉴于两者的内核相似性，本文将交替展示它们的基准测试结果。

1.3 Quack

Quack是由Tri Dao团队开发的、基于CuteDSL实现的极致优化GPU内核库。其README文档显示，在H100 GPU上，Quack针对此类归约内核的性能显著优于torch.compile。因此，我们将Quack作为行业顶尖（SOTA）性能基准，用以评估torch.compile的优化水平。

（图示说明：在bf16精度下，Quack在RMSNorm、Softmax、Cross-Entropy等核心算子上的性能全面领先于torch.compile及其他内核库。）

根据Quack的基准测试，优化前torch.compile的性能通常仅能达到Quack的50%左右。

1.4 torch.compile

以下阐述torch.compile为LayerNorm前向传播生成内核的通用逻辑（RMSNorm的实现思路与此完全一致）。我们假设输入的归约维度元素数量是连续的，在Inductor中这被称为内部归约。

（图示说明：RMSNorm/LayerNorm前向传播的分块并行计算流程。通过合理的分块策略适配GPU并行架构，优化访存效率。）

尽管内核生成逻辑涉及分块与并行调度，但其核心原理是直观的。

1.4 内核实现与优化策略

Inductor 生成的 LayerNorm/RMSNorm 前向传播内核遵循以下计算步骤：
1. 为输入的每一行维护大小为 R_BLOCK 的部分和；
2. 利用部分和计算均值与方差；
3. 依据层归一化公式对输入执行逐元素运算；
4. 存储逐元素运算的输出结果：
* 若开启逐元素仿射变换（elementwise_affine=True）且需要计算梯度（requires_grad=True），则额外存储均值与方差，供反向传播使用。

补充说明：若归约维度大小小于某一经验阈值（1024），Inductor 会生成持久化归约（persistent reduction） 内核。此时无需循环遍历归约维度，可直接计算均值。

在 H100 与 B200 显卡上对比 torch.compile 与 Quack 实现的 RMSNorm 前向传播时，我们复现了 torch.compile 性能远逊于 Quack 的问题。通过自动调优并优化 Inductor 默认参数后，torch.compile 在 H100 与 B200 上达到了行业顶尖（SOTA）性能。

实现该效果的核心优化手段如下：
* 基准测试中插入 torch._dynamo.reset()：确保 torch.compile 不启用自动动态形状推导；此前针对每个形状单独调用一次 torch.compile，会导致编译器判定使用动态形状。
* 修复不合理的自动调优配置决策：H100 与 B200 上的默认自动调优配置会做出次优选择，进而导致性能不佳；通过设置 mode='max-autotune' 可缓解该问题。我们对默认启发式规则做了多项改进：
1. 放大内部归约的 RBLOCK 大小；
2. 在持久化归约中放大 XBLOCK 大小，适用于归约元素数 ≤2048 的小尺寸归约场景；
3. 依据特定归约维度减少线程束数量（num_warps）。线程束数量过大时，往往会阻碍峰值向量化；峰值向量化是最大化飞行字节数、让内存受限型 workload 打满峰值内存带宽的关键，而 Blackwell 架构（B200）对内存带宽更为敏感，该问题影响更显著。

1.5 基准测试结果

下方展示 torch.compile 2.11 与 Quack（2026 年 3 月 24 日主干版本）在 Quack 基准测试形状，以及实际业务中常见的大 M、小 N 形状下的性能对比数据。结果表明，torch.compile 的性能基本与 Quack 持平。

测试中存在两类性能回退场景：
1. 在 N=384 时出现小幅性能回退，原因是 Triton 无法简洁地表示非 2 的整数次幂的分块大小；
2. H100 在 N 尺寸极大时出现大幅性能回退，原因是 Triton 暂不支持分布式共享内存的表达。

图：B200 GPU 在 bf16 精度下的 RMSNorm 性能测试，对比 Quack 与优化后的 Static torch.compile 方案。

图：H100 GPU 在 bf16 精度下的 RMSNorm 性能测试，对比 Quack 与优化后的 Static torch.compile 方案。

二、反向传播

LayerNorm/RMSNorm 的反向传播过程比前向传播更复杂。至少需要计算两类梯度：输入的梯度、权重的梯度；LayerNorm 中还需额外计算偏置梯度。

图：反向传播的梯度归约布局示意图，用于高效计算三类梯度。

为简化实现并规避复杂的数学公式，从性能角度考虑，这些梯度计算需要对反向传播的输入梯度（即前向传播上一步输出的梯度）的两个维度分别执行归约操作。

最朴素的实现方式是在独立内核中分别完成归约：一个内核计算输入的梯度，另一个内核计算权重和偏置的梯度。但该方案会导致同一输入被两次读取，鉴于归一化内核属于内存受限型任务，这会带来显著的额外延迟。

2.1 融合归约（Fused Reductions）

对于元素数通常不大、且单行数据可充分放入线程块的合理形状（一般 ≤16384），可以采用性能更优的融合内核，同时不会造成共享内存/寄存器占用爆炸。

简单来说，该内核会照常执行权重、偏置梯度的归约计算，同时为每一行并行完成输入梯度的列维度归约。已有多篇文献对这类融合方式进行了研究。

在 Inductor 中，我们用不同类型标识归约操作：
* 内部归约（INNER reduction）：对步长为 1 的连续维度执行归约；
* 外部归约（OUTER reduction）：对其余维度执行归约。

基于上述定义，融合内核会对同一输入张量同时执行内部归约与外部归约：其中内部归约对应输入梯度（连续维度），外部归约对应权重、偏置梯度。

拆分归约（Split Reduction）

在实际业务的大量形状中，批维度远大于归约维度。此时更优的方案是：对批维度的归约计算部分和，最终通过 torch.sum 汇总所有部分和，以此提升并行度。

Triton 官方教程中的层归一化示例就采用了拆分归约，但该实现通过原子操作加锁，让单个线程块负责处理单行数据；在批维度较大时，该方式性能极差，且会引发数值计算不一致问题。

图：反向传播中权重梯度的两阶段并行计算方案。

Inductor 当前也具备类似的拆分归约能力：与上述方案一致，会分配一块工作区张量存储部分和，但不使用原子操作；而是保证单个线程块协处理器（CTA）处理多行数据，并向工作区张量的唯一位置写入结果。

2.2 Inductor 生成的融合归一化反向传播内核

结合上文提到的融合归约与拆分归约方案，我们让 TorchInductor 能够自动生成达到行业顶尖水平的融合归一化反向传播内核。此外，编译器自动生成这类内核的方式，还支持更灵活的自动调优，以及与周边算子的自动融合能力。

由于该优化的核心难点是对同一输入执行不同归约顺序的归约融合，我们将该优化命名为混合阶数归约（MixOrderReduction）。

针对形状为 [M, N] 的输入，生成的内核执行流程如下：

2.2 混合阶数归约（MixOrderReduction）的实现

混合阶数归约的核心思想是将外部归约与内部归约融合到同一个内核中，避免对输入数据进行重复加载。其具体执行流程如下：

1. 按照预设的 SPLIT_SIZE 对归约维度进行分块；
2. 为每个分块分配工作区张量的一行，用于存储外部归约的中间结果（例如，每个分块的部分列和）；
3. 在每个分块内部，循环加载其包含的每一行数据：
   • 执行内部归约（例如，对整行求和，得到行和标量）；
   • 将该行数据与工作区中存储的中间结果合并，更新为新的部分归约结果；
4. 最后，启动一个轻量级的额外内核，将工作区中所有分块的部分归约结果汇总，得到外部归约的最终结果。由于此内核的输入张量尺寸极小，其执行开销可忽略不计。

在 Inductor 的代码生成逻辑中，识别到混合阶数归约模式后，会执行以下转换步骤：

1. 针对外部归约内核，将其内部的归约与 store_reduction 节点替换为新的 partial_accumulate 节点。该节点记录了待归约的数值和归约类型等信息。此变换将外部归约内核转换为一个逐元素操作内核（记为 PW1）。
2. 利用 Inductor 已有的循环重排序能力，对 PW1 进行循环顺序调整，得到新的逐元素内核（记为 PW2）。
3. 此时 PW2 的循环顺序已与内部归约内核保持一致，从而可以完成两者的内核融合。

2.3 拆分尺寸（Split-Size）自动调优

拆分尺寸 SPLIT_SIZE 是影响混合阶数归约内核性能的关键参数。

在 H100 GPU 上，针对 bfloat16 数据类型的输入，Liger 的 RMSNorm 反向传播内核默认性能为 0.417 TB/s；而将 SPLIT_SIZE 缩小 32 倍后，性能可提升至 1.912 TB/s。测试表明：
* 拆分尺寸选择不当，可能导致超过 2 倍的性能下降。
* 性能曲线近似呈抛物线形态。 因此，采用“持续将拆分尺寸放大 2 倍或缩小 1/2，直至找到性能峰值”的自动调优策略，可以高效地定位最优参数。

Inductor 原有的拆分归约特性会为提升性能而拆分外部归约，但其基于启发式规则选择的尺寸（图表中标记为 fused_split_reduction）往往效果不佳。我们让混合阶数归约忽略该尺寸，转而使用专属的启发式规则或自动调优机制，以选择更优的拆分尺寸。

2.4 软件流水线（Software Pipelining）

在优化反向传播内核以冲击峰值带宽的过程中，软件流水线（亦称加载预取）的优化价值得以凸显。

以往，流水线优化主要用于矩阵乘法（GEMM）、注意力机制等计算密集型任务，内存受限型任务通常无需此操作。Inductor 针对逐元素/归约内核以及 Liger 示例均未实现 num_stages（流水线阶段数）的自动调优。但我们注意到 Quack 内核中实现了类似的预取逻辑。

我们将 num_stages 作为通用自动调优参数加入 Inductor 内核，在部分输入形状下实现了显著加速，尤其在大 M、小 N 的场景中，应用于混合阶数归约后提速最高可达 20%。

2.5 基准测试结果

下方展示了混合阶数归约（MixOrderReduction）与 PyTorch 即时执行（eager）、旧版编译版本，以及 Quack、Liger 等开源基准的对比测试结果。所有测试均在 2025 年末、搭载 CUDA 12.9 的 750W B200 机器上完成。

从结果可以观察到：
* 搭载混合阶数归约的 torch.compile 比即时执行模式快 17.07 倍，而未搭载该优化的 torch.compile 仅比即时执行快 9.93 倍。
* 搭载混合阶数归约的 torch.compile 比 Liger 快 1.45 倍，比 Quack 快 1.34 倍。

鉴于内核逻辑相似，LayerNorm 的表现与 RMSNorm 基本一致。

结果呈现出与 RMSNorm 相同的趋势：
* 未搭载混合阶数归约的 torch.compile 相比 PyTorch 即时执行已有显著提速。
* 而搭载全新混合阶数归约方案的 torch.compile，相比旧版 torch.compile 实现了近 2 倍提速，性能更接近峰值内存带宽。

结论

本文围绕 PyTorch 的 torch.compile 在 LayerNorm 与 RMSNorm 归一化算子上的性能优化展开深入研究，最终实现了突破性成果：通过一系列针对性改进，让编译生成的归一化内核在 H100、B200 等新一代 GPU 上达到行业顶尖水平，全面对标甚至超越 Quack 手工优化内核与 Liger 等开源基准。

从具体表现来看，在反向传播环节搭载混合阶数归约方案后：相较原有 torch.compile 实现近 2 倍提速，比 Liger 快 1.45 倍、比 Quack 快 1.34 倍。前向传播也通过参数调优与自动调优策略，解决了非 2 的整数次幂分块等场景的性能回退问题，内存带宽利用率大幅提升。

这些成果的达成，源于对 torch.compile Inductor 后端核心痛点的精准突破：
* 前向传播中，通过优化内部归约的 RBLOCK 大小、调整持久化归约的 XBLOCK 配置、动态适配线程束数量，解决了原有编译内核硬件适配性差的问题。
* 反向传播中，创新性地提出混合阶数归约，融合内外归约逻辑并搭配拆分归约与软件流水线技术，规避了数据重复读取的效率损耗，同时通过拆分尺寸自动调优，避免了不当参数导致的性能下降。

更重要的是，本文优化方案保留了 torch.compile 的核心优势——自动融合周边逐元素与归约算子，这是手工编写内核难以实现的灵活性，能进一步提升大模型训练与推理的端到端效率。

该优化不仅为归一化算子提供了高性能编译方案，更验证了通过编译器后端优化、结合硬件特性的自动调优策略，能够充分发挥新一代 GPU 的硬件潜力，为深度学习框架中基础算子的编译优化提供了可复用的思路，对推动大模型高效部署具有重要的实践价值。

参考文献

[1] Layer Normalization, Jimmy Lei Ba, Jamie Ryan Kiros, Geoffrey E. Hinton: https://arxiv.org/abs/1607.06450

[2] Root Mean Square Layer Normalization, Biao Zhang, Rico Sennrich: https://arxiv.org/abs/1910.07467

[3] Quack 项目 (Dao-AILab): https://github.com/Dao-AILab/quack

[4] Liger-Kernel RMSNorm 实现: https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/rms_norm.py#L459

[5] Meta 提出的融合半持久化归一化反向传播内核: https://fb.workplace.com/groups/257735836456307/posts/999376505625566

[6] Quack 基于 CuteDSL 实现的融合 RMSNorm 内核: https://github.com/Dao-AILab/quack/blob/main/quack/rmsnorm.py

关注“鲸栖”小程序，掌握最新AI资讯