write-triton-rmsnorm-kernel

Skill: Write a Triton RMSNorm Kernel

Purpose

Guide the agent through implementing a correct, numerically stable RMSNorm kernel in Triton: y = x * rsqrt(mean(x², axis=-1) + eps) * weight. RMSNorm is the dominant normalization in modern decoder-only LLMs (LLaMA, Mistral, Qwen, Gemma, DeepSeek). This skill covers one-pass sum-of-squares with fp32 accumulation, the persistent kernel pattern when the hidden dim fits in a single tile, masking for non-divisible tails, the affine weight broadcast (no bias), and the backward pass. RMSNorm is structurally simpler than LayerNorm — no mean subtraction, no Welford — but the failure modes around fp16 squaring and weight pointer arithmetic still bite.

Use this when

• You are writing the normalization layer for an LLM inference engine (vLLM-style, TensorRT-LLM-style, or custom) and want to fuse the residual add or a downstream epilogue with the norm.
• You need RMSNorm forward + backward for training a LLaMA-family model and apex.normalization.FusedRMSNorm is not available on your target hardware (e.g., AMD CDNA).
• You need a fused residual + RMSNorm — the pre-norm pattern that dominates LLM blocks — and want to avoid materializing the residual sum to HBM. The kernel may also need to write the post-residual sum back as the next block's residual stream.
• You are debugging numerical drift between a PyTorch reference and a vendor kernel and need a clean Triton baseline to bisect against.

Do not use this when

• You are on PyTorch 2.4+ and torch.nn.functional.rms_norm (or a torch.compile'd nn.RMSNorm) is sufficient. The compiler fuses the read, square, reduce, scale, and weight broadcast.
• You are using a HuggingFace LLaMA / Mistral / Qwen model and the stock LlamaRMSNorm with torch.compile meets your perf bar. Only write a custom kernel if you need fusion or you are inside an inference engine that controls launch.
• The hidden dim is very small (< 256). Vendor and warp-level CUDA reductions outperform a Triton tile-based approach at this size.
• You need normalization over a non-trailing dimension. RMSNorm by convention normalizes the last dim; this skill assumes that.
• The model uses LayerNorm, not RMSNorm. Use the Triton LayerNorm skill — re-adding mean subtraction into an "RMSNorm" kernel changes semantics.