Implement Root Mean Square Layer Normalization (RMSNorm) in Triton. For each row of x∈RM×N and column j:
out[i,j]=N1∑kx[i,k]2+ϵx[i,j]⋅γ[j]Unlike LayerNorm, RMSNorm does not subtract the mean and has no bias. It only rescales each row by its root-mean-square and multiplies by a learned scale γ. This is the normalization used in modern LLMs such as LLaMA and Mistral because it saves the mean subtraction and the bias parameter without measurably hurting quality.
Input: x = [[1.0, 2.0, 3.0, 4.0]], gamma = [1, 1, 1, 1], eps = 1e-5
Output: [[0.3651, 0.7303, 1.0954, 1.4606]]
Input: x = [[2, 2], [3, 4]], gamma = [1, 2], eps = 1e-5
Output: [[1.0, 2.0], [0.8485, 2.2627]]
Each program owns one row. Load with other=0.0 on the column mask so masked lanes contribute zero to ∑x2 without special-casing the boundary.
Use tl.math.rsqrt on mean_sq+ϵ for the inverse standard deviation; it is faster and more accurate than the reciprocal of a square root.
Multiply by γ[j] only - no mean subtraction, no β. Reuse the same column mask on the final store.
torch.float32 on cuda. eps is a Python float in [10−6,10−1].BLOCK_SIZE as a tl.constexpr; pick it as triton.next_power_of_2(N) in the launcher.atol = 1e-2 with rtol = 1e-2 to absorb accumulation drift across the row reduction.Accepts: array
Accepts: array
Accepts: number
Submit a correct solution to see your kernel performance.
Implement Root Mean Square Layer Normalization (RMSNorm) in Triton. For each row of x∈RM×N and column j:
out[i,j]=N1∑kx[i,k]2+ϵx[i,j]⋅γ[j]Unlike LayerNorm, RMSNorm does not subtract the mean and has no bias. It only rescales each row by its root-mean-square and multiplies by a learned scale γ. This is the normalization used in modern LLMs such as LLaMA and Mistral because it saves the mean subtraction and the bias parameter without measurably hurting quality.
Input: x = [[1.0, 2.0, 3.0, 4.0]], gamma = [1, 1, 1, 1], eps = 1e-5
Output: [[0.3651, 0.7303, 1.0954, 1.4606]]
Input: x = [[2, 2], [3, 4]], gamma = [1, 2], eps = 1e-5
Output: [[1.0, 2.0], [0.8485, 2.2627]]
Each program owns one row. Load with other=0.0 on the column mask so masked lanes contribute zero to ∑x2 without special-casing the boundary.
Use tl.math.rsqrt on mean_sq+ϵ for the inverse standard deviation; it is faster and more accurate than the reciprocal of a square root.
Multiply by γ[j] only - no mean subtraction, no β. Reuse the same column mask on the final store.
torch.float32 on cuda. eps is a Python float in [10−6,10−1].BLOCK_SIZE as a tl.constexpr; pick it as triton.next_power_of_2(N) in the launcher.atol = 1e-2 with rtol = 1e-2 to absorb accumulation drift across the row reduction.Accepts: array
Accepts: array
Accepts: number
Submit a correct solution to see your kernel performance.