feat(ml-alpha): projection kernel (128->8 + layer-norm)

Single-block 8-thread kernel; thread j computes its own 128-dim dot
product, then thread 0 computes block-wide mean/var, then each thread
applies the per-output affine layer-norm. No atomicAdd; reductions are
single-thread (8 elements — negligible cost).

Tests (5/5 on sm_86) assert:
  - layer-norm zero-mean output under identity gain
  - layer-norm unit-variance output under identity gain
  - ln_bias shifts mean uniformly
  - ln_gain scales variance (var = gain^2)
  - finite output under zero input (variance clamp activates)

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

crates/ml-alpha/cuda/projection.cu

@@ -1,2 +1,51 @@
-// placeholder — real implementation in the task adding this kernel
-extern "C" __global__ void projection_stub() {}
+// projection.cu — h_128 -> z_8 with affine + layer-norm.
+//
+// pre[j]   = sum_i w[j, i] * h[i] + b[j]                  for j in 0..8
+// mean     = (1/8) * sum_j pre[j]
+// var      = (1/8) * sum_j (pre[j] - mean)^2
+// out[j]   = ln_gain[j] * (pre[j] - mean) * rsqrt(var + 1e-6) + ln_bias[j]
+//
+// Single block, 8 threads (one per output unit). Mean and variance
+// computed by a single thread (block-coordinated) to keep the kernel
+// simple — the cost is negligible for 8 outputs. No atomicAdd anywhere.
+
+extern "C" __global__ void projection_kernel(
+    const float* __restrict__ w,        // [8, 128]
+    const float* __restrict__ b,        // [8]
+    const float* __restrict__ ln_gain,  // [8]
+    const float* __restrict__ ln_bias,  // [8]
+    const float* __restrict__ h,        // [128]
+    float* __restrict__ out             // [8]
+) {
+    __shared__ float pre[8];
+    __shared__ float mean_var[2]; // [mean, var]
+
+    int j = threadIdx.x;
+    if (j < 8) {
+        float v = b[j];
+        for (int i = 0; i < 128; ++i) {
+            v += w[j * 128 + i] * h[i];
+        }
+        pre[j] = v;
+    }
+    __syncthreads();
+
+    if (j == 0) {
+        float s = 0.0f, ss = 0.0f;
+        for (int k = 0; k < 8; ++k) {
+            s += pre[k];
+            ss += pre[k] * pre[k];
+        }
+        const float mean = s * 0.125f; // /8
+        const float var = ss * 0.125f - mean * mean;
+        mean_var[0] = mean;
+        mean_var[1] = fmaxf(var, 1e-6f);
+    }
+    __syncthreads();
+
+    if (j < 8) {
+        const float mean = mean_var[0];
+        const float invstd = rsqrtf(mean_var[1]);
+        out[j] = ln_gain[j] * (pre[j] - mean) * invstd + ln_bias[j];
+    }
+}

crates/ml-alpha/src/heads.rs

@@ -13,6 +13,7 @@ use ml_core::device::MlDevice;
 use crate::pinned_mem::MappedF32Buffer;
 
 const HEADS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_horizon_heads.cubin"));
+const PROJ_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/projection.cubin"));
 
 pub const N_HORIZONS: usize = 5;
 pub const HORIZONS: [usize; N_HORIZONS] = [30, 100, 300, 1000, 6000];
@@ -38,6 +39,50 @@ pub struct MultiHorizonHeads;
 
 pub struct Projection;
 
