LayerNorm: One Block Owns One Row

LayerNorm sits before or after every attention and MLP block in a transformer: normalize each row to zero mean and unit variance, guarded by a small epsilon. The starter assigns one THREAD per row - with this launch that leaves nearly every thread idle and each working thread walking 1024 columns serially, three times. Rewrite it so one BLOCK owns one row and cooperates on both reductions.

Languages: CUDA C++, Rust, HIP C++

Lesson

LayerNorm(x) = (x - mean) * rsqrt(variance + eps), computed independently per row. Structurally it is two chained reductions (sum for the mean, then sum of squared deviations for the variance) followed by an elementwise pass - the same block-per-row shape as the softmax lab, with addition both times. This lab's two-pass form computes variance the numerically safe way: subtract the mean FIRST, then square. The tempting one-pass shortcut variance = E[x^2] - mean^2 is algebraically identical and numerically treacherous: this harness's data (values near 1000 with a spread of 0.1) makes float E[x^2] so large that the true variance drowns in rounding error - a kernel using the shortcut fails the correctness check outright. The honest one-pass method (Welford's algorithm) is the sequel lab.

The reused shared scratch needs the same barrier discipline as the softmax lab: reading the mean out of red[0] must be separated by a barrier from overwriting the scratch with squared-deviation partials.

LayerNorm subtracts the mean; RMSNorm does not. Transformer architectures choose different norms, so keep the requested formula and epsilon rather than substituting another architecture norm.

What you will learn

Map LayerNorm to the block-per-row pattern: two chained sum reductions plus a normalize pass.
Compute variance safely: subtract the mean before squaring.
Know why the E[x^2] - mean^2 shortcut catastrophically cancels on large-mean data.
Reuse shared reduction scratch across stages with correct barrier placement.

Starter code

    // One THREAD per row: correct, but at this one-block-per-row launch almost
    // every thread computes nothing, and the ones that do walk the row serially
    // three times.
    int row = blockIdx.x * blockDim.x + threadIdx.x;
    if (row < rows) {
        const float* src = X + row * cols;
        float mean = 0.0f;
        for (int c = 0; c < cols; ++c) mean += src[c];
        mean /= cols;

        float var = 0.0f;
        for (int c = 0; c < cols; ++c) {
            float d = src[c] - mean;
            var += d * d;
        }
        var /= cols;

        float invStd = rsqrtf(var + eps);
        for (int c = 0; c < cols; ++c) Y[row * cols + c] = (src[c] - mean) * invStd;
    }