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)?

Required reasoning process

1. Assign one program per row. row_idx = tl.program_id(0). For a (B, T, H) input, the grid is (B * T,). Compute the row base pointer as input_ptr + row_idx * row_stride.

2. Load the row in BLOCK_SIZE chunks with masking.

   col_offsets = tl.arange(0, BLOCK_SIZE)
   mask = col_offsets < H
   x = tl.load(x_ptr + row_idx * stride + col_offsets, mask=mask, other=0.0)
   x = x.to(tl.float32)
   

   Using other=0.0 is correct here (unlike softmax) because out-of-bounds positions should contribute 0 to the mean and sum-of-squares.

3. Compute the mean (LayerNorm only). mean = tl.sum(x, axis=0) / H. Note: divide by H (the full row length), not by the number of loaded elements. If H is not a multiple of BLOCK_SIZE, the masked positions loaded as 0.0 do not affect the sum, and dividing by H is still correct.

4. Subtract the mean (LayerNorm only). x_centered = x - mean. Apply the mask again to zero out positions beyond H.

5. Compute the variance. var = tl.sum(x_centered * x_centered, axis=0) / H. For RMSNorm, skip mean subtraction: var = tl.sum(x * x, axis=0) / H.

6. Compute rstd. rstd = 1.0 / tl.sqrt(var + epsilon).

7. Normalize. x_norm = x_centered * rstd (or x * rstd for RMSNorm).

8. Apply affine transform if gamma/beta are present.

   gamma = tl.load(gamma_ptr + col_offsets, mask=mask, other=1.0)
   beta  = tl.load(beta_ptr  + col_offsets, mask=mask, other=0.0)
   out   = x_norm * gamma + beta
   

   Gamma defaults to 1.0 and beta to 0.0 at out-of-bounds positions — these are identity values that produce no effect.

9. Store the output.

   tl.store(out_ptr + row_idx * out_stride + col_offsets, out.to(output_dtype), mask=mask)
   

10. If training, save mean and rstd. Store scalar mean and rstd to separate output buffers indexed by row_idx. These are needed in the backward pass.

11. For H > BLOCK_SIZE, loop over chunks. This complicates the two-pass approach: pass 1 accumulates partial sums for mean and variance across all chunks, pass 2 normalizes and stores. Alternatively, use Welford's online algorithm to compute mean and variance in a single pass.

Kernel design rules

• BLOCK_SIZE must be a power of 2 and declared tl.constexpr. Choose the smallest power of 2 >= H for the single-block case, or a fixed tile size for the loop case.
• All arithmetic for mean, variance, and rstd must use fp32. Cast inputs to fp32 immediately after loading.
• Epsilon placement: 1.0 / sqrt(var + epsilon) (epsilon inside sqrt) is the PyTorch convention. Placing epsilon outside (1.0 / (sqrt(var) + epsilon)) is numerically different and produces different outputs — clarify which convention the caller expects.
• Gamma and beta are 1D tensors of shape (H,). Their pointer arithmetic uses only col_offsets, not row_idx. They are shared across all rows. Do not add a row offset to gamma/beta pointers.
• Row stride must be passed as a kernel argument. Do not hardcode stride = H.
• For training, write mean and rstd to separate (N,) output buffers. The backward kernel reads these values to avoid recomputing them.
• Do not fuse the residual add as a masked read unless you have profiled and confirmed it reduces memory traffic. Fusing changes the memory footprint and can complicate the backward pass.

Correctness requirements

• The mean division denominator must be H (the full row length), not the number of non-masked positions. If a different convention is required (e.g., only normalize over valid tokens), the caller must specify this and the mask logic must be updated accordingly.
• Out-of-bounds positions must load as 0.0 so they contribute neither to the mean sum nor to the variance sum. Confirm other=0.0 in the load call, not other=-float('inf') (which would corrupt the mean).
• The variance must be computed as E[x^2] - E[x]^2 (or equivalently as E[(x - mean)^2]). Both are algebraically equivalent but the two-pass approach (subtract mean first, then compute variance) is numerically more stable and preferred.
• Epsilon must be strictly positive. Reject or warn on epsilon <= 0.
• For RMSNorm, the mean subtraction step must be absent. Including it would change the semantics to LayerNorm.
• Gamma and beta loading must not include any row offset — these parameters are row-invariant. Adding row_idx * H to the gamma pointer is a common bug that reads from garbage memory.
• The output store mask must match the input load mask (same col_offsets < H condition).

Performance requirements

