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 硬件通常有专门的快速近似指令。

代码实现

CUDA

#include <cuda_runtime.h>
#include <math.h>
__global__ void rms_norm_kernel(const float* input, float* output, int N, float gamma, float beta) {
    __shared__ float sdata[256];
    int tid = threadIdx.x; float sq = 0.0f;
    for (int i = blockIdx.x * blockDim.x + tid; i < N; i += gridDim.x * blockDim.x)
        sq += input[i] * input[i];
    sdata[tid] = sq; __syncthreads();
    for (int s = blockDim.x/2; s > 0; s >>= 1) {
        if (tid < s) sdata[tid] += sdata[tid+s]; __syncthreads();
    }
    float rms = sqrtf(sdata[0]/N + 1e-5f);
    for (int i = blockIdx.x * blockDim.x + tid; i < N; i += gridDim.x * blockDim.x)
        output[i] = gamma * input[i] / rms + beta;
}
extern "C" void solve(const float* input, float* output, int N, float gamma, float beta) {
    rms_norm_kernel<<<1, 256>>>(input, output, N, gamma, beta);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def rms_norm_kernel(input_ptr, output_ptr, N: tl.constexpr, gamma, beta, BLOCK: tl.constexpr):
    idx = tl.program_id(0)*BLOCK + tl.arange(0, BLOCK)
    mask = idx < N
    x = tl.load(input_ptr + idx, mask=mask, other=0.0)
    rms = tl.sqrt(tl.sum(x*x, axis=0)/N + 1e-5)
    tl.store(output_ptr + idx, gamma*x/rms + beta, mask=mask)