diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index e02434d57..416198024 100644 --- a/crates/ml-alpha/build.rs +++ b/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"); diff --git a/crates/ml-alpha/cuda/layer_norm.cu b/crates/ml-alpha/cuda/layer_norm.cu new file mode 100644 index 000000000..84200e00a --- /dev/null +++ b/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]; + } +}