RMSNorm: Fuse The Scale Into One Pass

Lesson

RMSNorm (Zhang & Sennrich, 2019) drops LayerNorm's mean subtraction: y_c = x_c * rsqrt(mean(x^2) + eps) * gamma_c. One reduction instead of two, no numerical cancellation hazard, measurably cheaper - which is why LLaMA-family decoders adopted it. That simplicity makes the traffic question stand alone: the naive shape computes the sum of squares in one block-stride traversal, then traverses the row a second time to apply the normalization and the per-column gamma scale. Both traversals stream the same 4 KB row through global memory.

The fusion is the same shared-row-cache discipline as the online-softmax and Welford labs: during the statistics pass, each thread also stores its elements into a shared array; after the tree reduction produces rsqrt(sumSq / cols + eps), the scale pass reads the SHARED copy. gamma is read once per column per block either way (it is broadcast-friendly and tiny next to the row stream). Traffic per row drops from 2 reads + 1 write to 1 read + 1 write.

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 starter already handles the reduction. The guided task is staging/reuse and fused scaling: y = x * rsqrt(mean(x*x) + epsilon) * weight.

What you will learn

• Know what RMSNorm computes and why LLaMA-family models prefer it over LayerNorm.
• Fuse a statistics pass and a scale pass by staging the row in shared memory.
• Apply a per-column affine parameter (gamma) inside the fused pass.
• Confirm the fusion via GLOBAL_MEMORY_SECTORS.

Starter code

    // One block per row: a parallel sum-of-squares reduction, then a SECOND
    // full read of the row from global memory to normalize and scale.
    __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 localSq = 0.0f;
    for (int c = tid; c < cols; c += blockDim.x) {
        float x = src[c];
        localSq += x * x;
    }
    red[tid] = localSq;
    __syncthreads();
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) red[tid] += red[tid + stride];
        __syncthreads();
    }
    float invRms = rsqrtf(red[0] / cols + eps);

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