Fused RMSNorm

What it is

RMSNorm normalizes a row by its own root-mean-square and then scales it by a learned vector γ. Arithmetically that is almost nothing: a square, a mean, an add, a reciprocal square root, a multiply, and another multiply. The cost is entirely in moving the row. Every one of those steps, run as its own PyTorch kernel, reads the row out of HBM and writes it back — so the operation runs at a fraction of what the card can move, and the arithmetic is never the limit.

Fusing means writing all of it as one kernel: read the row once, keep everything in registers and on-chip memory, write the result once. That is the whole idea, and the measurement that matters is not FLOPs but effective bandwidth — how close to the card’s memory ceiling the kernel gets.

How it works

One Triton program per row, so the entire reduction fits in a single block and no cross-block communication is needed. The row is loaded once in fp16 or bf16 and accumulated in fp32, because summing thousands of squared fp16 values in fp16 drifts. Block size and warp count are left to triton.autotune, which sweeps block sizes from 512 to 8,192 and 4 to 16 warps and picks per shape; a pruning hook drops any configuration whose block is narrower than the row.

The backward pass computes dx the same way, and accumulates dγ across blocks with tl.atomic_add straight into the shared gradient vector.

Results

On a T4 with a nominal 320 GB/s, the fused forward kernel peaks at about 240 GB/s — roughly three quarters of the card’s rated bandwidth, and level with torch.compile across the sweep. Eager PyTorch sits near 62 GB/s, which is the cost of the round trips rather than the maths.

The backward pass reaches about 54 GB/s against eager’s 20. That is a real speedup and still the weakest part of the project, for a reason worth stating plainly.

What didn’t work

The backward pass does not beat torch.compile, and the reason is the dγ accumulation. Every block atomically adds its partial weight gradient into one shared N-element vector. Under a high block count those atomics contend on the same addresses and serialize, and the kernel spends its time waiting rather than moving bytes. TorchInductor avoids this entirely: it writes each block’s partial into a workspace buffer and then launches a second, fully vectorized kernel to reduce them. No atomics, no contention.

Selecting Backward and Fused in the scene above shows this as a queue of carts at a single loading dock. It is not a subtle effect — it is the entire gap.

The forward parity is also worth reading precisely. The fused kernel matches torch.compile on the forward pass. It does not match it on the backward pass, and describing the project as reaching parity without saying which direction would overstate it.

Next

Implement the two-stage dγ reduction — block partials to a workspace, then a separate reduction kernel — and re-benchmark. The interactive above is built to take a fourth implementation in the backward mode when that exists, so the comparison can be made in the same place rather than in a new chart.