Capstone: a fused kernel from scratch

01 · The spec

One op, everywhere, and secretly three kernels

The target: y = rmsnorm(x + residual) * w — fused residual-add-then-RMSNorm, the exact op every transformer layer applies twice. In eager mode it's three kernels and seven row-trips of DRAM traffic (day 23's calculator): add reads x and the residual and writes a temporary; RMSNorm reads the temporary and writes another; the scale multiply reads and writes again. Fused, it needs three: read x, read the residual, write y — the weight vector rides along, L2-cached (day 10). It's the ideal capstone because it's memory-bound (so day 11 governs), reduction-shaped (day 16), and elementwise around the reduction (day 23) — the two patterns of this journey in one kernel.

02 · The design

Decisions first, code second

Applying day 27's checklist in the forward direction. Output tile: one row per program — the reduction is per-row, so a program must own a whole row; BLOCK is the next power of two above the width, masked (day 20). Accumulation: sum of squares of the residual sum, in FP32 even when inputs are FP16 — day 27's upcast lesson, pre-installed. Staging: none — the row fits in registers, so shared memory would be ceremony. Tuning: BLOCK and warp count as autotune configs keyed on width (day 22). Expected profile: three rows of traffic per program over near-zero math — DRAM bar high and flat, or something's wrong.

03 · The kernel

End to end, launch included

@triton.autotune(
    configs=[
        triton.Config({'BLOCK': 512}, num_warps=4),
        triton.Config({'BLOCK': 1024}, num_warps=4),
        triton.Config({'BLOCK': 2048}, num_warps=8),
    ],
    key=['n_cols'],
)
@triton.jit
def fused_add_rmsnorm(x_ptr, r_ptr, w_ptr, y_ptr,
                                 n_cols, eps, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    offs = tl.arange(0, BLOCK)
    mask = offs < n_cols

    # the two loads — the only DRAM reads in the whole op
    x = tl.load(x_ptr + row * n_cols + offs, mask=mask,
                other=0.0).to(tl.float32)
    r = tl.load(r_ptr + row * n_cols + offs, mask=mask,
                other=0.0).to(tl.float32)
    t = x + r                    # residual add: registers only

    var = tl.sum(t * t, axis=0) / n_cols
    rstd = 1.0 / tl.sqrt(var + eps)

    w = tl.load(w_ptr + offs, mask=mask).to(tl.float32)
    tl.store(y_ptr + row * n_cols + offs, t * rstd * w, mask=mask)


def add_rmsnorm(x, residual, w, eps=1e-6):
    y = torch.empty_like(x)
    n_rows, n_cols = x.shape
    # grid: one program per row (day 14's ceil-division, named)
    fused_add_rmsnorm[(n_rows,)](x, residual, w, y, n_cols, eps,
        BLOCK=triton.next_power_of_2(n_cols))
    return y

Count the journey in that listing: masked block loads (day 20), an FP32 upcast before the reduction (day 25/27), a compiler-lowered tree sum (day 16), constexpr specialization and an autotuned search (day 22), a named grid (day 14), zero intermediates (day 23), and an intensity estimate that predicted the profile before launch (day 11). Twenty-odd lines — every one of them a decision you've seen made explicitly somewhere in the past four weeks.

04 · Verify, then profile

Correctness is a benchmark too

First against a reference: torch.allclose of this output against F.rms_norm(x + residual) * w in FP32 — tolerances sized for the dtype, and both FP16 and FP32 inputs exercised, because the FP16 path is where the upcast earns its keep. Then against the eager chain: three kernels and seven row-trips collapse to one kernel and three, so expect ~2.3× wall-clock at large widths — and, better, the Nsight SOL panel should show the DRAM bar in the mid-80s, flat, with compute barely ticking. If it doesn't, you now know exactly which of the four numbers to read first (day 18) and which chapter to reopen. That — hypothesis, change, measurement, acceptance — is the whole job. One artifact, every habit applied on purpose: that's the capstone.

Next: day 30 — the last chapter: where this road continues, and why mastery here is a practice tied to hardware that won't sit still.