From d4e46aba94cdbed54f387538e4cc1611ae221fda Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 16 May 2026 21:51:48 +0200 Subject: [PATCH] feat(ml-alpha): multi_horizon_heads kernel (128->5 sigmoid) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-horizon P(up) at h ∈ {30, 100, 300, 1000, 6000} snapshots forward. Single-block 5-thread kernel; each thread is its own 128-dim dot product + sigmoid. No atomicAdd. Tests (5/5 pass on sm_86) assert invariants only: - sigmoid output ∈ [0, 1] for all heads - zero weights + zero bias → 0.5 exactly - bias = +20 → saturates near 1 - bias = -20 → saturates near 0 - per-head independence (mixed-bias configuration) Addendum updated to explicitly state no-CPU-mirror discipline per feedback_no_cpu_test_fallbacks.md. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/cuda/multi_horizon_heads.cu | 23 +++- crates/ml-alpha/src/heads.rs | 104 ++++++++++++++++++ crates/ml-alpha/src/lib.rs | 1 + crates/ml-alpha/tests/heads_bit_equiv.rs | 87 +++++++++++++++ ...026-05-16-ml-alpha-phase-a-api-addendum.md | 27 ++++- 5 files changed, 238 insertions(+), 4 deletions(-) create mode 100644 crates/ml-alpha/src/heads.rs create mode 100644 crates/ml-alpha/tests/heads_bit_equiv.rs diff --git a/crates/ml-alpha/cuda/multi_horizon_heads.cu b/crates/ml-alpha/cuda/multi_horizon_heads.cu index 9b543b682..1c5e9e7dd 100644 --- a/crates/ml-alpha/cuda/multi_horizon_heads.cu +++ b/crates/ml-alpha/cuda/multi_horizon_heads.cu @@ -1,2 +1,21 @@ -// placeholder — real implementation in the task adding this kernel -extern "C" __global__ void multi_horizon_heads_stub() {} +// multi_horizon_heads.cu +// +// 128-dim hidden -> 5 sigmoid logits, one per horizon +// {30, 100, 300, 1000, 6000} snapshots forward. Single-block kernel, +// 5 threads (one per head). Each thread computes its own dot product +// against the 128-dim hidden vector + bias, then sigmoid. + +extern "C" __global__ void multi_horizon_heads( + const float* __restrict__ w, // [5, 128] + const float* __restrict__ b, // [5] + const float* __restrict__ h, // [128] + float* __restrict__ probs // [5] +) { + int k = threadIdx.x; + if (k >= 5) return; + float z = b[k]; + for (int i = 0; i < 128; ++i) { + z += w[k * 128 + i] * h[i]; + } + probs[k] = 1.0f / (1.0f + expf(-z)); +} diff --git a/crates/ml-alpha/src/heads.rs b/crates/ml-alpha/src/heads.rs new file mode 100644 index 000000000..25f029059 --- /dev/null +++ b/crates/ml-alpha/src/heads.rs @@ -0,0 +1,104 @@ +//! Multi-horizon sigmoid heads (128 → 5) and projection (128 → 8 + layer-norm). +//! +//! Per spec Section 2 + 2026-05-16 stacked amendment, the heads sit on +//! the CfC hidden output and emit 5 sigmoid probabilities of "up" at +//! horizons {30, 100, 300, 1000, 6000} snapshots forward. + +use std::sync::Arc; + +use anyhow::{Context, Result}; +use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg}; +use ml_core::device::MlDevice; + +use crate::pinned_mem::MappedF32Buffer; + +const HEADS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_horizon_heads.cubin")); + +pub const N_HORIZONS: usize = 5; +pub const HORIZONS: [usize; N_HORIZONS] = [30, 100, 300, 1000, 6000]; +pub const HIDDEN_DIM: usize = 128; +pub const PROJ_DIM: usize = 8; + +#[derive(Clone, Debug)] +pub struct HeadsWeights { + pub w: Vec, // [5, 128] row-major + pub b: Vec, // [5] +} + +#[derive(Clone, Debug)] +pub struct ProjectionWeights { + pub w: Vec, // [8, 128] + pub b: Vec, // [8] + pub ln_gain: Vec, // [8] + pub ln_bias: Vec, // [8] +} + +/// Placeholder marker — concrete trunk integration in Task 9. +pub struct MultiHorizonHeads; + +pub struct Projection; + +pub fn multi_horizon_heads_gpu( + dev: &MlDevice, + w: &HeadsWeights, + h: &[f32], +) -> Result<[f32; N_HORIZONS]> { + assert_eq!(h.len(), HIDDEN_DIM); + assert_eq!(w.w.len(), N_HORIZONS * HIDDEN_DIM); + assert_eq!(w.b.len(), N_HORIZONS); + + let stream: &Arc = dev.cuda_stream().context("heads stream")?; + let ctx = dev.cuda_context().context("heads ctx")?; + let module = ctx.load_cubin(HEADS_CUBIN.to_vec()).context("load heads cubin")?; + let func = module.load_function("multi_horizon_heads").context("load function")?; + + let w_d = upload(stream, &w.w)?; + let b_d = upload(stream, &w.b)?; + let h_d = upload(stream, h)?; + let mut probs_d = stream.alloc_zeros::(N_HORIZONS).context("probs alloc")?; + + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (N_HORIZONS as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = stream.launch_builder(&func); + launch.arg(&w_d).arg(&b_d).arg(&h_d).arg(&mut probs_d); + unsafe { launch.launch(cfg).context("heads launch")?; } + + let v = download(stream, &probs_d)?; + let mut out = [0f32; N_HORIZONS]; + out.copy_from_slice(&v); + Ok(out) +} + +fn upload(stream: &Arc, host: &[f32]) -> Result> { + let n = host.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("heads upload staging: {e}"))?; + staging.write_from_slice(host); + let mut dst = stream.alloc_zeros::(n).context("heads upload alloc")?; + if n > 0 { + let nbytes = n * std::mem::size_of::(); + unsafe { + let (dst_ptr, _g) = dst.device_ptr_mut(stream); + cudarc::driver::result::memcpy_dtod_async(dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream()) + .context("heads upload DtoD")?; + } + } + Ok(dst) +} + +fn download(stream: &Arc, src: &CudaSlice) -> Result> { + let n = src.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("heads download staging: {e}"))?; + let nbytes = n * std::mem::size_of::(); + unsafe { + let (src_ptr, _g) = src.device_ptr(stream); + cudarc::driver::result::memcpy_dtod_async(staging.dev_ptr, src_ptr, nbytes, stream.cu_stream()) + .context("heads download DtoD")?; + } + stream.synchronize().context("heads download sync")?; + Ok(staging.read_all()) +} diff --git a/crates/ml-alpha/src/lib.rs b/crates/ml-alpha/src/lib.rs index 3b53403df..17007b516 100644 --- a/crates/ml-alpha/src/lib.rs +++ b/crates/ml-alpha/src/lib.rs @@ -26,6 +26,7 @@ // Task 12: pub mod data; // Task 16+: pub mod gate; pub mod cfc; +pub mod heads; pub mod isv; pub mod pinned; pub mod pinned_mem; diff --git a/crates/ml-alpha/tests/heads_bit_equiv.rs b/crates/ml-alpha/tests/heads_bit_equiv.rs new file mode 100644 index 000000000..98e96376b --- /dev/null +++ b/crates/ml-alpha/tests/heads_bit_equiv.rs @@ -0,0 +1,87 @@ +//! multi_horizon_heads GPU invariants. +//! +//! Per `feedback_no_cpu_test_fallbacks.md`: validation via analytically- +//! known synthetic inputs + property assertions. + +use approx::assert_relative_eq; +use ml_alpha::heads::{multi_horizon_heads_gpu, HeadsWeights, HIDDEN_DIM, N_HORIZONS}; +use ml_core::device::MlDevice; + +fn test_device() -> MlDevice { + MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests") +} + +#[test] +fn probs_are_in_unit_interval() { + let dev = test_device(); + let w = HeadsWeights { + w: vec![0.1; N_HORIZONS * HIDDEN_DIM], + b: vec![0.0; N_HORIZONS], + }; + let h: Vec = (0..HIDDEN_DIM).map(|i| 0.01 * i as f32).collect(); + let probs = multi_horizon_heads_gpu(&dev, &w, &h).expect("gpu"); + for k in 0..N_HORIZONS { + assert!((0.0..=1.0).contains(&probs[k]), "head {k} prob {} not in [0,1]", probs[k]); + } +} + +#[test] +fn zero_weights_and_bias_give_half() { + // logit = 0 -> sigmoid(0) = 0.5 + let dev = test_device(); + let w = HeadsWeights { + w: vec![0.0; N_HORIZONS * HIDDEN_DIM], + b: vec![0.0; N_HORIZONS], + }; + let h: Vec = (0..HIDDEN_DIM).map(|i| i as f32).collect(); + let probs = multi_horizon_heads_gpu(&dev, &w, &h).expect("gpu"); + for k in 0..N_HORIZONS { + assert_relative_eq!(probs[k], 0.5_f32, epsilon = 1e-6); + } +} + +#[test] +fn large_positive_bias_saturates_high() { + let dev = test_device(); + let w = HeadsWeights { + w: vec![0.0; N_HORIZONS * HIDDEN_DIM], + b: vec![20.0; N_HORIZONS], + }; + let h = vec![0.0; HIDDEN_DIM]; + let probs = multi_horizon_heads_gpu(&dev, &w, &h).expect("gpu"); + for k in 0..N_HORIZONS { + assert!(probs[k] > 0.999, "head {k} should saturate near 1 with bias=20, got {}", probs[k]); + } +} + +#[test] +fn large_negative_bias_saturates_low() { + let dev = test_device(); + let w = HeadsWeights { + w: vec![0.0; N_HORIZONS * HIDDEN_DIM], + b: vec![-20.0; N_HORIZONS], + }; + let h = vec![0.0; HIDDEN_DIM]; + let probs = multi_horizon_heads_gpu(&dev, &w, &h).expect("gpu"); + for k in 0..N_HORIZONS { + assert!(probs[k] < 0.001, "head {k} should saturate near 0 with bias=-20, got {}", probs[k]); + } +} + +#[test] +fn per_head_independence() { + // Set bias[0] = 5, bias[4] = -5, others = 0 with zero weights. + // sigmoid(5) ≈ 0.993, sigmoid(0) = 0.5, sigmoid(-5) ≈ 0.0067. + let dev = test_device(); + let w = HeadsWeights { + w: vec![0.0; N_HORIZONS * HIDDEN_DIM], + b: vec![5.0, 0.0, 0.0, 0.0, -5.0], + }; + let h = vec![0.0; HIDDEN_DIM]; + let probs = multi_horizon_heads_gpu(&dev, &w, &h).expect("gpu"); + assert!(probs[0] > 0.99 && probs[0] < 1.0); + assert_relative_eq!(probs[1], 0.5, epsilon = 1e-6); + assert_relative_eq!(probs[2], 0.5, epsilon = 1e-6); + assert_relative_eq!(probs[3], 0.5, epsilon = 1e-6); + assert!(probs[4] > 0.0 && probs[4] < 0.01); +} diff --git a/docs/superpowers/plans/2026-05-16-ml-alpha-phase-a-api-addendum.md b/docs/superpowers/plans/2026-05-16-ml-alpha-phase-a-api-addendum.md index 2c74b1915..25fa4d865 100644 --- a/docs/superpowers/plans/2026-05-16-ml-alpha-phase-a-api-addendum.md +++ b/docs/superpowers/plans/2026-05-16-ml-alpha-phase-a-api-addendum.md @@ -3,8 +3,31 @@ > Companion to `2026-05-16-ml-alpha-phase-a.md`. When the plan's binding-code > blocks reference `dev.htod_copy`, `dev.alloc_zeros`, `dev.load_cubin_function`, > `func.launch_async`, or `dev.default_stream`, **use the patterns in this -> addendum instead.** The plan's kernel `.cu` source, CPU oracles, tests, and -> assertion thresholds are unchanged. +> addendum instead.** The plan's kernel `.cu` source and assertion thresholds +> are unchanged. + +## Validation discipline override: no CPU mirrors + +Plan 1 incorrectly listed a "CPU oracle + bit-equiv test" step for every kernel +task. Per `feedback_no_cpu_test_fallbacks.md` — kernels must be validated via +**invariants on the GPU output**, not CPU mirrors. The user re-confirmed this +mid-execution. + +**Canonical validation pattern for each kernel:** + +1. Build analytically-known synthetic inputs (zeros, identities, magic constants + with known kernel outputs). +2. Run the GPU kernel. +3. Assert algebraic invariants on the output: ranges (e.g. sigmoid outputs + ∈ [0, 1]), fixed-points (dt=0 → identity), monotonicity (input ↑ → output ↑), + conservation (sum of softmax probs ≈ 1), bounded growth (output ≤ deterministic + bound from input). +4. Where forward + backward are paired (BCE, PPO loss), validate the gradient + via **finite-difference on the GPU forward** at a small set of perturbed + inputs — NOT against a CPU autograd. + +Whenever Plan 1's task body says "CPU oracle" or "bit-equiv test", read it as +"GPU-output invariant test." This addendum exists because Plan 1 was written against an older cudarc device-centric API. The vendored cudarc 0.19 (`vendor/cudarc/`) moved