Reductions: Softmax & Fused RMSNorm

1What a reduction is

So far every operation in your kernels was elementwise: one input element in, one output element out. x + y, x * 2, tl.exp(x) — the shape never changes.

But the moment you need a summary of a tile — the biggest value, the total, the average — you need the second pillar: a reduction.

Reductions are everywhere in a transformer, and all four of decode-lab’s target kernels lean on them:

Lesson 2’s matmul was secretly a reduction too: each output cell is the sum of products along K. tl.dot is just the fused, hardware-accelerated version of “multiply everything, then reduce”.

2tl.sum & friends: reduce, then broadcast back

Triton gives you four reduction primitives, all with the same shape rules:

tl.sum(x)                    # EVERYTHING → one scalar
tl.max(x)                    # (BLOCK_M, BLOCK_N) tile → scalar
tl.sum(x, axis=1)            # reduce along axis 1 only → (BLOCK_M,)
tl.max(x, axis=1)            # one value per row → (BLOCK_M,)
tl.min(x, axis=0)            # one value per column → (BLOCK_N,)
tl.mean(x, axis=1)           # average per row → (BLOCK_M,)

axis picks which dimension gets combined away; the other dimensions survive. The resulting column of scalars is exactly what you [:, None] back into a broadcast to shape the whole tile again (lesson 2’s trick, now doing real work):

m = tl.max(x, axis=1)        # (BLOCK_M,)  — one number per row
p = x - m[:, None]           # broadcast the column back over the tile
l = tl.sum(p, axis=1)        # (BLOCK_M,)  — one total per row
out = p / l[:, None]         # broadcast again

3Softmax: the two-pass dance

Softmax turns a row of raw scores into probabilities that sum to 1:

p_i = exp(x_i) / sum_j exp(x_j)

Problem: exp(1000) is Infinity in fp32. Attention scores routinely reach large magnitudes, so the naive formula silently NaNs. The fix is one of the most famous one-liners in numerical computing: subtract the row max first. Because exp(x − m) / Σ exp(x − m) = exp(x) / Σ exp(x) — the math is unchanged, but the exponents are now ≤ 0, so nothing overflows. This costs one extra pass over the row — which is free here, because the whole row already lives in the program’s registers:

@triton.jit
def softmax_kernel(x_ptr, out_ptr, M, N,
                   BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
    pid_m = tl.program_id(0)
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)      # (BLOCK_M,)
    offs_n = tl.arange(0, BLOCK_N)                        # (BLOCK_N,)

    row_mask = offs_m[:, None] < M
    col_mask = offs_n[None, :] < N
    x = tl.load(x_ptr + offs_m[:, None] * N + offs_n[None, :],
                mask=row_mask & col_mask, other=-float("inf"))

    m = tl.max(x, axis=1)                                 # (BLOCK_M,) row maxima
    m = tl.where(offs_m < M, m, 0.0)                      # fully-masked rows: no -inf
    p = tl.exp(x - m[:, None])                            # broadcast back over the row
    l = tl.sum(p, axis=1)                                 # (BLOCK_M,) row sums
    out = p / l[:, None]

    tl.store(out_ptr + offs_m[:, None] * N + offs_n[None, :], out,
             mask=row_mask & col_mask)

4Fused RMSNorm — the real decode-lab kernel

Now the payoff. RMSNorm is the normalization used by Qwen3 (and most modern LLMs): scale each row by the root mean square of its own values, then apply a learned per-column weight:

out_i = x_i / sqrt(mean(x²) + eps) * w_i

Why RMS and not LayerNorm? No mean subtraction, no batch statistics — just a per-token scale. Perfect for autoregressive decode, where a single token arrives at a time and there is no batch to average over. And it is exactly the reduction dance: x² elementwise, sum along the row, broadcast back, multiply.

Here is the kernel — this is the actual file in your repo, src/qwen3_lm/kernels/rmsnorm.py, lightly trimmed for the lesson:

@triton.jit
def rmsnorm_kernel(x_ptr, residual_ptr, weight_ptr, out_ptr, M, N, eps,
                   HAS_RESIDUAL: tl.constexpr, BLOCK: tl.constexpr):
    row = tl.program_id(0)                       # one program per row
    offs = tl.arange(0, BLOCK)
    mask = offs < N

    x = tl.load(x_ptr + row * N + offs, mask=mask, other=0.0).to(tl.float32)
    if HAS_RESIDUAL:                             # fused x + residual, no second kernel
        x += tl.load(residual_ptr + row * N + offs, mask=mask, other=0.0).to(tl.float32)
    weight = tl.load(weight_ptr + offs, mask=mask, other=0.0).to(tl.float32)

    variance = tl.sum(x * x, axis=0) / N         # mean of squares (0-padded)
    out = (x * tl.rsqrt(variance + eps) * weight).to(out_ptr.dtype.element_ty)
    tl.store(out_ptr + row * N + offs, out, mask=mask)

5Why fusing wins: the memory traffic ledger

What does the unfused torch reference do? (rmsnorm_eager in the repo):

variance = x.float().pow(2).mean(-1, keepdim=True)
out = weight * (x.float() * torch.rsqrt(variance + eps)).to(x.dtype)

Each of those torch operations is a separate kernel launch that reads or writes the full tensor from HBM (GPU memory). Count the trips for one row:

The fused kernel moves 3× less data — and data movement is the actual cost on a GPU.

6Homework — seriously, do this

Run lesson3_softmax_rmsnorm.py on Colab, then try each of these by modifying the code. Hints are collapsible — try it first, peek only if stuck.

@triton.jit
def minmax_kernel(x_ptr, out_ptr, M, N, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
    pid_m = tl.program_id(0)
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = tl.arange(0, BLOCK_N)
    row_mask = offs_m[:, None] < M
    col_mask = offs_n[None, :] < N
    x = tl.load(x_ptr + offs_m[:, None] * N + offs_n[None, :],
                mask=row_mask & col_mask, other=0.0)
    lo = tl.min(x, axis=1)
    hi = tl.max(x, axis=1)
    out = (x - lo[:, None]) / (hi[:, None] - lo[:, None])
    tl.store(out_ptr + offs_m[:, None] * N + offs_n[None, :], out,
             mask=row_mask & col_mask)

7Your learning journey