RMS Normalization

原始题目：LeetGPU - RMS Normalization

题目描述

编写一个 GPU 程序，对一维输入向量实现 RMS 归一化（Root Mean Square Normalization）的前向传播。给定形状为 $[N]$ 的输入张量，使用标量缩放（$\gamma$）和平移（$\beta$）参数计算归一化输出。

RMS 归一化计算：

$$\begin{aligned} rms &= \sqrt{\frac{1}{N}\sum_{i=1}^{N} x_i^2 + \epsilon} \\ \hat{x}_i &= \frac{x_i}{rms} \\ y_i &= \gamma \hat{x}_i + \beta \end{aligned}$$

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终结果必须存储在 output 张量中。

示例

示例 1

Input:  input = [1.0, 2.0, 3.0, 4.0] (N=4), gamma=1.0, beta=0.0, eps=1e-5
Output: [0.365, 0.730, 1.095, 1.461]

示例 2

Input:  input = [1.0, 2.0, 3.0] (N=3), gamma=1.0, beta=0.0, eps=1e-5
Output: [0.463, 0.926, 1.389]

约束条件

• $1 \le N \le 100{,}000$。
• $\epsilon = 10^{-5}$。
• $-100.0 \le input \le 100.0$，$0.1 \le \gamma \le 10.0$，$-10.0 \le \beta \le 10.0$。
• 性能测试在 $N = 100{,}000$ 的规模下进行。

解题思路

RMS Norm 被 LLaMA、Mistral 等现代 LLM 广泛使用，相比 LayerNorm 省去了均值减法，计算更快。与 BatchNorm 类似，需要先做平方求和规约得到 RMS 值，再做逐元素除法。rsqrt（平方根倒数）是性能瓶颈的指令，但 GPU 硬件通常有专门的快速近似指令。欢迎在 GitHub Discussions 分享你的解法。