CUDA Group Normalization Optimization Guide

// Three separate kernels
// 1. Compute group means
// 2. Compute group variances
// 3. Normalize with gamma/beta
void group_norm_naive(float* x, float* y, int N, int C, int HW, int G) {
    compute_group_means<<<N*G, 256>>>(x, means, N, C, HW, G);
    compute_group_vars<<<N*G, 256>>>(x, means, vars, N, C, HW, G);
    normalize_groups<<<N*C, 256>>>(x, y, means, vars, gamma, beta, N, C, HW, G);
}

__global__ void group_norm_opt(float* __restrict__ x, float* __restrict__ y,
                                  float* __restrict__ gamma, float* __restrict__ beta,
                                  int C, int HW, int G, float eps) {
    extern __shared__ float smem[];
    int n = blockIdx.x;
    int g = blockIdx.y;
    int cpg = C / G;
    int group_size = cpg * HW;

    float sum = 0, sum_sq = 0;
    for (int i = threadIdx.x; i < group_size; i += blockDim.x) {
        int c = g * cpg + i / HW;
        float val = x[n * C * HW + c * HW + i % HW];
        sum += val;
        sum_sq += val * val;
    }

    // Block reduction
    sum = blockReduceSum(sum);
    sum_sq = blockReduceSum(sum_sq);

    __shared__ float s_mean, s_var;
    if (threadIdx.x == 0) {
        s_mean = sum / group_size;
        s_var = sum_sq / group_size - s_mean * s_mean;
    }
    __syncthreads();

    float inv_std = rsqrtf(s_var + eps);
    for (int i = threadIdx.x; i < group_size; i += blockDim.x) {
        int c = g * cpg + i / HW;
        int idx = n * C * HW + c * HW + i % HW;
        y[idx] = (x[idx] - s_mean) * inv_std * gamma[c] + beta[c];
    }
}