Write Cuda Layernorm Kernel

Skill: Write CUDA LayerNorm Kernel

Purpose

Guide the agent through designing and implementing a correct, efficient CUDA LayerNorm (and RMSNorm) kernel, covering mean/variance computation strategies, Welford online accumulation, epsilon placement, affine transform application, backward pass structure, and decomposition for non-power-of-two hidden dimensions.

Use this when

• You need a custom LayerNorm kernel with a non-standard normalization axis, fused activation, or non-standard epilogue
• You are implementing RMSNorm (no mean subtraction, only variance normalization) and need a custom kernel
• You need the backward pass and require explicit control over the gradient computation for numerical accuracy or fusion
• The hidden dimension size does not match the assumptions of library implementations, or you need custom handling of non-power-of-two sizes
• You are fusing LayerNorm with a preceding or following operation and need to avoid extra memory round-trips

Do not use this when

• Standard forward-only LayerNorm with affine parameters on common hidden dimensions: use torch.nn.LayerNorm (which calls into cuDNN or a well-tuned CUDA kernel) or APEX's FusedLayerNorm
• You are using PyTorch with autograd: the built-in LayerNorm already has a correct, optimized backward pass
• The normalization is applied over batch dimensions (BatchNorm): this skill covers LayerNorm (normalization over the last D dimensions within a single sample)

Inputs the agent should gather first

• Input shape: [batch, seq, hidden] or [N, D] where D is the normalization dimension. What is the hidden dimension size? Is it always the same or variable at runtime?
• Normalization axis: which dimension(s) are normalized? LayerNorm normalizes over the last D dimensions; confirm this assumption.
• Dtype: fp32, fp16, or bf16 for input/output. What dtype are the gamma/beta parameters? (Usually fp32 even when input is fp16.)
• Affine parameters: are gamma (scale) and beta (shift) parameters present (LayerNorm) or absent (plain normalization)?
• Epsilon: what value? Typically 1e-5 for LayerNorm, 1e-6 for RMSNorm. Is it applied inside or outside the sqrt?
• Forward only vs forward + backward: does the kernel need to support gradient computation? If yes, what tensors need to be saved from the forward pass?
• RMSNorm variant: should mean subtraction be skipped (RMSNorm only divides by the root mean square, not the standard deviation)?
• Hardware target: SM architecture.

Required reasoning process

1. Determine the normalization work per row. Each row has D elements (the hidden dimension). A single row is normalized independently of all other rows. Assign one thread block (or one warp for small D) to each row. Grid size = number of rows = N * seq_len (or batch * seq if 3D input).

2. Choose the accumulation strategy for mean and variance.

   • Two-pass: first pass computes the mean mu = sum(x) / D; second pass computes the variance var = sum((x - mu)^2) / D. Requires two reads of the row. Straightforward to implement correctly.
   • Welford online algorithm: single pass computes both mean and variance simultaneously using the recurrence mean_new = mean_old + (x - mean_old) / n, M2_new = M2_old + (x - mean_old) * (x - mean_new), with var = M2 / D. Requires one read of the row. Numerically more stable than the naive E[x^2] - E[x]^2 formula, which suffers from catastrophic cancellation when variance is small relative to the mean. For production kernels: use Welford for its single-pass and numerical stability properties. For pedagogical or simple cases: two-pass is easier to verify.

3. Design the warp+block reduction for Welford. Welford accumulators (count, mean, M2) must be merged at the warp and block level. Parallel Welford merge: given two accumulators (n_a, mean_a, M2_a) and (n_b, mean_b, M2_b), the combined accumulator is:

   n = n_a + n_b
   delta = mean_b - mean_a
   mean = mean_a + delta * n_b / n
   M2 = M2_a + M2_b + delta * delta * n_a * n_b / n
   

   This merge is used at each step of the warp shuffle reduction and the smem block reduction. For simple sum reductions (two-pass approach), standard sum reductions suffice.

4. Handle non-power-of-two D. The reduction pattern must correctly handle D values that are not multiples of the block size. Threads with index threadIdx.x + k * blockDim.x >= D must not read from global memory and must contribute the identity accumulator to the Welford merge.

5. Compute the normalization. After the reduction:

   mean = accumulated_sum / D   (or from Welford directly)
   var = M2 / D                 (from Welford, or second-pass sum-of-squared-deviations / D)
   inv_std = rsqrtf(var + epsilon)
   

   Each element is then normalized: y_i = (x_i - mean) * inv_std. For RMSNorm: skip mean computation; inv_std = rsqrtf(mean_of_squares + epsilon).

