From 800116cebd2d4cb794f2b5529152a6a123f9946f Mon Sep 17 00:00:00 2001 From: Kunal Mansukhani <24417742+kunal-mansukhani@users.noreply.github.com> Date: Fri, 16 Oct 2026 09:27:22 -0700 Subject: [PATCH] Add Fused Residual Add and LayerNorm with Saved Statistics challenge (hard #143) --- .../challenge.html | 104 ++++++++++++ .../challenge.py | 152 ++++++++++++++++++ .../starter/starter.cu | 6 + .../starter/starter.cute.py | 20 +++ .../starter/starter.jax.py | 17 ++ .../starter/starter.mojo | 22 +++ .../starter/starter.pytorch.py | 18 +++ .../starter/starter.triton.py | 20 +++ 8 files changed, 359 insertions(+) create mode 100644 challenges/hard/143_fused_residual_add_layernorm/challenge.html create mode 100644 challenges/hard/143_fused_residual_add_layernorm/challenge.py create mode 100644 challenges/hard/143_fused_residual_add_layernorm/starter/starter.cu create mode 100644 challenges/hard/143_fused_residual_add_layernorm/starter/starter.cute.py create mode 100644 challenges/hard/143_fused_residual_add_layernorm/starter/starter.jax.py create mode 100644 challenges/hard/143_fused_residual_add_layernorm/starter/starter.mojo create mode 100644 challenges/hard/143_fused_residual_add_layernorm/starter/starter.pytorch.py create mode 100644 challenges/hard/143_fused_residual_add_layernorm/starter/starter.triton.py diff --git a/challenges/hard/143_fused_residual_add_layernorm/challenge.html b/challenges/hard/143_fused_residual_add_layernorm/challenge.html new file mode 100644 index 00000000..7ce7b5b4 --- /dev/null +++ b/challenges/hard/143_fused_residual_add_layernorm/challenge.html @@ -0,0 +1,104 @@ +

+ Implement the fused pre-norm residual block entry used in GPT-2, BERT, and most LayerNorm + transformers: add the sublayer output x to the residual stream + residual, apply Layer Normalization with learned gamma and + beta, and in the same pass write out everything the next layer and the backward pass + need. All inputs are float32; x and residual have shape + (N, C) (one row per token), gamma and beta have shape + (C,). You must produce four outputs: the normalized output y + (N, C), the updated residual stream residual_out (N, C), + and the per-row mean and reciprocal standard deviation rstd, each of + shape (N,). +

+ + + + + + + + + + x + + residual + + + + + + + h + + row stats over C + μ, rstd + + + residual_out = h + + y + + mean, rstd + one read of x and residual, four outputs + + +

+ With \(\varepsilon\) = eps, for each row \(i\) and column \(j\): +

+

+ \[ + \begin{align} + h_{i,j} &= x_{i,j} + \text{residual}_{i,j} \\ + \mu_i &= \frac{1}{C}\sum_{k=0}^{C-1} h_{i,k}, \qquad + r_i = \frac{1}{\sqrt{\frac{1}{C}\sum_{k=0}^{C-1} \left(h_{i,k} - \mu_i\right)^2 + \varepsilon}} \\ + y_{i,j} &= \left(h_{i,j} - \mu_i\right) r_i \, \gamma_j + \beta_j + \end{align} + \] +

+

+ Store \(h\) in residual_out, \(\mu\) in mean, and \(r\) in + rstd. The variance is the biased (population) variance. Rows may have a large mean + relative to their spread, so compute the variance in a numerically stable way. +

+ +

Implementation Requirements

+ + +

Example

+

+ Input (N = 2, C = 4, eps = 1e-5): + \[ + x = \begin{bmatrix} 1.0 & 2.0 & 3.0 & 4.0 \\ 0.0 & 0.0 & 2.0 & 2.0 \end{bmatrix}, \quad + \text{residual} = \begin{bmatrix} 1.0 & 0.0 & -1.0 & 0.0 \\ 1.0 & 1.0 & -1.0 & -1.0 \end{bmatrix} + \] + \[ + \gamma = \begin{bmatrix} 1.0 & 2.0 & 1.0 & 0.5 \end{bmatrix}, \quad + \beta = \begin{bmatrix} 0.0 & 0.0 & 1.0 & 0.0 \end{bmatrix} + \] + Output: + \[ + y = \begin{bmatrix} -0.5773 & -1.1547 & 0.4227 & 0.8660 \\ 0.0000 & 0.0000 & 1.0000 & 0.0000 \end{bmatrix}, \quad + \text{residual\_out} = \begin{bmatrix} 2.0 & 2.0 & 2.0 & 4.0 \\ 1.0 & 1.0 & 1.0 & 1.0 \end{bmatrix} + \] + \[ + \text{mean} = \begin{bmatrix} 2.5 & 1.0 \end{bmatrix}, \quad + \text{rstd} = \begin{bmatrix} 1.1547 & 316.2278 \end{bmatrix} + \] +

+ +

Constraints

+ diff --git a/challenges/hard/143_fused_residual_add_layernorm/challenge.py b/challenges/hard/143_fused_residual_add_layernorm/challenge.py new file mode 100644 index 00000000..56be722e --- /dev/null +++ b/challenges/hard/143_fused_residual_add_layernorm/challenge.py @@ -0,0 +1,152 @@ +import ctypes +from typing import Any, Dict, List + +import torch +from core.challenge_base import ChallengeBase + + +class Challenge(ChallengeBase): + name = "Fused Residual Add and LayerNorm with Saved Statistics" + # y is O(|gamma| * 4 + |beta|). A tree/Welford reduction in float32 is accurate to ~1e-6 + # relative; the one-pass E[h^2] - E[h]^2 variance loses ~1e-2 on the large-offset tests, + # and dropping eps, beta, or the residual add is off by far more than 1e-3. + atol = 1e-03 + rtol = 1e-03 + num_gpus = 1 + access_tier = "free" + + def reference_impl( + self, + x: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + beta: torch.Tensor, + y: torch.Tensor, + residual_out: torch.Tensor, + mean: torch.Tensor, + rstd: torch.Tensor, + N: int, + C: int, + eps: float, + ): + assert x.shape == residual.shape == y.shape == residual_out.shape == (N, C) + assert gamma.shape == beta.shape == (C,) + assert mean.shape == rstd.shape == (N,) + for t in (x, residual, gamma, beta, y, residual_out, mean, rstd): + assert t.dtype == torch.float32 + + # The sum is rounded to float32 exactly as the solution must store it; statistics + # are computed in float64 from that rounded sum for a stable reference. + h = x + residual + hd = h.double() + mu = hd.mean(dim=1) + var = ((hd - mu[:, None]) ** 2).mean(dim=1) + r = torch.rsqrt(var + eps) + out = (hd - mu[:, None]) * r[:, None] * gamma.double() + beta.double() + residual_out.copy_(h) + y.copy_(out.float()) + mean.copy_(mu.float()) + rstd.copy_(r.float()) + + def get_solve_signature(self) -> Dict[str, tuple]: + return { + "x": (ctypes.POINTER(ctypes.c_float), "in"), + "residual": (ctypes.POINTER(ctypes.c_float), "in"), + "gamma": (ctypes.POINTER(ctypes.c_float), "in"), + "beta": (ctypes.POINTER(ctypes.c_float), "in"), + "y": (ctypes.POINTER(ctypes.c_float), "out"), + "residual_out": (ctypes.POINTER(ctypes.c_float), "out"), + "mean": (ctypes.POINTER(ctypes.c_float), "out"), + "rstd": (ctypes.POINTER(ctypes.c_float), "out"), + "N": (ctypes.c_int, "in"), + "C": (ctypes.c_int, "in"), + "eps": (ctypes.c_float, "in"), + } + + def _outputs(self, N, C): + # Pre-filled with garbage so solutions that skip any output (or any row of it) fail. + dev, f32 = self.device, torch.float32 + return { + "y": torch.full((N, C), 123.0, device=dev, dtype=f32), + "residual_out": torch.full((N, C), 123.0, device=dev, dtype=f32), + "mean": torch.full((N,), 123.0, device=dev, dtype=f32), + "rstd": torch.full((N,), 123.0, device=dev, dtype=f32), + } + + def _make(self, N, C, eps=1e-5, seed=0, scale=1.0, offset=0.0, g_range=(0.5, 1.5)): + torch.manual_seed(seed) + dev, f32 = self.device, torch.float32 + x = torch.randn(N, C, device=dev, dtype=f32) * scale + offset + residual = torch.randn(N, C, device=dev, dtype=f32) * scale + gamma = torch.empty(C, device=dev, dtype=f32).uniform_(*g_range) + beta = torch.empty(C, device=dev, dtype=f32).uniform_(-0.5, 0.5) + return { + "x": x, + "residual": residual, + "gamma": gamma, + "beta": beta, + **self._outputs(N, C), + "N": N, + "C": C, + "eps": eps, + } + + def generate_example_test(self) -> Dict[str, Any]: + dev, f32 = self.device, torch.float32 + return { + "x": torch.tensor([[1.0, 2.0, 3.0, 4.0], [0.0, 0.0, 2.0, 2.0]], device=dev, dtype=f32), + "residual": torch.tensor( + [[1.0, 0.0, -1.0, 0.0], [1.0, 1.0, -1.0, -1.0]], device=dev, dtype=f32 + ), + "gamma": torch.tensor([1.0, 2.0, 1.0, 0.5], device=dev, dtype=f32), + "beta": torch.tensor([0.0, 0.0, 1.0, 0.0], device=dev, dtype=f32), + **self._outputs(2, 4), + "N": 2, + "C": 4, + "eps": 1e-5, + } + + def generate_functional_test(self) -> List[Dict[str, Any]]: + tests = [] + + # Hand-checkable example: residual_out = [[2,2,2,4],[1,1,1,1]]. + tests.append(self.generate_example_test()) + + # C = 1: variance is 0, so y = beta and rstd = 1/sqrt(eps) (eps must be inside sqrt). + tests.append(self._make(1, 1, seed=1)) + + # Constant rows (x + residual identical across C): y = beta exactly. + c = self._make(3, 8, seed=2) + c["x"].fill_(2.5) + c["residual"].fill_(-1.0) + tests.append(c) + + # Large eps relative to the variance: catches implementations that ignore eps. + tests.append(self._make(4, 16, eps=1.0, seed=3, scale=0.3)) + + # Zero residual: plain LayerNorm, residual_out must equal x. + zr = self._make(5, 32, seed=4) + zr["residual"].zero_() + tests.append(zr) + + # Negative gamma, non-power-of-two widths, ragged tails. + tests.append(self._make(7, 30, seed=5, g_range=(-2.0, -0.5))) + tests.append(self._make(33, 1000, seed=6, scale=3.0, g_range=(-1.5, 1.5))) + + # Large DC offset (mean ~500, std ~1.4): the one-pass E[h^2] - E[h]^2 variance in + # float32 suffers catastrophic cancellation; requires two-pass or Welford. + tests.append(self._make(64, 1024, seed=7, offset=500.0)) + tests.append(self._make(17, 3000, seed=8, offset=-300.0, g_range=(-1.0, 1.0))) + + # Large magnitude (|h| ~ 100) and a tall, odd row count. + tests.append(self._make(4099, 256, seed=9, scale=70.0)) + + # Realistic LLM shapes: GPT-2 XL (1600) and a 5120-wide model. + tests.append(self._make(512, 1600, seed=10)) + tests.append(self._make(256, 5120, seed=11, offset=20.0)) + + return tests + + def generate_performance_test(self) -> Dict[str, Any]: + # 16,384 tokens (e.g. batch 4 x seq 4,096) with hidden size 4,096. + return self._make(16384, 4096, seed=0) diff --git a/challenges/hard/143_fused_residual_add_layernorm/starter/starter.cu b/challenges/hard/143_fused_residual_add_layernorm/starter/starter.cu new file mode 100644 index 00000000..feb1bfc9 --- /dev/null +++ b/challenges/hard/143_fused_residual_add_layernorm/starter/starter.cu @@ -0,0 +1,6 @@ +#include + +// x, residual, gamma, beta, y, residual_out, mean, rstd are device pointers +extern "C" void solve(const float* x, const float* residual, const float* gamma, const float* beta, + float* y, float* residual_out, float* mean, float* rstd, int N, int C, + float eps) {} diff --git a/challenges/hard/143_fused_residual_add_layernorm/starter/starter.cute.py b/challenges/hard/143_fused_residual_add_layernorm/starter/starter.cute.py new file mode 100644 index 00000000..0f1e2ce5 --- /dev/null +++ b/challenges/hard/143_fused_residual_add_layernorm/starter/starter.cute.py @@ -0,0 +1,20 @@ +import cutlass +import cutlass.cute as cute + + +# x, residual, gamma, beta, y, residual_out, mean, rstd are tensors on the GPU +@cute.jit +def solve( + x: cute.Tensor, + residual: cute.Tensor, + gamma: cute.Tensor, + beta: cute.Tensor, + y: cute.Tensor, + residual_out: cute.Tensor, + mean: cute.Tensor, + rstd: cute.Tensor, + N: cute.Int32, + C: cute.Int32, + eps: cute.Float32, +): + pass diff --git a/challenges/hard/143_fused_residual_add_layernorm/starter/starter.jax.py b/challenges/hard/143_fused_residual_add_layernorm/starter/starter.jax.py new file mode 100644 index 00000000..3cd11980 --- /dev/null +++ b/challenges/hard/143_fused_residual_add_layernorm/starter/starter.jax.py @@ -0,0 +1,17 @@ +import jax +import jax.numpy as jnp + + +# x, residual, gamma, beta are tensors on device +@jax.jit +def solve( + x: jax.Array, + residual: jax.Array, + gamma: jax.Array, + beta: jax.Array, + N: int, + C: int, + eps: float, +) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]: + # return (y, residual_out, mean, rstd) directly + pass diff --git a/challenges/hard/143_fused_residual_add_layernorm/starter/starter.mojo b/challenges/hard/143_fused_residual_add_layernorm/starter/starter.mojo new file mode 100644 index 00000000..1cbc1627 --- /dev/null +++ b/challenges/hard/143_fused_residual_add_layernorm/starter/starter.mojo @@ -0,0 +1,22 @@ +from max.gpu.host import DeviceContext +from max.gpu import block_dim, block_idx, thread_idx +from std.memory import UnsafePointer +from std.math import ceildiv + + +# x, residual, gamma, beta, y, residual_out, mean, rstd are device pointers +@export +def solve( + x: UnsafePointer[Float32, MutAnyOrigin], + residual: UnsafePointer[Float32, MutAnyOrigin], + gamma: UnsafePointer[Float32, MutAnyOrigin], + beta: UnsafePointer[Float32, MutAnyOrigin], + y: UnsafePointer[Float32, MutAnyOrigin], + residual_out: UnsafePointer[Float32, MutAnyOrigin], + mean: UnsafePointer[Float32, MutAnyOrigin], + rstd: UnsafePointer[Float32, MutAnyOrigin], + N: Int32, + C: Int32, + eps: Float32, +) raises: + pass diff --git a/challenges/hard/143_fused_residual_add_layernorm/starter/starter.pytorch.py b/challenges/hard/143_fused_residual_add_layernorm/starter/starter.pytorch.py new file mode 100644 index 00000000..5baa1081 --- /dev/null +++ b/challenges/hard/143_fused_residual_add_layernorm/starter/starter.pytorch.py @@ -0,0 +1,18 @@ +import torch + + +# x, residual, gamma, beta, y, residual_out, mean, rstd are tensors on the GPU +def solve( + x: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + beta: torch.Tensor, + y: torch.Tensor, + residual_out: torch.Tensor, + mean: torch.Tensor, + rstd: torch.Tensor, + N: int, + C: int, + eps: float, +): + pass diff --git a/challenges/hard/143_fused_residual_add_layernorm/starter/starter.triton.py b/challenges/hard/143_fused_residual_add_layernorm/starter/starter.triton.py new file mode 100644 index 00000000..8648aa6b --- /dev/null +++ b/challenges/hard/143_fused_residual_add_layernorm/starter/starter.triton.py @@ -0,0 +1,20 @@ +import torch +import triton +import triton.language as tl + + +# x, residual, gamma, beta, y, residual_out, mean, rstd are tensors on the GPU +def solve( + x: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + beta: torch.Tensor, + y: torch.Tensor, + residual_out: torch.Tensor, + mean: torch.Tensor, + rstd: torch.Tensor, + N: int, + C: int, + eps: float, +): + pass