write-triton-layernorm-kernel

Skill: Write a Triton LayerNorm Kernel

Purpose

Guide the agent through implementing a correct, numerically stable row-wise layer normalization kernel in Triton. This covers mean and variance computation with fp32 accumulation, epsilon handling, affine transform with gamma/beta, the RMSNorm variant, masking for hidden dimensions not divisible by BLOCK_SIZE, and pointer arithmetic for 1D affine parameters applied to 2D or higher-rank inputs.

---

Use this when

• You need a fused forward LayerNorm that avoids separate mean, variance, normalize, and scale passes — i.e., a single kernel reading each row once (or twice for a two-pass approach).
• You need an RMSNorm variant (no mean subtraction, only RMS scaling) that is not available in your framework's kernel library.
• You need a custom LayerNorm that fuses a downstream or upstream operation (e.g., fusing the residual add into the LayerNorm input).
• You require the backward pass and intend to write a custom autograd function — knowing how the forward is structured is a prerequisite.
• torch.nn.LayerNorm with torch.compile is not achieving the expected fusion or is producing numerical issues you need to diagnose.

---

Do not use this when

• The normalized shape maps to a standard torch.nn.LayerNorm call and no fusion is needed. torch.compile will fuse the LayerNorm efficiently.
• The hidden dimension is very small (< 64). Warp-level reductions in CUDA (via vendor libraries) are more efficient at this size.
• You need training with a custom backward pass for a non-standard normalization variant. Prefer implementing the full custom kernel with saved statistics before committing to a Triton forward-only version.
• The normalization is over a non-contiguous dimension. This skill covers row-wise normalization (last dimension). Normalizing over other axes requires a different decomposition.

---

Inputs the agent should gather first

Before writing any code, confirm:

1. Input shape — (N, H) for 2D, or (B, T, H) for sequence inputs. Which dimension is normalized? (Assume last dimension H unless stated otherwise.)
2. Hidden dimension H — fixed or dynamic? Is H a power of 2? Is H always divisible by BLOCK_SIZE?
3. Input dtype — fp16, bf16, or fp32. Accumulation for mean and variance must be fp32.
4. Affine transform — does the kernel apply learned gamma and beta parameters? Are they 1D tensors of shape (H,)?
5. RMSNorm vs LayerNorm — LayerNorm subtracts mean and divides by std; RMSNorm skips mean subtraction and uses RMS of activations.
6. Epsilon value — typically 1e-5 or 1e-6. Confirm whether it is added inside or outside the square root (both are used in practice — they differ numerically).
7. Forward-only or training — if training, the mean and inverse-std (rstd) must be saved for the backward pass.
8. Residual fusion — is there a residual to add before normalization (pre-norm pattern)?