kunal-mansukhani · Oct 16, 2026 · Oct 16, 2026
@@ -0,0 +1,104 @@
+<p>
+  Implement the fused pre-norm residual block entry used in GPT-2, BERT, and most LayerNorm
+  transformers: add the sublayer output <code>x</code> to the residual stream
+  <code>residual</code>, apply Layer Normalization with learned <code>gamma</code> and
+  <code>beta</code>, and in the same pass write out everything the next layer and the backward pass
+  need. All inputs are float32; <code>x</code> and <code>residual</code> have shape
+  <code>(N, C)</code> (one row per token), <code>gamma</code> and <code>beta</code> have shape
+  <code>(C,)</code>. You must produce four outputs: the normalized output <code>y</code>
+  <code>(N, C)</code>, the updated residual stream <code>residual_out</code> <code>(N, C)</code>,
+  and the per-row <code>mean</code> and reciprocal standard deviation <code>rstd</code>, each of
+  shape <code>(N,)</code>.
+</p>
+
+<svg width="640" height="200" viewBox="0 0 640 200" xmlns="http://www.w3.org/2000/svg"
+     style="display:block; margin:20px auto;">
+  <defs>
+    <marker id="ah" markerWidth="8" markerHeight="6" refX="7" refY="3" orient="auto">
+      <polygon points="0 0,8 3,0 6" fill="#999"/>
+    </marker>
+  </defs>
+  <rect width="640" height="200" fill="#222" rx="8"/>
+  <rect x="20" y="40" width="100" height="36" fill="#1a3a5c" stroke="#3a6a9c" rx="4"/>
+  <text x="70" y="63" fill="#ddd" font-family="monospace" font-size="13" text-anchor="middle">x</text>
+  <rect x="20" y="110" width="100" height="36" fill="#1a3a5c" stroke="#3a6a9c" rx="4"/>
+  <text x="70" y="133" fill="#ddd" font-family="monospace" font-size="13" text-anchor="middle">residual</text>
+  <circle cx="175" cy="93" r="16" fill="#333" stroke="#999"/>
+  <text x="175" y="98" fill="#ddd" font-family="monospace" font-size="16" text-anchor="middle">+</text>
+  <line x1="120" y1="58" x2="162" y2="84" stroke="#999" stroke-width="1.5" marker-end="url(#ah)"/>
+  <line x1="120" y1="128" x2="162" y2="102" stroke="#999" stroke-width="1.5" marker-end="url(#ah)"/>
+  <line x1="191" y1="93" x2="250" y2="93" stroke="#999" stroke-width="1.5" marker-end="url(#ah)"/>
+  <text x="220" y="84" fill="#aaa" font-family="monospace" font-size="11" text-anchor="middle">h</text>
+  <rect x="252" y="70" width="160" height="46" fill="#2d5a27" rx="4"/>
+  <text x="332" y="90" fill="#ddd" font-family="monospace" font-size="12" text-anchor="middle">row stats over C</text>
+  <text x="332" y="106" fill="#ddd" font-family="monospace" font-size="12" text-anchor="middle">&#956;, rstd</text>
+  <line x1="412" y1="93" x2="470" y2="93" stroke="#999" stroke-width="1.5" marker-end="url(#ah)"/>
+  <rect x="472" y="20" width="150" height="34" fill="#7a4a1a" rx="4"/>
+  <text x="547" y="42" fill="#ddd" font-family="monospace" font-size="12" text-anchor="middle">residual_out = h</text>
+  <rect x="472" y="76" width="150" height="34" fill="#7a4a1a" rx="4"/>
+  <text x="547" y="98" fill="#ddd" font-family="monospace" font-size="12" text-anchor="middle">y</text>
+  <rect x="472" y="132" width="150" height="34" fill="#7a4a1a" rx="4"/>
+  <text x="547" y="154" fill="#ddd" font-family="monospace" font-size="12" text-anchor="middle">mean, rstd</text>
+  <text x="320" y="190" fill="#999" font-family="monospace" font-size="11" text-anchor="middle">one read of x and residual, four outputs</text>
+</svg>
+
+<p>
+  With \(\varepsilon\) = <code>eps</code>, for each row \(i\) and column \(j\):
+</p>
+<p>
+  \[
+  \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}
+  \]
+</p>
+<p>
+  Store \(h\) in <code>residual_out</code>, \(\mu\) in <code>mean</code>, and \(r\) in
+  <code>rstd</code>. 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.
+</p>
+
+<h2>Implementation Requirements</h2>
+<ul>
+  <li>Use only native features (external libraries are not permitted)</li>
+  <li>The <code>solve</code> function signature must remain unchanged</li>
+  <li>The results must be stored in <code>y</code>, <code>residual_out</code>, <code>mean</code>,
+      and <code>rstd</code></li>
+  <li>The outputs are not zero-initialized; every element must be written</li>
+</ul>
+
+<h2>Example</h2>
+<p>
+  Input (<code>N</code> = 2, <code>C</code> = 4, <code>eps</code> = 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}
+  \]
+</p>
+
+<h2>Constraints</h2>
+<ul>
+  <li>1 &le; <code>N</code> &le; 65,536</li>
+  <li>1 &le; <code>C</code> &le; 16,384</li>
+  <li>All tensors are float32</li>
+  <li>Values of <code>x</code> and <code>residual</code> are in [&minus;1,000, 1,000]; values of
+      <code>gamma</code> are in [&minus;2, 2] and <code>beta</code> in [&minus;0.5, 0.5]</li>
+  <li>1e-6 &le; <code>eps</code> &le; 1.0</li>
+  <li>Performance is measured with <code>N</code> = 16,384, <code>C</code> = 4,096</li>
+</ul>
@@ -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)
@@ -0,0 +1,6 @@
+#include <cuda_runtime.h>
+
+// 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) {}
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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