Norm backward in Triton — the dγ/dβ reduction

The forward RMSNorm of lesson 11 was a clean row-per-program job. Its backward splits into two kernels with opposite personalities: dx is per-row (one program per row, easy — you just recompute the row statistic) and dγ/dβ is a reduction over every token in the batch (the hard cross-row sum). This lesson is the Triton code for both halves, including the classic lock / grouped-partial reduction the official layernorm tutorial uses.

The question this lesson answers

Forward norm parallelizes perfectly: every row is independent, no program talks to another. So why is the backward the kernel people get wrong, and where exactly does the cross-program coordination sneak back in?

The gradients (the half-page you reverse)

For LayerNorm, per row x ∈ ℝ^d with μ = mean(x), σ = sqrt(var(x)+ε), x̂ = (x−μ)/σ, output y = γ·x̂ + β. Given dy, write dŷ = dy·γ. Then:

dx = (1/σ)·[ dŷ − mean(dŷ) − x̂·mean(dŷ·x̂) ]      (per-row, two in-row reductions)
dγ = Σ over all tokens  dy·x̂                       (per-feature, cross-row)
dβ = Σ over all tokens  dy                          (per-feature, cross-row)

RMSNorm drops the mean: x̂ = x/σ with σ = sqrt(mean(x²)+ε), so its dx loses the mean(dŷ) term and keeps mean(dŷ·x̂). The shape of the problem is identical. The two in-row reductions are why a normalized output's gradient depends on the whole row — but each row stays independent, so dx is still row-per-program. The dγ/dβ lines are the trouble: their reduction axis is the token axis, the very axis dx parallelized over.

The dx kernel — recompute the stat, don't store x̂

The forward saved one scalar per row (rstd = 1/σ) — same as lesson 11's wrapper. The backward reads x and dy, recomputes x̂ in registers, and never touches a stored normalized row. Real Triton:

@triton.jit
def layernorm_bwd_dx_kernel(
    dy_ptr, x_ptr, w_ptr, rstd_ptr, mean_ptr,
    dx_ptr,
    dw_partial_ptr, db_partial_ptr,        # [GROUPS, d] workspace
    stride_row, N, GROUPS,
    BLOCK: tl.constexpr,
):
    row  = tl.program_id(0)
    offs = tl.arange(0, BLOCK)
    mask = offs < N

    x_row  = x_ptr  + row * stride_row
    dy_row = dy_ptr + row * stride_row

    # masked loads, promote to fp32 immediately
    x    = tl.load(x_row  + offs, mask=mask, other=0.0).to(tl.float32)
    dy   = tl.load(dy_row + offs, mask=mask, other=0.0).to(tl.float32)
    w    = tl.load(w_ptr  + offs, mask=mask, other=0.0).to(tl.float32)

    # recompute the row stat (one extra reduction over x — already in regs)
    mean = tl.load(mean_ptr + row)           # LayerNorm; RMSNorm skips this
    rstd = tl.load(rstd_ptr + row)
    xhat = (x - mean) * rstd                  # RMSNorm: xhat = x * rstd

    # dy_hat = dy * gamma, then the two IN-ROW reductions in fp32
    dyhat = dy * w
    c1 = tl.sum(dyhat,        axis=0) / N      # mean(dy_hat)
    c2 = tl.sum(dyhat * xhat, axis=0) / N      # mean(dy_hat * xhat)
    dx = (dyhat - c1 - xhat * c2) * rstd       # RMSNorm: drop the c1 term

    tl.store(dx_ptr + row * stride_row + offs, dx.to(tl.bfloat16), mask=mask)

    # accumulate this row's contribution into the group's partial dw/db
    g = row % GROUPS
    tl.atomic_add(dw_partial_ptr + g * N + offs, (dy * xhat), mask=mask)
    tl.atomic_add(db_partial_ptr + g * N + offs, dy,          mask=mask)

The two tl.sum calls are the in-row reductions — warp shuffle plus SMEM, exactly the reduction primitive of lesson 06, run twice. mask guards the N-vs-BLOCK tail so other=0.0 lanes contribute nothing. Everything that feeds a reduction is fp32; only the dx store casts back to bf16.

The hard half: dγ/dβ across the whole batch

Look at the last three lines of the dx kernel. Each program owns one row, but it writes into a [GROUPS, d] buffer shared with other programs — the cross-program communication that forward norm never needed. There are three ways to combine the per-row partials, and they are the same three from lesson 28's table; here is what each looks like in Triton.

1 · tl.atomic_add straight to a [d] buffer

# every row atomically adds into the SAME d slots
tl.atomic_add(dw_ptr + offs, dy * xhat, mask=mask)
tl.atomic_add(db_ptr + offs, dy,        mask=mask)

Simplest to write, no second kernel. But with B·T programs hammering d slots, contention is brutal on hot features, and — fatally — the order of fp32 atomic adds is nondeterministic, so two identical runs produce bitwise-different gradients. That is a debugging nightmare in training (lesson 28 / the CUDA track's reproducibility argument).

2 · grouped partial buffers + a tiny reduce kernel

This is what the kernel above does. Instead of B·T writers per slot, you fix GROUPS (say 256) partial rows; row r accumulates into group r mod GROUPS. Now at most GROUPS programs touch a slot, not B·T. Then a second, trivially small kernel sums the groups:

@triton.jit
def dw_reduce_kernel(dw_partial_ptr, db_partial_ptr, dw_ptr, db_ptr,
                     GROUPS, N, BLOCK_N: tl.constexpr, BLOCK_G: tl.constexpr):
    col   = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N)
    cmask = col < N
    accw = tl.zeros((BLOCK_N,), tl.float32)
    accb = tl.zeros((BLOCK_N,), tl.float32)
    for g0 in range(0, GROUPS, BLOCK_G):       # walk the GROUPS axis in tiles
        rows = g0 + tl.arange(0, BLOCK_G)
        rmask = rows < GROUPS
        m = rmask[:, None] & cmask[None, :]
        accw += tl.sum(tl.load(dw_partial_ptr + rows[:, None]*N + col[None, :],
                               mask=m, other=0.0), axis=0)
        accb += tl.sum(tl.load(db_partial_ptr + rows[:, None]*N + col[None, :],
                               mask=m, other=0.0), axis=0)
    tl.store(dw_ptr + col, accw.to(tl.bfloat16), mask=cmask)
    tl.store(db_ptr + col, accb.to(tl.bfloat16), mask=cmask)

