# LayerNorm 前向与反向传播

## 实现目标

`layernorm.py` 实现 LayerNorm 的前向传播和三个输入梯度，并通过
`torch.autograd.Function` 接入 PyTorch 计算图。实现以整行能够放入 SRAM 为前提，
用较受限的输入范围换取融合执行效率。

## 核心数据流

1. `_layer_norm_forward()` 为每一行计算均值和方差倒数 `rstd`，保存二者供反向
   使用，再完成归一化及仿射变换 `x_hat * w + b`。
2. `LayerNorm.forward()` 将任意前导维度展平为 M 行，选择二次幂块大小，分配输出
   与中间统计量并启动前向内核。
3. `_layer_norm_backward_input()` 每个 PID 处理一行，计算 `grad_x`，同时形成该行对
   `grad_weight` 和 `grad_bias` 的贡献。
4. 多个 PID 通过分组锁和原子操作，将参数梯度累加到较小的中间缓冲区。
5. `_layer_norm_backward_params()` 再沿中间缓冲区归约，产生最终参数梯度。

## 关键细节与限制

- 最后一个特征维度必须连续，且填充后的整行不能超过内核采用的 SRAM 上限。
- 均值、方差和梯度归约使用 FP32，以减小低精度输入的累加误差。
- `mean` 和 `rstd` 在前向保存，避免反向时重复计算统计量。
- 参数梯度拆为两个内核，是为了降低写入冲突并使用更合适的并行规模。

## 测试与基准

```bash
python3 layernorm.py
python3 layernorm.py --benchmark
```

测试分别比较输出、输入梯度、权重梯度和偏置梯度；基准测试比较 PyTorch 与 Triton
的前向或反向耗时。

## 配套视频

[bilibili](https://www.bilibili.com/video/BV1fMyWBgERM?spm_id_from=333.788.videopod.episodes&vd_source=e98b669ccbafff4b5aa59dd6303b722f&p=9)