6. Apply the affine transform. If gamma and beta are present: y_i = y_i * gamma[i] + beta[i]. The gamma and beta arrays have size D; each thread reads the gamma/beta value for its assigned position. These are typically loaded from global memory once per thread per call.

7. Epsilon placement. Epsilon must be added inside the sqrt (i.e., rsqrtf(var + epsilon), not rsqrtf(var) + epsilon). The latter formula does not prevent division by zero and provides incorrect normalization when variance is near zero.

8. Handle the backward pass (if required). The LayerNorm backward computes:

   • dL/dgamma_i = sum over rows of (normalized_x_i * dL/dy_i)
   • dL/dbeta_i = sum over rows of dL/dy_i
   • dL/dx_i = (1/D) * inv_std * (D * dL/dy_i_hat - sum(dL/dy_hat) - x_hat_i * sum(dL/dy_hat * x_hat)) where x_hat_i = (x_i - mean) * inv_std is the normalized input. The backward pass needs the normalized input x_hat or the saved mean and inv_std from the forward pass. Save those in the forward pass if backward is needed.

9. Handle fp16/bf16 inputs. Convert fp16/bf16 inputs to fp32 immediately on load. All accumulation (sum, sum-of-squares, mean, variance) must be in fp32. Gamma and beta should be fp32 regardless of input dtype. The output is converted back to fp16/bf16 at the store step.

Kernel design rules

• Always use rsqrtf(var + epsilon) — epsilon inside the sqrt argument. Never add epsilon outside the sqrt.
• For fp16/bf16 inputs: accumulate entirely in fp32. Convert on load, convert back on store.
• Welford parallel merge must be used when combining warp-level accumulators into block-level accumulators. Do not sum means independently — it produces the wrong result when different threads hold different counts.
• Thread block size: 128 or 256 threads for D >= 128. For D < 32, use one warp per row (blockDim.x = 32, grid = N). For D in [32, 128], one warp or 128-thread block per row.
• Gamma and beta reads should be coalesced: consecutive threads read consecutive elements of gamma and beta. This is naturally the case when threadIdx.x indexes directly into the D dimension.
• Do not use __syncthreads more than necessary. For Welford with a pure warp (D <= 32), no smem is needed at all.
• If the hidden dimension D exceeds the block size, each thread processes multiple elements in a stride loop: for (int i = threadIdx.x; i < D; i += blockDim.x). This must be done consistently for both the accumulation pass and the normalization write pass.
• RMSNorm: compute mean_sq = sum(x_i^2) / D, then inv_rms = rsqrtf(mean_sq + epsilon). Apply y_i = x_i * inv_rms * gamma[i]. No beta term for RMSNorm.

Correctness requirements

• Epsilon placement: epsilon must be inside the variance sqrt argument. Verify by inspection that the formula is rsqrtf(var + eps).
• Welford merge correctness: the parallel merge formula must be used when combining partial Welford accumulators across warps or threads. Averaging the means without weighted combination is incorrect and produces wrong variance for non-uniform load distributions.
• fp16 variance underflow: for fp16 inputs with small activation values, the variance computed in fp16 may underflow to 0, making inv_std = rsqrtf(0 + epsilon) = 1/sqrt(epsilon), which is a valid but imprecise result. This is prevented by accumulating in fp32.
• Thread synchronization: the smem write (warp leaders writing Welford partials) must be followed by __syncthreads() before any reads.
• Boundary handling: for D not a multiple of blockDim.x, threads with out-of-range indices must not read gamma/beta or the input tensor, and must contribute identity accumulators (count=0 for Welford).
• Gamma/beta indexing: gamma[i] and beta[i] are indexed by position within the D dimension, not by the global thread index. Verify that the index into gamma/beta is threadIdx.x + k * blockDim.x (the position within the row), not the global thread index.
• Backward pass saved tensors: if the backward pass recomputes x_hat from saved mean and inv_std, verify that the saved tensors use the same epsilon and have the same dtype as used in the backward computation.

Performance requirements

