Add RMS Norm Backward challenge (medium #142)

[claude]

Summary

Adds RMS Norm Backward (challenges/medium/142_rms_norm_backward), the RMSNorm training kernel used in every LLaMA/Mistral/Qwen block. Given x (N, C), weight (C,) and upstream dy (N, C), the solver computes dx (N, C) and dweight (C,). The per-row rstd is not saved from the forward pass, so the solver has to recompute it.

Number 142 is the next one after the merged 140/141. 134–139 belong to the open PR #349, and #236 covers multi-GPU vector add, so neither number nor topic overlaps.

What a solver learns

• Two reductions along different axes in one kernel. dx needs per-row reductions (Σx² and Σ dy·w·x) and dweight needs a reduction down each column over all N tokens. Doing both in a single pass over memory is the core difficulty of norm-backward kernels in Apex, Liger and Triton's layer-norm tutorial.
• Cutting cross-block synchronization in a gradient reduction. dweight is a reduction over the whole batch. Solvers choose between per-element atomics, a column-per-thread kernel, and per-block register partials followed by a small second-stage reduce.
• Algebraic fusion. Both row statistics come out of one block reduction because Σ g·x̂ = r·Σ g·x, so dx = r·g − x·r³·Σ(g·x)/C.

Novelty compared with the closest challenges

• 50 RMS Normalization, 83 Fused Residual Add + RMS Norm, 113 Layer Normalization, 131 Fused LayerNorm Linear GELU: forward only. Each output row depends on one row reduction and there is no cross-row work. The backward adds the projection term (dx depends on Σ_k g_k x̂_k) and, more importantly, a reduction over the batch axis that crosses blocks.
• 111 Softmax Attention Backward: the backward is mostly GEMMs. This challenge is bandwidth-bound and is about combining reductions in two directions.
• 127 Fused AdamW: a fused element-wise optimizer step with no reductions.
• 4 Reduction / 136 (PR [DO NOT MERGE] Kernel Interview Track: curriculum + 6 'why is this slow' challenges (134-139) #349) row sum: single-axis reductions with no coupling between axes.

Naive vs. optimized

• Naive (correct): kernel A, one block per row, computes rstd and the projection and writes dx. Kernel B, one thread per column, loops over all N rows, recomputes or reloads rstd, and accumulates dy·x̂. Alternatively, kernel A atomicAdds into dweight for every element.
• Optimized (the validation solution): a persistent grid of about 4×#SM blocks. Each block walks a strided set of rows and keeps the row of x/dy in registers. One fused block reduction gives Σx² and Σ dy·w·x, then the block writes dx and accumulates dy·x̂ into per-thread register partials for the columns that thread owns. At the end each block writes one C-length partial, and a tiny kernel sums the G × C partials. Results match the naive version up to fp32 summation order.

How the benchmark exposes it

Performance is measured with N = 16,384, C = 4,096 (batch 4 × seq 4,096 at LLaMA-7B hidden size). That is 256 MB per tensor and 768 MB in total, which fits 5× in 16 GB. The op is purely memory-bound: about 2 FLOPs per byte.

• The minimum traffic is reading x and dy once and writing dx once (3·N·C·4 B ≈ 805 MB). The two-kernel naive version reads x and dy a second time (+537 MB, about 1.7× the traffic).
• With only C = 4,096 columns, a column-per-thread dweight kernel launches just 4,096 threads, about 3 warps per SM on a 40-SM T4, and each runs a serial loop of 16,384 iterations. That leaves the GPU latency-bound and badly underused.
• Per-element atomicAdd puts 67M atomics on only 4,096 addresses. Per-block register partials need just G·C (about 0.6M) writes plus a reduction smaller than 3 MB.
• dy is correlated with x in the benchmark, so the projection term is O(1) and can't be skipped without failing correctness.

Functional tests and the mistakes they catch

12 cases. Outputs are pre-filled with 123.0, so every element has to be written. Each mistake below was checked locally by running it as a mutant against the reference.

Mistake	Caught by
Dropping the projection term − x̂·mean(g·x̂) (treating RMSNorm as a pure scale)	N=C=1 case (exact dx ≈ 0), correlated-dy cases, almost all others
Using dy instead of dy·weight in dx	example, negative-weight case, most others
Using x instead of x̂ in dweight	large-magnitude x (σ=50) case, most others
Missing the outer rstd factor in dx	large-magnitude case, most others
Ignoring eps, or adding it outside the sqrt	all-zero x case (otherwise inf/nan), eps = 1.0 case
Dividing by N instead of C	non-square 100×30, 255×1000 and others
Accumulating into dweight without initializing it (atomics onto an assumed-zero buffer)	every case (garbage-filled outputs)
Overwriting rather than accumulating partials across row blocks	4099×256 (ragged multi-block) and 1024×2048
Column-tile tails / non-power-of-two C	C = 1, 3, 30, 1000

Other coverage: zero inputs (all-zero x, all-zero dy, which must give exactly 0), negative weights and values, powers of two (16×64, 32×128, 256, 1024×2048) and non-powers of two.

Independent check: the example values were worked out by hand (eps≈0: dx = [[1.375, 1.625, 0.875, 0.625], [0, 0, 0, 1]], dweight = [3, -1, 1, -1]), checked with float64 scalar loops, and matched by the reference. Every functional case was also compared against float64 torch.autograd of the forward pass. The largest absolute deviation was 8.3e-4, on the 4,099-row case, where dweight values are about 4e3, so it is well within the relative tolerance.

Tolerance (atol = rtol = 1e-3): fp32 reordering error is about 1e-6 relative, while each mistake in the table shifts results by at least 1e-2 on its targeted case.

Validation

• scripts/run_challenge.py --language cuda on NVIDIA Tesla T4: --action run gave ✓ Test passed, and --action submit gave ✓ All tests passed, including the performance test. The solution file is not committed.
• pre-commit run --all-files passes, and scripts/validate_challenges.py reports 0 errors.
• All CLAUDE.md checklist items were verified: HTML starts with <p>, has <h2> sections, includes an SVG on #222, the bmatrix example matches generate_example_test(), and the perf bullet matches. All 6 starters have exactly one parameter comment (no parenthetical, since this is medium); the JAX starter also has # return output tensor directly. Allocations use device=self.device, and the Mojo starter uses 1.1 style.

🤖 Generated with Claude Code

[github-actions]

✅ Review bot: low risk, safe to merge. Lint and structure checks passed; no sensitive paths or grader edits; 161 lines changed.

• ID 142 is unused. The open PR's 134–139 are not in the base list, and the topic (RMSNorm backward with a cross-row dweight reduction) does not duplicate any existing challenge (forward-only 50/83/113/131, 111 is GEMM-heavy).
• challenge.py defines the 5 required methods, has no init, and sets name/atol/rtol/num_gpus/access_tier. It adds only two private helpers (_outputs, _make). Imports are limited to ctypes, typing, torch and the challenge base, with nothing running at import time.
• reference_impl asserts shapes and dtypes with no device checks, and uses only standard torch ops. I verified the example by hand (dx and dweight match the html), and the formulas match the html.
• All tensor allocations use device=self.device. Outputs are pre-filled with 123.0, which forces solutions to write every element.
• Functional tests cover N=1/C=1, a single row, all-zero x, large eps, zero dy, negative weights, large x magnitude, non-power-of-two shapes, a ragged row count and a realistic shape. A 1e-3 tolerance is reasonable for float32.
• The performance test (N=16384, C=4096) uses about 805 MB for x, dy and dx plus small buffers, well under the 3.2 GB limit.
• The starters compile as stubs with only the parameter comment, and the Mojo starter uses max.gpu imports and MutAnyOrigin. The html constraints and example match the signature, and the diff is complete with no injection or obfuscation.