foxhunt/crates/ml-alpha/cuda/layer_norm.cu

// layer_norm.cu — per-row LayerNorm over a [N_rows, HIDDEN] tensor.
//
// y[r, j] = gain[j] * (x[r, j] - mean_r) / sqrt(var_r + eps) + bias[j]
//
// Forward saves `mean_r` and `inv_std_r = 1 / sqrt(var_r + eps)` to a
// scratch buffer of shape [N_rows, 2] so the backward kernel doesn't
// re-derive them. One block per row; block tree-reduce for mean + var.
// Per `feedback_no_atomicadd.md`: shared-mem tree-reduce only, no atomics.

#define LN_HIDDEN     128
#define LN_BLOCK      128
#define LN_EPS        1e-5f

extern "C" __global__ void layer_norm_fwd(
    const float* __restrict__ x,        // [N_rows, HIDDEN]
    const float* __restrict__ gain,     // [HIDDEN]
    const float* __restrict__ bias,     // [HIDDEN]
    int n_rows,
    float* __restrict__ y,              // [N_rows, HIDDEN]
    float* __restrict__ stats           // [N_rows, 2] — col 0=mean, col 1=inv_std
) {
    int r = blockIdx.x;
    int tid = threadIdx.x;
    if (r >= n_rows) return;

    __shared__ float ssum[LN_BLOCK];
    __shared__ float ssq[LN_BLOCK];
    __shared__ float s_mean;
    __shared__ float s_inv_std;

    const float* x_row = x + (long long)r * LN_HIDDEN;
    float my_sum = 0.0f, my_sq = 0.0f;
    if (tid < LN_HIDDEN) {
        float v = x_row[tid];
        my_sum = v;
        my_sq = v * v;
    }
    ssum[tid] = my_sum;
    ssq[tid] = my_sq;
    __syncthreads();

    for (int s = LN_BLOCK / 2; s > 0; s >>= 1) {
        if (tid < s) {
            ssum[tid] += ssum[tid + s];
            ssq[tid]  += ssq[tid + s];
        }
        __syncthreads();
    }
    if (tid == 0) {
        const float n = (float)LN_HIDDEN;
        const float mean = ssum[0] / n;
        const float var  = ssq[0] / n - mean * mean;
        s_mean    = mean;
        s_inv_std = rsqrtf(var + LN_EPS);
        stats[(long long)r * 2 + 0] = mean;
        stats[(long long)r * 2 + 1] = s_inv_std;
    }
    __syncthreads();

    if (tid < LN_HIDDEN) {
        float normalised = (x_row[tid] - s_mean) * s_inv_std;
        y[(long long)r * LN_HIDDEN + tid] = gain[tid] * normalised + bias[tid];
    }
}

// Backward: given grad_y[N, HIDDEN] and saved mean+inv_std per row,
// produce grad_x[N, HIDDEN] and per-row grad_gain / grad_bias
// scratch tensors. A separate reducer (`layer_norm_reduce_param_grads`)
// sums the per-row scratch into the canonical [HIDDEN] grad_gain /
// grad_bias buffers.
//
// Per-row writes from this kernel are race-free because each block
// owns one row; the cross-row sum needs the reducer (no atomicAdd
// per `feedback_no_atomicadd.md`).
//
//   grad_x[r, j] = inv_std * (gain[j] * grad_y[r, j]
//                              - mean(gain[j'] * grad_y[r, j'])
//                              - normalised[r, j] * mean(gain[j'] * grad_y[r, j'] * normalised[r, j']))
//   grad_gain_per_row[r, j] = grad_y[r, j] * normalised[r, j]
//   grad_bias_per_row[r, j] = grad_y[r, j]

extern "C" __global__ void layer_norm_bwd(
    const float* __restrict__ x,                // [N_rows, HIDDEN]
    const float* __restrict__ gain,             // [HIDDEN]
    const float* __restrict__ stats,            // [N_rows, 2] from fwd
    const float* __restrict__ grad_y,           // [N_rows, HIDDEN]
    int n_rows,
    float* __restrict__ grad_x,                 // [N_rows, HIDDEN] (overwrite)
    float* __restrict__ grad_gain_per_row,      // [N_rows, HIDDEN] (overwrite)
    float* __restrict__ grad_bias_per_row       // [N_rows, HIDDEN] (overwrite)
) {
    int r = blockIdx.x;
    int tid = threadIdx.x;
    if (r >= n_rows) return;

    __shared__ float s_a[LN_BLOCK];  // sum of gain[j] * grad_y[r, j]
    __shared__ float s_b[LN_BLOCK];  // sum of gain[j] * grad_y[r, j] * normalised[r, j]
    __shared__ float s_mean_a;
    __shared__ float s_mean_b;

    const float mean    = stats[(long long)r * 2 + 0];
    const float inv_std = stats[(long long)r * 2 + 1];
    const float gy = (tid < LN_HIDDEN) ? grad_y[(long long)r * LN_HIDDEN + tid] : 0.0f;
    const float xv = (tid < LN_HIDDEN) ? x[(long long)r * LN_HIDDEN + tid] : 0.0f;
    const float gn = (tid < LN_HIDDEN) ? gain[tid] : 0.0f;
    const float normalised = (xv - mean) * inv_std;

    s_a[tid] = gn * gy;
    s_b[tid] = gn * gy * normalised;
    __syncthreads();
    for (int s = LN_BLOCK / 2; s > 0; s >>= 1) {
        if (tid < s) {
            s_a[tid] += s_a[tid + s];
            s_b[tid] += s_b[tid + s];
        }
        __syncthreads();
    }
    if (tid == 0) {
        const float n = (float)LN_HIDDEN;
        s_mean_a = s_a[0] / n;
        s_mean_b = s_b[0] / n;
    }
    __syncthreads();

    if (tid < LN_HIDDEN) {
        grad_x[(long long)r * LN_HIDDEN + tid] =
            inv_std * (gn * gy - s_mean_a - normalised * s_mean_b);
        grad_gain_per_row[(long long)r * LN_HIDDEN + tid] = gy * normalised;
        grad_bias_per_row[(long long)r * LN_HIDDEN + tid] = gy;
    }
}

// Reducer: sum [N_rows, HIDDEN] per-row tensor along axis 0 into
// [HIDDEN]. One block per output column; block tree-reduce across rows
// (no atomicAdd). Stride-loop within each block so we handle any
// N_rows up to ~64K (limited only by block-thread count of LN_BLOCK).
extern "C" __global__ void layer_norm_reduce_param_grads(
    const float* __restrict__ per_row,   // [N_rows, HIDDEN]
    int n_rows,
    float* __restrict__ out              // [HIDDEN] (overwrite, NOT +=)
) {
    int j = blockIdx.x;
    int tid = threadIdx.x;
    if (j >= LN_HIDDEN) return;

    __shared__ float ssum[LN_BLOCK];
    float my_sum = 0.0f;
    for (int r = tid; r < n_rows; r += LN_BLOCK) {
        my_sum += per_row[(long long)r * LN_HIDDEN + j];
    }
    ssum[tid] = my_sum;
    __syncthreads();
    for (int s = LN_BLOCK / 2; s > 0; s >>= 1) {
        if (tid < s) ssum[tid] += ssum[tid + s];
        __syncthreads();
    }
    if (tid == 0) {
        out[j] = ssum[0];
    }
}