• LayerNorm is memory-bandwidth-bound for large hidden dimensions. Target at least 80% of device peak memory bandwidth.
• A single read of the input row is preferred (Welford). Two reads (two-pass) is acceptable for clarity but costs roughly 2x memory bandwidth for the input.
• Vectorized loads: use float4 for fp32 inputs where D is divisible by 4, or half2/half4 for fp16 inputs where alignment allows. This reduces load instruction count and improves effective bandwidth.
• For the forward pass, gamma and beta reads add D * 2 extra reads (assuming they fit in L2 cache across calls). For large batch sizes, they will be cached; for single-sample inference, they may not be. This is typically a small fraction of total memory traffic.
• Occupancy: LayerNorm kernels use a small amount of smem (one Welford triplet per warp per block). Occupancy should be high (50–100% theoretical). If register usage is high due to Welford state, consider annotating with __launch_bounds__.
• For the common transformer case (D = 768, 1024, 2048, 4096): ensure the thread block size divides D or handles the non-divisible case with a stride loop. Power-of-two block sizes (128, 256) are preferred because they simplify the reduction tree.

Output format

The final response must include:

1. Decomposition decision: rows-per-block assignment, thread block size, justification.
2. Algorithm choice: Welford online or two-pass, with reasons.
3. Kernel code: complete, compilable CUDA kernel for the forward pass. If backward is requested, a second kernel.
4. Warp+block Welford reduction helper: device function implementing the parallel Welford merge.
5. Host dispatch: grid/block dims, kernel call, saved tensor list for backward.
6. Correctness notes: epsilon placement, fp16 accumulation strategy, boundary handling.
7. RMSNorm variant (if requested): separate kernel or a compile-time template flag.
8. Known limitations: maximum D for the chosen block size, non-divisible D behavior, backward pass numerical sensitivity.

Common failure modes

• Variance underflow at fp16: computing variance in fp16 when input values are small (e.g., near-zero activations) causes var = 0, making inv_std = 1/sqrt(epsilon). This is mathematically valid but loses all information. Fix: always accumulate variance in fp32.
• Incorrect epsilon placement (outside sqrt): 1.0f / (sqrtf(var) + epsilon) does not prevent division by zero when var is exactly 0 — epsilon must be inside the argument to rsqrtf. This is a subtle but important difference.
• Wrong gamma/beta broadcast: gamma/beta are 1D tensors of size D, not size N * D. A common bug indexes them with the global thread index instead of the position-within-row index, causing gamma[0] to be applied to the first element of every row, gamma[1] to the second element, etc. — but only if the global thread index happens to match the D index, which it may not for all rows.
• Non-parallel-merge of Welford accumulators: summing partial means (each warp's local mean) without weighting by the count produces a biased estimate of the true mean when each warp processes a different number of elements (which happens when D is not a multiple of blockDim.x * numWarps). Fix: use the weighted parallel Welford merge formula.
• Missing __syncthreads between smem write and read: warp leaders write Welford partials to smem, but the first warp starts reading before all warps have written. Fix: __syncthreads() after the write, before the read.
• Two-pass mean drift: computing the variance in the second pass using (x_i - mean)^2 with a mean that was computed in floating-point can accumulate errors for large D. Welford avoids this by maintaining a numerically stable running estimate.
• Backward pass shape mismatch: gamma gradients must be summed over the batch and sequence dimensions (all rows). A common bug reduces only over the batch but not the sequence, producing wrong gamma gradients for 3D inputs.

Review checklist

• Is epsilon placed inside the rsqrtf argument — rsqrtf(var + epsilon) — not outside?
• For fp16/bf16 inputs: is all accumulation (sum, variance, Welford state) in fp32?
• Is the Welford parallel merge formula used when combining partial accumulators across warps?
• Is gamma/beta indexed by the position within the D dimension (not the global thread index)?
• Are threads beyond position D excluded from global reads and contributing identity accumulators?
• Is there a __syncthreads() after warp leaders write Welford partials to smem?
• For D not divisible by blockDim.x: does the stride loop handle the final partial tile correctly?
• For RMSNorm: is mean subtraction absent, and is the accumulator computing sum-of-squares (not sum)?
• For the backward pass: are mean and inv_std saved from the forward pass, or is x_hat explicitly saved?
• Has the kernel been validated against torch.nn.LayerNorm on: D=64, D=128, D=512, D=768, D=1024, D=2049 (non-power-of-two), small variance (near-constant rows), zero-variance rows?
• Is APEX FusedLayerNorm or cuDNN LayerNorm considered and explicitly noted for standard use cases?