This cuts contention from B·T writers per slot to at most GROUPS, and the reduce kernel sums groups in a fixed order. But note the subtlety: the partial rows above are still built with atomic_add, and multiple rows map to the same group (row mod GROUPS), so the order in which they accumulate is not fixed — the partials are not bitwise-reproducible even though the reduce step is. For full determinism you remove the atomics entirely: give each program a disjoint contiguous block of rows, accumulate that block's partial in registers (it owns those rows, no one else writes the slot), and store one partial row per block. Then both the partial accumulation and the fixed-order reduce are deterministic. That disjoint-block variant — not the grouped-atomic_add one — is what PyTorch's native layernorm, Apex, and Liger ship. The workspace is small (a few MB; see the worked numbers) and the second pass is sub-microsecond.

3 · the historical lock + partial buffer (the official Triton tutorial)

Before grouped buffers became the norm, Triton's own layernorm tutorial used a lock: a small [GROUPS] integer array of spin-locks plus a [GROUPS, d] partial buffer. A program spins on tl.atomic_cas to acquire its group's lock, does a read-modify-write of that partial row, then releases:

# acquire: spin until the lock for this group is free
lock = lock_ptr + g
while tl.atomic_cas(lock, 0, 1) == 1:
    pass
count = count_ptr + g
n = tl.load(count)
if n == 0:                      # first writer initializes
    tl.atomic_xchg(count, 1)
else:                           # later writers accumulate
    partial = tl.load(dw_partial_ptr + g*N + offs, mask=mask, other=0.0)
    new = partial + dy * xhat
    ...                         # store back
tl.atomic_xchg(lock, 0)         # release

The lock makes the read-modify-write atomic and the accumulation order fixed (deterministic), and it still needs the second reduce kernel over the [GROUPS, d] partials. It works, but it serializes every program competing for the same group's lock — on hot groups that is real contention. The modern grouped-partial approach drops the lock entirely: each group's partial is touched by atomic_add (commutative, no read-modify-write you have to guard) or, if you want bitwise determinism without atomics, you assign each program a disjoint group so no slot is ever shared. The trade is the classic one:

Activation backward — fuse it into the epilogue

SiLU/GELU sit between matmuls in the MLP, and their backward is elementwise: dx = dy · f'(preact). Don't round-trip dpreact through HBM — recompute the derivative from the saved pre-activation inside the matmul-backward epilogue, mirroring how the forward fused activation into the matmul epilogue (lesson 08). For SiLU, f'(x) = σ(x)·(1 + x·(1 − σ(x))):

# inside the dgrad epilogue — preact was saved by the forward
s   = tl.sigmoid(preact)                       # recompute, don't store
dsilu = s * (1.0 + preact * (1.0 - s))         # SiLU'
dx  = dy * dsilu                                # fused, no extra HBM pass

One saved tensor (the pre-activation), one fused multiply, zero extra HBM round-trips. Storing the derivative itself would cost a full d-sized read for nothing — the same recompute reflex as the norm stat.

Worked numbers — B·T = 16384, d = 4096, bf16, H100 (3.35 TB/s)

• dx: reads dy + x (2 · 16384 · 4096 · 2 B = 0.27 GB), writes dx (0.13 GB), recomputes σ on chip. Bandwidth-bound: ~0.4 GB / 3.35 TB/s ≈ 120 µs. The two tl.sum reductions are register/SMEM work hidden under the load latency — free.
• dγ/dβ partials: folded into the same dx pass (the kernel above writes them while it already has dy and xhat in registers), so no extra HBM read. The only added cost is the workspace: 256 groups · 4096 · 4 B = 4 MB, plus a sub-microsecond reduce kernel that reads it back.
• If you stored x̂ instead of recomputing: +0.13 GB extra read ≈ +40 µs, a 33% slowdown on a bandwidth-bound kernel — for zero benefit. Recompute always wins for a norm.

Interactive · dγ/dβ reduction — atomic contention vs two-pass workspace

Set the batch rows, feature width, and group count. The widget shows the atomic-to-[d] contention (writers per feature slot) against the grouped two-pass workspace cost, and flags the determinism trade.

Where this leaves you

You now have the two backward personalities in Triton code: the per-row dx kernel that recomputes its stat with masked fp32 reductions, and the cross-token dγ/dβ reduction with its grouped fp32 workspace and tiny reduce kernel — plus the lock/atomic alternatives and why the grouped two-pass one wins on determinism. That, with the matmul backward (lesson 19) and the fused cross-entropy backward (lesson 20), is the backward-pass core of training in Triton.

This is the last lesson of the Triton track. The hardest backward of all — recomputing the entire N² score matrix, carrying the D = rowsum(dO ∘ O) correction, running the two-GEMM rule inside a streaming K-loop — is FlashAttention backward, and the depth lives on the CUDA side: gpu_kernels lesson 29. Everything you built here (recompute-the-stat, fp32 grad accumulators, grouped reductions) is a prerequisite for reading it.