Fused RMSNorm (Triton)

A fused RMSNorm kernel written in Triton, with a custom backward pass. RMSNorm is the normalization layer used in Llama, Mistral, Qwen and most modern LLMs.

y = x / sqrt(mean(x^2) + eps) * weight

Usage

import torch
from rmsnorm_triton import rmsnorm, RMSNorm

x = torch.randn(8, 4096, device="cuda", dtype=torch.float16, requires_grad=True)
w = torch.ones(4096, device="cuda", dtype=torch.float16, requires_grad=True)

y = rmsnorm(x, w)          # functional, differentiable
y.sum().backward()

layer = RMSNorm(4096).cuda().half()   # drop-in nn.Module

What it does

• Forward: one Triton program per row. Loads the row once, computes in fp32, writes once, and saves rstd for backward.
• Backward: computes dx per row and accumulates dw in per-program partial buffers (no atomics), reduced with a final .sum(0).
• num_warps and programs-per-SM are picked from the row size, tuned by a sweep on a T4.

Benchmarks

Hardware: NVIDIA T4 (40 SMs). dtype: float16. Rows: 4096. Times in ms (lower is better), measured with triton.testing.do_bench. Baseline: torch.nn.functional.rms_norm.

Forward

Forward + backward

The forward pass is at parity with PyTorch's fused op for small rows and faster for large rows. Most of the end-to-end gain comes from the fused backward pass.

Correctness

Forward and gradients are checked against a plain PyTorch fp32 reference, including odd sizes (N=100, N=777) and 3D inputs. See test_rmsnorm.py.

Limitations

• Benchmarked only on a T4 with float16. Other GPUs may prefer different num_warps settings.
• bfloat16 is expected to work but has not been tested (the T4 does not support it).
• Rows longer than 8192 have not been benchmarked. Rows above 65536 are rejected.
• Weight gradients are accumulated in fp32 partial buffers, which costs extra memory (num_programs x N floats).