The agent must reason through the following before finalizing:

• Memory access pattern. LayerNorm reads each element once (or twice in a two-pass implementation) and writes once. For fp16 inputs at (BT, H) with H=1024, total traffic is ~3 * BTH2 bytes. The kernel should approach HBM bandwidth limits.
• Two-pass vs Welford. Two-pass (separate mean then variance scan) requires reading the row twice. Welford's online algorithm computes mean and variance in a single pass using running statistics. For H that fits in one BLOCK_SIZE, both fit in registers and there is no difference. For H > BLOCK_SIZE requiring a loop, Welford saves one read pass.
• Register pressure from fp32 accumulation. Storing fp32 versions of fp16 inputs doubles register use. For BLOCK_SIZE=1024 fp16 inputs cast to fp32, each program uses ~4KB of registers for the data alone. This limits occupancy on the SM.
• Gamma/beta loads. These 1D parameter arrays are read by every row program. They will quickly reside in L2 cache after the first few programs. Do not preload them into shared memory manually — Triton does not expose shared memory explicitly, and L2 reuse will handle it.
• Benchmark comparison. A well-written Triton LayerNorm should match torch.nn.LayerNorm (compiled) for H >= 512. For H < 512, vendor implementations tend to use warp-level reductions that are harder to match.

Output format

The agent should produce:

1. The Triton kernel function with @triton.jit, taking x_ptr, out_ptr, gamma_ptr, beta_ptr, mean_ptr (optional), rstd_ptr (optional), stride, H, epsilon, and BLOCK_SIZE: tl.constexpr as arguments.
2. A Python launcher that computes the grid as (N_rows,), handles contiguity, extracts strides, and passes epsilon.
3. RMSNorm variant (if requested) as a separate kernel or a conditional code path clearly marked.
4. A correctness test comparing against torch.nn.LayerNorm with torch.allclose(atol=1e-4) for fp16 inputs and atol=1e-5 for fp32 inputs.
5. Comment explaining epsilon placement convention used (inside vs outside sqrt).

Common failure modes

• fp16 variance underflow. For small activations (magnitude < 1e-2), x*x in fp16 underflows to 0. The variance is 0, rstd blows up or is set to an arbitrary large value, and the normalized output is garbage. Always compute variance in fp32.
• Epsilon outside vs inside sqrt. 1/sqrt(var + eps) and 1/(sqrt(var) + eps) behave differently when var is near 0. PyTorch uses inside. Using outside will not match PyTorch's reference output and can produce inf when var=0 and eps is small.
• Adding row offset to gamma pointer. Gamma is a 1D tensor. gamma_ptr + row_idx * H + col_offsets is wrong. The correct access is gamma_ptr + col_offsets. This bug reads from unrelated memory and produces incorrect scale factors.
• Dividing by number of loaded elements instead of H. When H is not a multiple of BLOCK_SIZE, some positions are masked. If you count tl.sum(mask.to(tl.float32)) as the denominator instead of H, you get a different (incorrect) mean when H is not power-of-2 aligned.
• Two-pass variance without reloading. Storing x_centered to global memory between pass 1 and pass 2 wastes bandwidth. For H > BLOCK_SIZE, use Welford's algorithm to avoid re-reading.
• Wrong rstd storage for backward. Storing sqrt(var + epsilon) instead of 1/sqrt(var + epsilon) as rstd means the backward kernel must do an extra division. Confirm the convention matches what the backward kernel expects.
• Missing mask on output store. Writing beyond the valid H positions corrupts adjacent rows in memory if the allocation is tightly packed, or wastes memory traffic.

Review checklist

• BLOCK_SIZE is a power of 2 and declared tl.constexpr.
• Input is cast to fp32 immediately after loading, before any arithmetic.
• Out-of-bounds load uses other=0.0 (not -inf), contributing 0 to mean and variance.
• Mean denominator is H (full row length), not count of valid elements.
• Variance is computed in fp32 before taking sqrt.
• Epsilon is placed inside the sqrt (PyTorch convention), or the alternative is explicitly documented.
• Gamma and beta loads use col_offsets only — no row offset added.
• Output store applies the same boundary mask as the input load.
• Row stride is passed as a kernel argument, not computed as H.
• For training: mean and rstd are stored to separate output buffers indexed by row_idx.
• RMSNorm variant (if implemented) omits the mean subtraction step entirely.
• Correctness test compares against torch.nn.LayerNorm at both fp16 and fp32.
• No performance claims without a benchmark comparison.