+pub fn projection_gpu(
+    dev: &MlDevice,
+    w: &ProjectionWeights,
+    h: &[f32],
+) -> Result<[f32; PROJ_DIM]> {
+    assert_eq!(h.len(), HIDDEN_DIM);
+    assert_eq!(w.w.len(), PROJ_DIM * HIDDEN_DIM);
+    assert_eq!(w.b.len(), PROJ_DIM);
+    assert_eq!(w.ln_gain.len(), PROJ_DIM);
+    assert_eq!(w.ln_bias.len(), PROJ_DIM);
+
+    let stream: &Arc<CudaStream> = dev.cuda_stream().context("proj stream")?;
+    let ctx = dev.cuda_context().context("proj ctx")?;
+    let module = ctx.load_cubin(PROJ_CUBIN.to_vec()).context("load proj cubin")?;
+    let func = module.load_function("projection_kernel").context("load function")?;
+
+    let w_d = upload(stream, &w.w)?;
+    let b_d = upload(stream, &w.b)?;
+    let g_d = upload(stream, &w.ln_gain)?;
+    let n_d = upload(stream, &w.ln_bias)?;
+    let h_d = upload(stream, h)?;
+    let mut out_d = stream.alloc_zeros::<f32>(PROJ_DIM).context("proj out alloc")?;
+
+    let cfg = LaunchConfig {
+        grid_dim: (1, 1, 1),
+        block_dim: (PROJ_DIM as u32, 1, 1),
+        shared_mem_bytes: 0,
+    };
+    let mut launch = stream.launch_builder(&func);
+    launch
+        .arg(&w_d)
+        .arg(&b_d)
+        .arg(&g_d)
+        .arg(&n_d)
+        .arg(&h_d)
+        .arg(&mut out_d);
+    unsafe { launch.launch(cfg).context("proj launch")?; }
+
+    let v = download(stream, &out_d)?;
+    let mut out = [0f32; PROJ_DIM];
+    out.copy_from_slice(&v);
+    Ok(out)
+}
+
 pub fn multi_horizon_heads_gpu(
     dev: &MlDevice,
     w: &HeadsWeights,

crates/ml-alpha/tests/projection_bit_equiv.rs

@@ -0,0 +1,111 @@
+//! projection (128 -> 8 + layer-norm) GPU invariants.
+
+use approx::assert_relative_eq;
+use ml_alpha::heads::{projection_gpu, ProjectionWeights, HIDDEN_DIM, PROJ_DIM};
+use ml_core::device::MlDevice;
+use rand::{Rng, SeedableRng};
+use rand_chacha::ChaCha8Rng;
+
+fn test_device() -> MlDevice {
+    MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests")
+}
+
+#[test]
+fn layer_norm_output_has_near_zero_mean_under_identity_gain() {
+    let dev = test_device();
+    let mut r = ChaCha8Rng::seed_from_u64(0x90F1);
+    let w = ProjectionWeights {
+        w: (0..PROJ_DIM * HIDDEN_DIM).map(|_| r.gen_range(-0.1..0.1)).collect(),
+        b: vec![0.0; PROJ_DIM],
+        ln_gain: vec![1.0; PROJ_DIM],
+        ln_bias: vec![0.0; PROJ_DIM],
+    };
+    let h: Vec<f32> = (0..HIDDEN_DIM).map(|_| r.gen_range(-1.0..1.0)).collect();
+    let out = projection_gpu(&dev, &w, &h).expect("gpu");
+    let mean: f32 = out.iter().sum::<f32>() / PROJ_DIM as f32;
+    assert!(
+        mean.abs() < 1e-4,
+        "layer-norm should produce zero-mean output, got mean={mean}"
+    );
+}
+
+#[test]
+fn layer_norm_output_has_near_unit_variance_under_identity_gain() {
+    let dev = test_device();
+    let mut r = ChaCha8Rng::seed_from_u64(0x91F2);
+    let w = ProjectionWeights {
+        w: (0..PROJ_DIM * HIDDEN_DIM).map(|_| r.gen_range(-0.5..0.5)).collect(),
+        b: vec![0.0; PROJ_DIM],
+        ln_gain: vec![1.0; PROJ_DIM],
+        ln_bias: vec![0.0; PROJ_DIM],
+    };
+    let h: Vec<f32> = (0..HIDDEN_DIM).map(|_| r.gen_range(-1.0..1.0)).collect();
+    let out = projection_gpu(&dev, &w, &h).expect("gpu");
+    let mean: f32 = out.iter().sum::<f32>() / PROJ_DIM as f32;
+    let var: f32 = out.iter().map(|v| (v - mean).powi(2)).sum::<f32>() / PROJ_DIM as f32;
+    // Variance should be close to 1 (with biased-variance estimator divided
+    // by N — the kernel uses the same biased estimator).
+    assert!(
+        (var - 1.0_f32).abs() < 1e-3,
+        "layer-norm output variance should be ≈1, got {var}"
+    );
+}
+
+#[test]
+fn ln_bias_shifts_output_uniformly() {
+    let dev = test_device();
+    let mut r = ChaCha8Rng::seed_from_u64(0x92F3);
+    let bias_offset = 3.0_f32;
+    let w_with = ProjectionWeights {
+        w: (0..PROJ_DIM * HIDDEN_DIM).map(|_| r.gen_range(-0.1..0.1)).collect(),
+        b: vec![0.0; PROJ_DIM],
+        ln_gain: vec![1.0; PROJ_DIM],
+        ln_bias: vec![bias_offset; PROJ_DIM],
+    };
+    let h: Vec<f32> = (0..HIDDEN_DIM).map(|_| r.gen_range(-1.0..1.0)).collect();
+    let out = projection_gpu(&dev, &w_with, &h).expect("gpu");
+    let mean: f32 = out.iter().sum::<f32>() / PROJ_DIM as f32;
+    assert_relative_eq!(mean, bias_offset, epsilon = 1e-4);
+}
+
+#[test]
+fn ln_gain_scales_centered_output() {
+    let dev = test_device();
+    let mut r = ChaCha8Rng::seed_from_u64(0x93F4);
+    let gain = 2.5_f32;
+    let w = ProjectionWeights {
+        w: (0..PROJ_DIM * HIDDEN_DIM).map(|_| r.gen_range(-0.3..0.3)).collect(),
+        b: vec![0.0; PROJ_DIM],
+        ln_gain: vec![gain; PROJ_DIM],
+        ln_bias: vec![0.0; PROJ_DIM],
+    };
+    let h: Vec<f32> = (0..HIDDEN_DIM).map(|_| r.gen_range(-1.0..1.0)).collect();
+    let out = projection_gpu(&dev, &w, &h).expect("gpu");
+    let mean: f32 = out.iter().sum::<f32>() / PROJ_DIM as f32;
+    let var: f32 = out.iter().map(|v| (v - mean).powi(2)).sum::<f32>() / PROJ_DIM as f32;
+    // After layer-norm var ≈ 1; after gain=2.5 var ≈ 6.25.
+    assert!(
+        (var - gain * gain).abs() < 1e-2,
+        "variance should be gain^2 ≈ 6.25, got {var}"
+    );
+}
+
+#[test]
+fn output_finite_under_zero_input() {
+    // Zero hidden -> pre = b -> all equal -> var = 0 -> 1e-6 clamp kicks in.
+    let dev = test_device();
+    let w = ProjectionWeights {
+        w: vec![1.0; PROJ_DIM * HIDDEN_DIM],
+        b: vec![3.0; PROJ_DIM],
+        ln_gain: vec![1.0; PROJ_DIM],
+        ln_bias: vec![0.0; PROJ_DIM],
+    };
+    let h = vec![0.0; HIDDEN_DIM];
+    let out = projection_gpu(&dev, &w, &h).expect("gpu");
+    // All pre values are equal (3.0), so after layer-norm everything is 0
+    // (zero variance clamped to 1e-6, then (pre - mean) = 0 → 0).
+    for v in &out {
+        assert!(v.is_finite(), "output must be finite even at zero variance, got {}", v);
+        assert!(v.abs() < 1e-3, "output should be ≈0 when all pre values are equal, got {}", v);
+    }
+}