1.36. RMS Normalization Operation

Calculates RMS normalization by computing the square root of the mean squared values along the feature dimension and normalizing the input tensor with numerical stability using a small epsilon.

Progress

Step 2 of 2

Reward

1.80x

Optimization Timeline

Step 2 Details

Reasoning

The RMSNorm operation computes, for each spatial position (b, i, j) in a tensor x of shape (B, F, H, W), the RMS over the feature axis: rms = sqrt((1/F * ∑₍f₌₀₎^(F−1) x[b,f,i,j]²) + eps) and then normalizes each element via: y[b,f,i,j] = x[b,f,i,j] / rms. Our initial fused CUDA kernel loops over the F dimension in two separate loops (one for accumulation, one for normalization). Since F is often moderate (e.g., 64) and usually divisible by 4, we can unroll the inner loops by a factor of 4, which should reduce the loop overhead and increase throughput by doing fewer iterations while taking advantage of vectorized loads from global memory. We modify the CUDA kernel to check if (F % 4 == 0). If so, we unroll the loops by processing 4 channels at a time; otherwise, we fall back to the standard loop. This optimization will genuinely help reduce the arithmetic and loop overhead in each thread, leading to better runtime as observed in microbenchmarks. The Python module remains largely unchanged, as it simply loads our extension and wraps it in a CUDAModel.

#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

// RMS normalization kernel with loop unrolling if F (num_features) is divisible by 4.
// Each thread handles one spatial location (i.e., one pair (i,j) for each batch element).
// Let spatial = H * W, and the input tensor has shape (B, F, H, W). The element x[b, f, i, j] is accessed
// as: x[(b * F + f) * spatial + pos], where pos = i * W + j.
__global__ void rms_norm_kernel(const float* __restrict__ x,
                                float* __restrict__ y,
                                int B,
                                int F,
                                int H,
                                int W,
                                float eps) {
    int spatial = H * W;
    int total = B * spatial;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < total) {
        int b = idx / spatial;
        int pos = idx % spatial;
        float sum = 0.0f;
        
        // Unroll if F is divisible by 4.
        if ((F & 3) == 0) {
            for (int f = 0; f < F; f += 4) {
                int base = (b * F + f) * spatial + pos;
                float v0 = x[base];
                float v1 = x[base + spatial];
                float v2 = x[base + 2 * spatial];
                float v3 = x[base + 3 * spatial];
                sum += v0 * v0 + v1 * v1 + v2 * v2 + v3 * v3;
            }
        } else {
            for (int f = 0; f < F; f++) {
                int index = (b * F + f) * spatial + pos;
                float v = x[index];
                sum += v * v;
            }
        }
        
        float mean = sum / (float)F;
        float norm = sqrtf(mean + eps);
        
        // Normalize the input.
        if ((F & 3) == 0) {
            for (int f = 0; f < F; f += 4) {
                int base = (b * F + f) * spatial + pos;
                float v0 = x[base];
                float v1 = x[base + spatial];
                float v2 = x[base + 2 * spatial];
                float v3 = x[base + 3 * spatial];
                y[base] = v0 / norm;
                y[base + spatial] = v1 / norm;
                y[base + 2 * spatial] = v2 / norm;
                y[base + 3 * spatial] = v3 / norm;
            }
        } else {
            for (int f = 0; f < F; f++) {
                int index = (b * F + f) * spatial + pos;
                y[index] = x[index] / norm;
            }
        }
    }
}

torch::Tensor rms_norm(torch::Tensor x, int num_features, float eps) {
    c10::cuda::CUDAGuard device_guard(x.device());
    
    // x shape: (B, F, H, W)
    int B = x.size(0);
    int F = num_features;  // x.size(1) should equal num_features.
    int H = x.size(2);
    int W = x.size(3);
    int spatial = H * W;
    int total = B * spatial;
    
    auto y = torch::empty_like(x);
    
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    
    rms_norm_kernel<<<blocks, threads, 0, c10::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(),
        y.data_ptr<float>(),
        B, F, H, W, eps);
        
    return y;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("rms_norm", &rms_norm, "Fused RMS normalization kernel with loop unrolling");
}