Implement the LayerNorm backward pass in Triton. Given the forward input x∈RM×N, the scale vector γ∈RN, and the upstream gradient dy∈RM×N, compute dx, dγ, and dβ:
x^[i,j]=σi2+ϵx[i,j]−μi,dynorm[i,j]=dy[i,j]⋅γ[j] dx[i,j]=σi2+ϵ1⋅(dynorm[i,j]−N1k∑dynorm[i,k]−x^[i,j]⋅N1k∑dynorm[i,k]⋅x^[i,k]) dγ[j]=i∑dy[i,j]⋅x^[i,j],dβ[j]=i∑dy[i,j]One program per row handles the dx computation. The reductions for dγ and dβ accumulate across rows, so each program adds its row contribution into the length-N output buffers with an atomic add after the launcher pre-zeros them. This mirrors the structure of the official Triton tutorial for fused LayerNorm.
Input: x = [[1.0, 2.0, 3.0, 4.0]], gamma = [1, 1, 1, 1], dy = [[1, 0, 0, 0]], eps = 1e-5
Output: dx = [[0.6708, 0.2236, -0.2236, -0.6708]], dgamma = [-1.3416, -0.4472, 0.4472, 1.3416], dbeta = [1, 0, 0, 0]
Input: x = [[1, 2], [3, 4]], gamma = [1, 1], dy = [[1, 1], [1, 1]], eps = 1e-5
Output: dx = [[0, 0], [0, 0]], dgamma = [-2, 2], dbeta = [2, 2]
Recompute the row statistics inside the kernel; do not pass forward-pass scratch. The two row-level reductions you need are c1=N1∑kdynorm[i,k] and c2=N1∑kdynorm[i,k]⋅x^[i,k].
With c1 and c2 in hand the formula collapses to dx[i,j]=rstd⋅(dynorm[i,j]−c1−x^[i,j]⋅c2). Apply once across the row in registers.
dγ and dβ are reductions across rows. Each program contributes its row via \texttt{tl.atomic_add} on the length-N buffers with the column-tail mask.
torch.float32 on cuda. eps is a Python float in [10−6,10−1].BLOCK_SIZE as a tl.constexpr chosen as triton.next_power_of_2(N) in the launcher.atol = 1e-2 with rtol = 1e-2 because both the row reduction and the cross-row atomic reduction accumulate float32 error.Accepts: array
Accepts: array
Accepts: array
Accepts: number
Submit a correct solution to see your kernel performance.
Implement the LayerNorm backward pass in Triton. Given the forward input x∈RM×N, the scale vector γ∈RN, and the upstream gradient dy∈RM×N, compute dx, dγ, and dβ:
x^[i,j]=σi2+ϵx[i,j]−μi,dynorm[i,j]=dy[i,j]⋅γ[j] dx[i,j]=σi2+ϵ1⋅(dynorm[i,j]−N1k∑dynorm[i,k]−x^[i,j]⋅N1k∑dynorm[i,k]⋅x^[i,k]) dγ[j]=i∑dy[i,j]⋅x^[i,j],dβ[j]=i∑dy[i,j]One program per row handles the dx computation. The reductions for dγ and dβ accumulate across rows, so each program adds its row contribution into the length-N output buffers with an atomic add after the launcher pre-zeros them. This mirrors the structure of the official Triton tutorial for fused LayerNorm.
Input: x = [[1.0, 2.0, 3.0, 4.0]], gamma = [1, 1, 1, 1], dy = [[1, 0, 0, 0]], eps = 1e-5
Output: dx = [[0.6708, 0.2236, -0.2236, -0.6708]], dgamma = [-1.3416, -0.4472, 0.4472, 1.3416], dbeta = [1, 0, 0, 0]
Input: x = [[1, 2], [3, 4]], gamma = [1, 1], dy = [[1, 1], [1, 1]], eps = 1e-5
Output: dx = [[0, 0], [0, 0]], dgamma = [-2, 2], dbeta = [2, 2]
Recompute the row statistics inside the kernel; do not pass forward-pass scratch. The two row-level reductions you need are c1=N1∑kdynorm[i,k] and c2=N1∑kdynorm[i,k]⋅x^[i,k].
With c1 and c2 in hand the formula collapses to dx[i,j]=rstd⋅(dynorm[i,j]−c1−x^[i,j]⋅c2). Apply once across the row in registers.
dγ and dβ are reductions across rows. Each program contributes its row via \texttt{tl.atomic_add} on the length-N buffers with the column-tail mask.
torch.float32 on cuda. eps is a Python float in [10−6,10−1].BLOCK_SIZE as a tl.constexpr chosen as triton.next_power_of_2(N) in the launcher.atol = 1e-2 with rtol = 1e-2 because both the row reduction and the cross-row atomic reduction accumulate float32 error.Accepts: array
Accepts: array
Accepts: array
Accepts: number
Submit a correct solution to see your kernel performance.