Welford LayerNorm: Mean And Variance In One Read

Experimental — correctness is graded; the performance target awaits GPU calibration. The two-pass LayerNorm reads every row twice before it can normalize - and a third time to do it. Welford's algorithm (1962) maintains a running (count, mean, M2) triple that yields BOTH the mean and the variance in a single pass, with the numerical robustness the naive E[x^2] shortcut lacks. Chan's merge formula combines partial triples, so the block can tree-combine per-thread results. Add a shared row cache and the row touches global memory once in, once out - this is how production LayerNorms (PyTorch/Apex) are written.

Languages: CUDA C++, Rust

Lesson

Welford's one-pass update, per new value x: count += 1; delta = x - mean; mean += delta / count; M2 += delta * (x - mean). The subtlety is that M2 accumulates squared deviations from the RUNNING mean, and the two delta factors (before and after the mean update) are exactly what keeps it unbiased and numerically stable - unlike E[x^2] - mean^2, which on this harness's data (mean ~1000, spread ~0.1) loses the entire variance to float rounding and fails the correctness judge.

Parallelizing needs one more piece: merging two partial triples (nA, meanA, m2A) and (nB, meanB, m2B). Chan, Golub & LeVeque's formula does it: n = nA + nB; delta = meanB - meanA; mean = meanA + delta * nB / n; M2 = m2A + m2B + delta^2 * nA * nB / n. That makes Welford composable across the block's tree combine - the same composability trick as the online softmax's (m, d) pair, applied to second moments. With the row staged in shared memory during the single read pass, the normalize stage costs no further global reads.

Metric contract: algorithmic reads/writes count source-level sweeps. GLOBAL_MEMORY_SECTORS counts SASS-derived issued/requested sectors in the selected profiling scope, including repeated launches. It is cache-blind: it is neither L2 traffic nor device-memory bytes. Trace elapsed time is a separate measurement; a global load may hit cache.

The row cache must fit the supported row width. Welford combines mean/variance stably; reduced input requests do not automatically imply the same reduction in physical memory traffic.

What you will learn

Apply Welford's one-pass (count, mean, M2) update and explain why it is numerically robust.
Merge partial Welford triples with Chan's formula inside a shared-memory tree combine.
Stage the row in shared memory during the single read pass and normalize from the staged copy.
Verify the traffic win via GLOBAL_MEMORY_SECTORS, mirroring the online-softmax lab.

Starter code

    // The two-pass solution from the previous lab: numerically safe, but the
    // row crosses global memory three times before the write.
    __shared__ float red[256];

    int row = blockIdx.x;
    if (row >= rows) return;
    int tid = threadIdx.x;
    const float* src = X + row * cols;
    float* dst = Y + row * cols;

    float localSum = 0.0f;
    for (int c = tid; c < cols; c += blockDim.x) localSum += src[c];
    red[tid] = localSum;
    __syncthreads();
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) red[tid] += red[tid + stride];
        __syncthreads();
    }
    float mean = red[0] / cols;
    __syncthreads();

    float localVar = 0.0f;
    for (int c = tid; c < cols; c += blockDim.x) {
        float d = src[c] - mean;
        localVar += d * d;
    }
    red[tid] = localVar;
    __syncthreads();
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) red[tid] += red[tid + stride];
        __syncthreads();
    }
    float invStd = rsqrtf(red[0] / cols + eps);

    for (int c = tid; c < cols; c += blockDim.x) {
        dst[c] = (src[c] - mean) * invStd;
    }