What RMSNorm is, and why fusing both passes matters
RMSNorm takes a row x ∈ ℝN, divides it by its root mean square, and scales by a learned gain w:
There is no mean subtraction and no bias, unlike LayerNorm. So the arithmetic per element is a couple of multiplies and one reciprocal square root. That is nothing. The cost is moving x and y through HBM. RMSNorm is memory bound, and for a memory bound op the only thing that moves the wall clock is how many times you touch memory.
The eager PyTorch version touches it many times. This line of math compiles to a chain of separate CUDA kernels: square, mean, add epsilon, rsqrt, two multiplies. Each one reads the activation back from HBM and writes a result out again. On a memory bound op those round trips are the runtime.
variance = x.pow(2).mean(dim=-1, keepdim=True) x = x * torch.rsqrt(variance + eps) return weight * x
A fused kernel reads each row once, does all the math in registers, and writes the result once. The speedup you should expect is roughly the number of HBM round trips you removed, a prediction to carry into the results. Fusing the backward adds one more win: save one float of normalization scale per row in the forward, and the backward reads it back instead of recomputing the whole reduction. The cost is a tiny O(M) write, rounding error next to the O(MN) activation traffic.
Every kernel is an HBM round trip
Same memory bus, same time axis. Press play and watch the fused kernel finish while eager is still grinding through reads and writes.
Four choices, each with its tradeoff stated
a.One program per row, single launch
Each Triton program owns one row and keeps the whole row resident in registers and SRAM. The forward reads x once and writes y once. The tile width is fixed per call to next_power_of_2(N) so a single tile spans the row, with the tail masked when N isn't a power of two.
The tradeoff. The whole row must fit in one tile, so the design caps the hidden dimension. I guard it at 65536 and raise rather than silently produce garbage. Real hidden dims (2048 to 8192) sit comfortably inside that, but it's a real limit, a very wide last dimension would need a multi-pass variant. I'd rather state the limit than pretend the kernel is general.
b.fp32 reductions regardless of storage dtype
Every reduction accumulates in fp32 even when x is fp16 or bf16. The kernel upcasts on load and only rounds back down at the final store. The sum of N squares is exactly where low precision bites: you're adding N positive numbers, the running total grows while increments stay small, and bf16's 8 mantissa bits stop representing small additions, error that is systematic, not averaging out.
Memory bound, not compute bound, the extra fp32 registers aren't the bottleneck, the bus is. A correctness win at no cost.
In a compute bound kernel the register pressure could lower occupancy. That just isn't the situation here.
c.Atomic-free weight gradient
The two gradients, with x̂ᵢ = xᵢ·rstd and dywᵢ = dyᵢ·wᵢ:
dx is a per-row reduction, one program per row, no coordination. dw is the awkward one: it sums one contribution from every row into a single length-N vector. The obvious atomic-add into a shared buffer works, but atomics on the same addresses serialize, and contention worsens as you add parallelism.
Instead, each program gets its own private accumulator of length N and walks a contiguous chunk of rows with a grid-stride loop. No two programs ever write the same address. A small torch.sum reduces the [n_programs, N] partial buffer down to [N].
The program count is capped at the SM count, so the partial buffer stays small. The quiet extra win: dx and dw share the same x and dy loads, so the backward reads each once even though it produces two gradients.
d.Autotuning
Autotuned per N over num_warps ∈ {1,2,4,8,16,32} and num_stages ∈ {1,2,4}, how many warps cooperate on one row (memory-level parallelism) and how deep the load pipeline is. BLOCK_SIZE is deliberately pinned, not searched, because correctness requires BLOCK_SIZE ≥ N. The honest caveat: the search space is small, it covers the knobs I was confident mattered, which is part of why there's bandwidth left on the table.
The whole thing sits behind a torch.autograd.Function that flattens any […, N] input to [M, N] and reshapes back. So rmsnorm(x, w) is a drop-in differentiable op.
How I measured
The headline metric is not FLOP/s. For a memory bound op FLOP/s is misleading; the number that means something is effective bandwidth as a fraction of what the card can deliver. Timing goes through triton.testing.do_bench, warm up, many reps, median, cold launches discarded, so autotuning and one-time compilation stay out of the measurement.
For bandwidth I count only the activation traffic the kernel is obliged to move. Forward is read x + write y = 2MN elements. Backward is read x, read dy, write dx = 3MN. The weight, the saved rstd, and the partial buffer are O(M+N), negligible, so I leave them out. Padding the byte count with traffic the op doesn't need would inflate the number.
Note: effective bandwidth is a model, counted bytes over measured time, not a DRAM-counter reading. Both fp16 and bf16 store 2 bytes/element, so dtype changes the measured time, not the byte count.
Compared against an fp32 oracle, not another kernel
I don't bit-match against an eager fp16 implementation, two correct fp16 kernels can differ in the last bit from rounding order, and matching one spelling tests the wrong thing. Instead I compare against an fp32 oracle: computed end-to-end in fp32, rounded to fp16/bf16 only at the very end. That gives an honest claim, "within X of the true value", rather than "identical to some other code I wrote."
The backward is checked the same way, analytic Triton gradients against autograd gradients of the oracle. gradcheck wants float64, which Triton doesn't support, so the fp32 oracle comparison is the right tool. The sweep covers hidden dims {2048, 4096, 8192} × row counts {512, 1024, 2048, 4096} × {fp16, bf16}, plus non-power-of-two and 3D shapes, plus the gradient checks.