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:
M = 64, N = 256, eps = 1e-6
True
- 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: out=x⋅rstd⋅weight. - 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
Background Knowledge
The problem involves implementing RMSNorm, a normalization technique used in deep learning models, particularly in LLaMA-style models. RMSNorm normalizes the input data by the root-mean-square of the row, which is calculated as ms = sum(x^2)/N, where x is the input data and N is the number of elements in the row. The normalized output is then obtained by multiplying the input data with the reciprocal of the square root of ms plus a small value eps, and a learnable weight.
To understand this problem, it's essential to have a basic knowledge of linear algebra and deep learning concepts, such as normalization techniques, tensor operations, and GPU acceleration. The problem also involves using the Triton programming language, which is a Python-based language for writing high-performance GPU code. Additionally, the problem requires comparing the implemented RMSNorm kernel with a PyTorch reference implementation, which involves understanding how to use PyTorch and its tensor operations.
The mathematical formulation of RMSNorm can be expressed as out = x / sqrt(mean(x^2) + eps) * weight, where x is the input data, mean(x^2) is the mean of the squared input data, eps is a small value for numerical stability, and weight is a learnable parameter. This formulation involves element-wise operations, reductions, and linear transformations, which are common in deep learning computations.
Algorithm/Approach
The general approach to solving this problem involves:
- Implementing the RMSNorm formula using Triton's tensor operations and GPU acceleration
- Allocating input data on the GPU and launching the Triton kernel
- Comparing the output of the Triton kernel with a PyTorch reference implementation using torch.allclose
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
Editor locked
The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.