Implement row-wise RMS normalization with a learnable scale. Given an input matrix of shape (M,N) and a length-N scale vector γ, normalize each row by its root-mean-square then apply the per-column scale. For row i:
rmsi=N1j=0∑N−1input[i,j]2+ϵ output[i,j]=rmsiinput[i,j]γ[j]Unlike LayerNorm, RMSNorm does not subtract the mean and has no bias term. The normalization is taken over the last axis (each row independently). All matrices are row-major.
Input: M = 1, N = 4, input = [[1, 2, 3, 4]], gamma = [1, 1, 1, 1], eps = 1e-5
Output: [[0.3651, 0.7303, 1.0954, 1.4606]]
The row RMS is 7.5≈2.7386; each element is divided by it.
Input: M = 2, N = 2, input = [[3, 4], [1, 1]], gamma = [1, 1], eps = 1e-5
Output: [[0.8485, 1.1314], [1.0, 1.0]]
Each row is normalized by its own RMS independently.
One block per row is a clean mapping: each block reduces the sum of squares for its row in shared memory.
rsqrtf(meanSquare + eps) gives the inverse RMS in a single fast instruction.
Accepts: array
Accepts: array
Accepts: number
Submit a correct solution to see your kernel performance.
Implement row-wise RMS normalization with a learnable scale. Given an input matrix of shape (M,N) and a length-N scale vector γ, normalize each row by its root-mean-square then apply the per-column scale. For row i:
rmsi=N1j=0∑N−1input[i,j]2+ϵ output[i,j]=rmsiinput[i,j]γ[j]Unlike LayerNorm, RMSNorm does not subtract the mean and has no bias term. The normalization is taken over the last axis (each row independently). All matrices are row-major.
Input: M = 1, N = 4, input = [[1, 2, 3, 4]], gamma = [1, 1, 1, 1], eps = 1e-5
Output: [[0.3651, 0.7303, 1.0954, 1.4606]]
The row RMS is 7.5≈2.7386; each element is divided by it.
Input: M = 2, N = 2, input = [[3, 4], [1, 1]], gamma = [1, 1], eps = 1e-5
Output: [[0.8485, 1.1314], [1.0, 1.0]]
Each row is normalized by its own RMS independently.
One block per row is a clean mapping: each block reduces the sum of squares for its row in shared memory.
rsqrtf(meanSquare + eps) gives the inverse RMS in a single fast instruction.
Accepts: array
Accepts: array
Accepts: number
Submit a correct solution to see your kernel performance.