Triton RMSNorm Forward Kernel
Problem Statement
Implement RMSNorm over the last dimension: out = x / sqrt(mean(x^2) + eps) * weight. Unlike LayerNorm, there is no mean subtraction and no bias.
Background
RMSNorm (used in LLaMA-style models) normalizes by the root-mean-square of the row. Per row: ms = sum(x^2)/N, rstd = 1/sqrt(ms + eps), out = x * rstd * weight.
Your Task
Implement rmsnorm_kernel and run(M=64, N=256, eps=1e-6) comparing to a torch reference.
How it is tested
Your solution must define a top-level function run(...) that allocates inputs on the GPU, launches your Triton kernel, and returns a boolean from torch.allclose(triton_out, torch_reference, ...). The grader prints run(...); the expected output is True.
Example
Input
M = 64, N = 256, eps = 1e-6
Output
True
Reasoning
- The
rmsnorm_kernelfunction is called withM = 64,N = 256, andeps = 1e-6, which represents the number of rows, columns, and a small value for numerical stability, respectively. - For each row, the function calculates the mean of the squared values
x^2and computesms = sum(x^2)/N, thenrstd = 1/sqrt(ms + eps). - The output
outis calculated by multiplying each elementxwithrstdand a learnedweight: . - The result from the
rmsnorm_kernelfunction is compared to a PyTorch reference implementation usingtorch.allclose, which checks if the two outputs are element-wise equal within a tolerance, resulting in the outputTrueif they match.
Constraints
- No mean subtraction, no bias
- ms = sum(x^2)/N; rstd = 1/tl.sqrt(ms + eps)
- out = x * rstd * weight