diff --git a/crates/ml-alpha/cuda/projection.cu b/crates/ml-alpha/cuda/projection.cu index 263b2ad25..e6761df69 100644 --- a/crates/ml-alpha/cuda/projection.cu +++ b/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]; + } +} diff --git a/crates/ml-alpha/src/heads.rs b/crates/ml-alpha/src/heads.rs index 25f029059..8089f112c 100644 --- a/crates/ml-alpha/src/heads.rs +++ b/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 = 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::(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, diff --git a/crates/ml-alpha/tests/projection_bit_equiv.rs b/crates/ml-alpha/tests/projection_bit_equiv.rs new file mode 100644 index 000000000..9fee7f725 --- /dev/null +++ b/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 = (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::() / 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 = (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::() / PROJ_DIM as f32; + let var: f32 = out.iter().map(|v| (v - mean).powi(2)).sum::() / 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 = (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::() / 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 = (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::() / PROJ_DIM as f32; + let var: f32 = out.iter().map(|v| (v - mean).powi(2)).sum::() / 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); + } +}