feat(ml-alpha): LayerNorm backward + per-row param-grad reducer kernels

crates/ml-alpha/cuda/layer_norm.cu

@@ -62,3 +62,99 @@ extern "C" __global__ void layer_norm_fwd(
         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];
+    }
+}