算子优化实战：手写 Triton Kernel，将 LayerNorm 算子的执行时间压缩 50%

标签：#Triton #CUDA #AICompiler #HPC #PyTorch #LayerNorm

---

📉 前言：为什么 PyTorch 原生算子还不够快？

PyTorch 的torch.nn.LayerNorm虽然底层调用了 cuDNN 或 ATen 的优化实现，但在处理特定 Shape 或与其他算子结合时，依然存在性能损耗：

1. 显存带宽瓶颈：LayerNorm 需要读取输入 ，计算均值 和方差 ，然后再读取一次 进行归一化。多次读写 Global Memory 是最大的开销。
2. 通用性妥协：原生库为了兼容各种极端 Shape，往往会有很多分支判断逻辑，牺牲了特定场景的极致性能。

Triton 的核心思路：将所有计算（均值、方差、归一化、仿射变换）融合在一个 Kernel 中，数据只从 Global Memory 读一次，写一次，中间全部在 GPU 的 SRAM（片上内存）中完成。

---

🧠 一、 LayerNorm 的数学原理与并行策略

LayerNorm 的公式如下：

其中，计算是按行 (Row-wise)独立的。

• 并行策略：每一行（Row）由一个 Triton Progr