Implement the forward pass of LayerNorm in Triton. For each row of the input matrix x∈RM×N, compute the mean and variance across the feature dimension, then normalize and apply a learned affine transform with vectors γ,β∈RN.
out[i,j]=σi2+ϵx[i,j]−μi⋅γ[j]+β[j]where:
μi=N1j∑x[i,j],σi2=N1j∑(x[i,j]−μi)2This is the row-parallel pattern: one program per row, the entire row held in a single tile, statistics computed with in-register reductions, and normalized affine output stored in place. Computing variance from centered values avoids cancellation while keeping the row in registers.
Input: x = [[1.0, 2.0, 3.0, 4.0]], gamma = [1, 1, 1, 1], beta = [0, 0, 0, 0], eps = 1e-5
Output: [[-1.3416, -0.4472, 0.4472, 1.3416]]
Input: x = [[1, 2], [3, 4]], gamma = [2, 2], beta = [1, 1], eps = 1e-5
Output: [[-1.0, 3.0], [-1.0, 3.0]]
Each program owns one row. Form column offsets across the row and load with other=0.0 so masked lanes contribute zero to the mean reduction.
After computing the mean, form centered = tl.where(mask, x - mean, 0.0), then compute tl.sum(centered * centered, axis=0) / N. The explicit mask prevents padded lanes from contributing (−μ)2.
Multiply the centered values by tl.math.rsqrt(var+ϵ) before applying γ and β. 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 in the row reductions.Accepts: array
Accepts: array
Accepts: array
Accepts: number
Submit a correct solution to see your kernel performance.
Implement the forward pass of LayerNorm in Triton. For each row of the input matrix x∈RM×N, compute the mean and variance across the feature dimension, then normalize and apply a learned affine transform with vectors γ,β∈RN.
out[i,j]=σi2+ϵx[i,j]−μi⋅γ[j]+β[j]where:
μi=N1j∑x[i,j],σi2=N1j∑(x[i,j]−μi)2This is the row-parallel pattern: one program per row, the entire row held in a single tile, statistics computed with in-register reductions, and normalized affine output stored in place. Computing variance from centered values avoids cancellation while keeping the row in registers.
Input: x = [[1.0, 2.0, 3.0, 4.0]], gamma = [1, 1, 1, 1], beta = [0, 0, 0, 0], eps = 1e-5
Output: [[-1.3416, -0.4472, 0.4472, 1.3416]]
Input: x = [[1, 2], [3, 4]], gamma = [2, 2], beta = [1, 1], eps = 1e-5
Output: [[-1.0, 3.0], [-1.0, 3.0]]
Each program owns one row. Form column offsets across the row and load with other=0.0 so masked lanes contribute zero to the mean reduction.
After computing the mean, form centered = tl.where(mask, x - mean, 0.0), then compute tl.sum(centered * centered, axis=0) / N. The explicit mask prevents padded lanes from contributing (−μ)2.
Multiply the centered values by tl.math.rsqrt(var+ϵ) before applying γ and β. 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 in the row reductions.Accepts: array
Accepts: array
Accepts: array
Accepts: number
Submit a correct solution to see your kernel performance.