feat(ml-alpha): LayerNorm forward kernel for trunk-pre-CfC normalisation

crates/ml-alpha/build.rs

@@ -17,13 +17,13 @@ const KERNELS: &[&str] = &[
     "adamw_step",
     "grad_norm",
     "horizon_lambda",                // ISV-driven per-horizon gradient scaler (EMA + lambda)
+    "layer_norm",                    // Phase 1: trunk pre-CfC normalisation
 ];
 
-// Cache bust v5 (2026-05-17): horizon_lambda.cu rewrote the lambda
-// formula from ratio-of-mean to z-score-normalised. Old cubins compute
-// a numerically different lambda; running them against the new Rust
-// loop would silently apply the wrong gradient scaler. Force nvcc
-// recompile on the cluster's /cargo-target PVC.
+// Cache bust v6 (2026-05-17): Phase 1 model capacity — LayerNorm
+// kernel added (cuda/layer_norm.cu) + multi_horizon_heads.cu will get
+// 2-layer MLP heads. Old cubins don't have the new symbols. Force
+// fresh nvcc compile on the cluster's /cargo-target PVC.
 
 fn main() {
     println!("cargo:rerun-if-changed=build.rs");

crates/ml-alpha/cuda/layer_norm.cu

@@ -0,0 +1,64 @@
+// 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];
+    }
+}