From 8c335caef70ef6cf72537799f4b2ea3fbce4112b Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 18 May 2026 10:36:42 +0200 Subject: [PATCH] feat(ml-alpha): per-horizon residual head kernel + numgrad parity (C22) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Companion kernel to C21's per_horizon_attention_pool. Computes a per-horizon scalar residual from each horizon's context vector: residual[b, h] = Σ_d w_res[h, d] * context_h[b, h, d] + bias_res[h] Designed to be added (behind a learnable α-gate) to the existing multi_horizon_heads logit output — keeps the existing GRN head kernel completely unchanged. The per-horizon attention pool's contribution flows through this lightweight projection without weight-shape changes elsewhere or checkpoint-V2-bumping. Path A integration sketch (deferred to follow-up commit C23): alpha_logit_per_horizon = existing_head(h_K)[h] # from current path + tanh(α[h]) * residual_kernel(context_h)[h] where α[h] is a learnable 5-vector init'd to 0 (no effect at start). Training discovers per-horizon whether the residual contributes. This is a strict superset of the existing path — α=0 → bit-identical to today. Backward kernel produces: d_w_res — per-block scratch [B, N_HORIZONS, HIDDEN_DIM] for host reduce_axis0 → shared [N_HORIZONS, HIDDEN_DIM] d_bias_res — per-block scratch [B, N_HORIZONS], same reduction d_context_h — per-batch indexed; += chained with attention bwd Single-writer discipline preserved (no atomicAdd per feedback_no_atomicadd.md); horizon loop inside the per-batch block. Numgrad parity test: - B=3, N_HORIZONS=5, HIDDEN_DIM=128 fixture. - Loss = Σ residual_out (so d_residual = 1). - Probes 8 random w_res indices, all 5 bias_res entries, 8 random context_h indices via central-difference at ±eps=1e-2. - All within 5e-2 rel-tol or 5e-3 abs-floor. - Passes on RTX 3050. Same scope discipline as C21: kernel + binding + numgrad first; trainer wiring + α-gate + smoke training + A/B sweep follow once both kernels are individually validated (now done). Closes the second kernel-correctness portion of #203. Remaining: C23: trainer wiring (capture attn_pool fwd into the graph; sum residual into existing head output with α-gate) C24: CheckpointV1 → V2 bump (add q_h, w_res, bias_res, alpha fields) C25: 1-epoch smoke (assert no NaN, loss decreases vs baseline) C26: 30-epoch × 3-fold A/B (#204) — decision gate per spec §0 Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/build.rs | 1 + .../cuda/per_horizon_residual_head.cu | 95 ++++++++++ crates/ml-alpha/src/lib.rs | 1 + .../ml-alpha/src/per_horizon_residual_head.rs | 113 ++++++++++++ .../per_horizon_residual_head_numgrad.rs | 163 ++++++++++++++++++ 5 files changed, 373 insertions(+) create mode 100644 crates/ml-alpha/cuda/per_horizon_residual_head.cu create mode 100644 crates/ml-alpha/src/per_horizon_residual_head.rs create mode 100644 crates/ml-alpha/tests/per_horizon_residual_head_numgrad.rs diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 858574199..848541e35 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -21,6 +21,7 @@ const KERNELS: &[&str] = &[ "variable_selection", // Phase 2D: TFT-style per-feature gating "attention_pool", // Phase 3: learned context summary at CfC k=0 "per_horizon_attention_pool", // C21: per-horizon variant; ships behind a config flag + "per_horizon_residual_head", // C22: per-horizon scalar residual on top of multi_horizon_heads logits "reduce_axis0", // Phase B: cross-batch param-grad reducer ]; diff --git a/crates/ml-alpha/cuda/per_horizon_residual_head.cu b/crates/ml-alpha/cuda/per_horizon_residual_head.cu new file mode 100644 index 000000000..dd2f0140c --- /dev/null +++ b/crates/ml-alpha/cuda/per_horizon_residual_head.cu @@ -0,0 +1,95 @@ +// per_horizon_residual_head.cu — per-horizon scalar residual from contexts (C22). +// +// Produces a per-horizon scalar residual that the trainer adds (behind +// a learnable α-gate) to the existing multi_horizon_heads logit output. +// Keeps the existing head kernel unchanged — the per-horizon attention +// pool from C21 contributes via this lightweight projection that the +// PerceptionTrainer can opt into without weight-shape changes elsewhere. +// +// Forward math (per sample b, per horizon h): +// residual[b, h] = Σ_d w_res[h, d] * context_h[b, h, d] + bias_res[h] +// +// One block per batch; thread tid handles dimension d. Reduction across +// d uses block tree-reduce (no atomicAdd per feedback_no_atomicadd.md). +// Inner loop over N_HORIZONS sequentially within the block (same pattern +// as per_horizon_attention_pool.cu). +// +// Saved for backward: nothing (residual is linear in its inputs — bwd +// can recompute from w_res / context_h). + +#define PHR_HIDDEN_DIM 128 +#define PHR_BLOCK 128 +#define PHR_N_HORIZONS 5 + +extern "C" __global__ void per_horizon_residual_head_fwd( + const float* __restrict__ context_h, // [B, N_HORIZONS, HIDDEN_DIM] + const float* __restrict__ w_res, // [N_HORIZONS, HIDDEN_DIM] + const float* __restrict__ bias_res, // [N_HORIZONS] + int n_batch, + float* __restrict__ residual_out // [B, N_HORIZONS] +) { + int b = blockIdx.x; + int tid = threadIdx.x; + if (b >= n_batch || tid >= PHR_BLOCK) return; + + __shared__ float s_red[PHR_BLOCK]; + + for (int h = 0; h < PHR_N_HORIZONS; ++h) { + // Per-thread partial: w_res[h, tid] * context_h[b, h, tid]. + const float v = (tid < PHR_HIDDEN_DIM) + ? w_res[h * PHR_HIDDEN_DIM + tid] + * context_h[(long long)b * PHR_N_HORIZONS * PHR_HIDDEN_DIM + + (long long)h * PHR_HIDDEN_DIM + tid] + : 0.0f; + s_red[tid] = v; + __syncthreads(); + for (int s = PHR_BLOCK / 2; s > 0; s >>= 1) { + if (tid < s) s_red[tid] += s_red[tid + s]; + __syncthreads(); + } + if (tid == 0) { + residual_out[(long long)b * PHR_N_HORIZONS + h] = s_red[0] + bias_res[h]; + } + __syncthreads(); + } +} + +// Backward: +// d_w_res[h, d] += Σ_b context_h[b, h, d] * grad_residual[b, h] +// d_bias_res[h] += Σ_b grad_residual[b, h] +// d_context_h[b,h,d] = w_res[h, d] * grad_residual[b, h] +// +// One block per batch; thread tid owns column d. Writes: +// d_context_h[b, h, d] — sole writer per (b, h, d). No race. +// d_w_res_scratch[b, h, d] — per-block scratch; host reduce_axis0 collapses +// across batches to recover the shared [N_HORIZONS, HIDDEN_DIM] grad. +// d_bias_res_scratch[b, h] — per-block scratch; same reduction pattern. +extern "C" __global__ void per_horizon_residual_head_bwd( + const float* __restrict__ context_h, // [B, N_HORIZONS, HIDDEN_DIM] + const float* __restrict__ w_res, // [N_HORIZONS, HIDDEN_DIM] + const float* __restrict__ grad_residual, // [B, N_HORIZONS] + int n_batch, + float* __restrict__ d_w_res_scratch, // [B, N_HORIZONS, HIDDEN_DIM] (+=) + float* __restrict__ d_bias_res_scratch, // [B, N_HORIZONS] (+=) + float* __restrict__ d_context_h // [B, N_HORIZONS, HIDDEN_DIM] (+= chained) +) { + int b = blockIdx.x; + int tid = threadIdx.x; + if (b >= n_batch || tid >= PHR_BLOCK) return; + + for (int h = 0; h < PHR_N_HORIZONS; ++h) { + const float g = grad_residual[(long long)b * PHR_N_HORIZONS + h]; + if (tid < PHR_HIDDEN_DIM) { + const long long ctx_idx = (long long)b * PHR_N_HORIZONS * PHR_HIDDEN_DIM + + (long long)h * PHR_HIDDEN_DIM + tid; + const float ctx = context_h[ctx_idx]; + const float w = w_res[h * PHR_HIDDEN_DIM + tid]; + + d_w_res_scratch[ctx_idx] += g * ctx; + d_context_h[ctx_idx] += g * w; + } + if (tid == 0) { + d_bias_res_scratch[(long long)b * PHR_N_HORIZONS + h] += g; + } + } +} diff --git a/crates/ml-alpha/src/lib.rs b/crates/ml-alpha/src/lib.rs index 018fb12ae..d64df1feb 100644 --- a/crates/ml-alpha/src/lib.rs +++ b/crates/ml-alpha/src/lib.rs @@ -31,6 +31,7 @@ pub mod eval; pub mod heads; pub mod isv; pub mod per_horizon_attention_pool; +pub mod per_horizon_residual_head; pub mod pinned; pub mod pinned_mem; pub mod trainer; diff --git a/crates/ml-alpha/src/per_horizon_residual_head.rs b/crates/ml-alpha/src/per_horizon_residual_head.rs new file mode 100644 index 000000000..df8ef9ce2 --- /dev/null +++ b/crates/ml-alpha/src/per_horizon_residual_head.rs @@ -0,0 +1,113 @@ +//! Per-horizon residual head kernel host wrapper (C22). +//! +//! See `crates/ml-alpha/cuda/per_horizon_residual_head.cu` for kernel +//! math. v1 of this binding: standalone forward + backward, validated +//! via numgrad parity test. Trainer integration (gating the residual +//! addition behind a learnable α-scalar) is a follow-up commit. + +use anyhow::{Context, Result}; +use cudarc::driver::{ + CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg, +}; +use std::sync::Arc; + +pub const PHR_HIDDEN_DIM: usize = 128; +pub const PHR_BLOCK: usize = 128; +pub const PHR_N_HORIZONS: usize = 5; + +const CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/per_horizon_residual_head.cubin")); + +pub struct PerHorizonResidualHead { + _module: Arc, + fwd_fn: CudaFunction, + bwd_fn: CudaFunction, + stream: Arc, +} + +impl PerHorizonResidualHead { + pub fn new(ctx: &Arc, stream: Arc) -> Result { + let module = ctx + .load_cubin(CUBIN.to_vec()) + .context("load per_horizon_residual_head cubin")?; + let fwd_fn = module + .load_function("per_horizon_residual_head_fwd") + .context("load per_horizon_residual_head_fwd")?; + let bwd_fn = module + .load_function("per_horizon_residual_head_bwd") + .context("load per_horizon_residual_head_bwd")?; + Ok(Self { + _module: module, + fwd_fn, + bwd_fn, + stream, + }) + } + + /// Forward. context_h: [B, N_HORIZONS, HIDDEN_DIM]. + /// w_res: [N_HORIZONS, HIDDEN_DIM]. bias_res: [N_HORIZONS]. + /// residual_out: [B, N_HORIZONS] (written). + pub fn forward( + &self, + context_h: &CudaSlice, + w_res: &CudaSlice, + bias_res: &CudaSlice, + n_batch: i32, + residual_out: &mut CudaSlice, + ) -> Result<()> { + let cfg = LaunchConfig { + grid_dim: (n_batch as u32, 1, 1), + block_dim: (PHR_BLOCK as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.fwd_fn); + unsafe { + launch + .arg(context_h) + .arg(w_res) + .arg(bias_res) + .arg(&n_batch) + .arg(residual_out) + .launch(cfg) + .context("per_horizon_residual_head_fwd")?; + } + self.stream.synchronize()?; + Ok(()) + } + + /// Backward. d_w_res_scratch + d_bias_res_scratch are per-block + /// scratch tensors; reduce across batch with `reduce_axis0` to get + /// the shared `d_w_res[N_HORIZONS, HIDDEN_DIM]` + `d_bias[N_HORIZONS]`. + /// d_context_h is per-batch indexed; updates +=. + pub fn backward( + &self, + context_h: &CudaSlice, + w_res: &CudaSlice, + grad_residual: &CudaSlice, + n_batch: i32, + d_w_res_scratch: &mut CudaSlice, + d_bias_res_scratch: &mut CudaSlice, + d_context_h: &mut CudaSlice, + ) -> Result<()> { + let cfg = LaunchConfig { + grid_dim: (n_batch as u32, 1, 1), + block_dim: (PHR_BLOCK as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.bwd_fn); + unsafe { + launch + .arg(context_h) + .arg(w_res) + .arg(grad_residual) + .arg(&n_batch) + .arg(d_w_res_scratch) + .arg(d_bias_res_scratch) + .arg(d_context_h) + .launch(cfg) + .context("per_horizon_residual_head_bwd")?; + } + self.stream.synchronize()?; + Ok(()) + } +} diff --git a/crates/ml-alpha/tests/per_horizon_residual_head_numgrad.rs b/crates/ml-alpha/tests/per_horizon_residual_head_numgrad.rs new file mode 100644 index 000000000..b82654cea --- /dev/null +++ b/crates/ml-alpha/tests/per_horizon_residual_head_numgrad.rs @@ -0,0 +1,163 @@ +//! Numerical-gradient parity check for the per-horizon residual head +//! kernel (C22). Mirrors the C21 numgrad pattern. +//! +//! Three grads checked: +//! d_w_res [N_HORIZONS, HIDDEN_DIM] +//! d_bias_res [N_HORIZONS] +//! d_context_h [B, N_HORIZONS, HIDDEN_DIM] +//! +//! Loss = Σ residual_out (so d_residual = 1 everywhere). + +use anyhow::Result; +use cudarc::driver::CudaSlice; +use ml_alpha::per_horizon_residual_head::{PerHorizonResidualHead, PHR_HIDDEN_DIM, PHR_N_HORIZONS}; +use ml_core::device::MlDevice; +use rand::Rng; +use rand::SeedableRng; +use rand_chacha::ChaCha8Rng; + +const B: usize = 3; + +fn try_dev() -> Option { + match MlDevice::cuda(0) { + Ok(d) => Some(d), + Err(e) => { + eprintln!("skipping: cuda device unavailable ({e})"); + None + } + } +} + +fn alloc_upload(stream: &std::sync::Arc, host: &[f32]) -> CudaSlice { + let mut buf = stream.alloc_zeros::(host.len()).expect("alloc"); + stream.memcpy_htod(host, &mut buf).expect("htod"); + buf +} + +fn download(stream: &std::sync::Arc, src: &CudaSlice) -> Vec { + let mut out = vec![0.0f32; src.len()]; + stream.memcpy_dtoh(src, out.as_mut_slice()).expect("dtoh"); + out +} + +fn forward_loss( + head: &PerHorizonResidualHead, + stream: &std::sync::Arc, + context: &[f32], + w_res: &[f32], + bias: &[f32], +) -> f32 { + let ctx_d = alloc_upload(stream, context); + let w_d = alloc_upload(stream, w_res); + let b_d = alloc_upload(stream, bias); + let mut out_d = stream.alloc_zeros::(B * PHR_N_HORIZONS).unwrap(); + head.forward(&ctx_d, &w_d, &b_d, B as i32, &mut out_d).unwrap(); + download(stream, &out_d).iter().sum() +} + +#[test] +#[ignore = "requires CUDA"] +fn forward_then_backward_matches_central_difference() -> Result<()> { + let Some(dev) = try_dev() else { return Ok(()); }; + let ctx_h = dev.cuda_context()?.clone(); + let stream = dev.cuda_stream()?.clone(); + let head = PerHorizonResidualHead::new(&ctx_h, stream.clone())?; + + let mut rng = ChaCha8Rng::seed_from_u64(0xBEEF_F00D); + let context_host: Vec = (0..B * PHR_N_HORIZONS * PHR_HIDDEN_DIM) + .map(|_| rng.gen_range(-0.5..0.5)).collect(); + let w_host: Vec = (0..PHR_N_HORIZONS * PHR_HIDDEN_DIM) + .map(|_| rng.gen_range(-0.1..0.1)).collect(); + let bias_host: Vec = (0..PHR_N_HORIZONS) + .map(|_| rng.gen_range(-0.05..0.05)).collect(); + + // Analytical backward with grad_residual = 1 everywhere. + let ctx_d = alloc_upload(&stream, &context_host); + let w_d = alloc_upload(&stream, &w_host); + let b_d = alloc_upload(&stream, &bias_host); + let mut out_d = stream.alloc_zeros::(B * PHR_N_HORIZONS)?; + head.forward(&ctx_d, &w_d, &b_d, B as i32, &mut out_d)?; + + let grad_res_host = vec![1.0f32; B * PHR_N_HORIZONS]; + let grad_res_d = alloc_upload(&stream, &grad_res_host); + let mut d_w_scratch_d = stream.alloc_zeros::(B * PHR_N_HORIZONS * PHR_HIDDEN_DIM)?; + let mut d_b_scratch_d = stream.alloc_zeros::(B * PHR_N_HORIZONS)?; + let mut d_ctx_d = stream.alloc_zeros::(B * PHR_N_HORIZONS * PHR_HIDDEN_DIM)?; + head.backward( + &ctx_d, &w_d, &grad_res_d, B as i32, + &mut d_w_scratch_d, &mut d_b_scratch_d, &mut d_ctx_d, + )?; + let d_w_scratch = download(&stream, &d_w_scratch_d); + let d_b_scratch = download(&stream, &d_b_scratch_d); + let d_ctx = download(&stream, &d_ctx_d); + + // Reduce d_w_scratch across batch → [N_HORIZONS, HIDDEN_DIM]. + let mut d_w_analytical = vec![0.0f32; PHR_N_HORIZONS * PHR_HIDDEN_DIM]; + for b in 0..B { + for i in 0..PHR_N_HORIZONS * PHR_HIDDEN_DIM { + d_w_analytical[i] += d_w_scratch[b * PHR_N_HORIZONS * PHR_HIDDEN_DIM + i]; + } + } + // Reduce d_b_scratch across batch → [N_HORIZONS]. + let mut d_b_analytical = vec![0.0f32; PHR_N_HORIZONS]; + for b in 0..B { + for h in 0..PHR_N_HORIZONS { + d_b_analytical[h] += d_b_scratch[b * PHR_N_HORIZONS + h]; + } + } + + let eps = 1e-2f32; + let tol = 5e-2f32; + + // 1) Probe d_w_res at random indices. + for _ in 0..8 { + let idx = rng.gen_range(0..PHR_N_HORIZONS * PHR_HIDDEN_DIM); + let mut wp = w_host.clone(); wp[idx] += eps; + let mut wm = w_host.clone(); wm[idx] -= eps; + let lp = forward_loss(&head, &stream, &context_host, &wp, &bias_host); + let lm = forward_loss(&head, &stream, &context_host, &wm, &bias_host); + let fd = (lp - lm) / (2.0 * eps); + let an = d_w_analytical[idx]; + let abs_err = (fd - an).abs(); + let rel_err = abs_err / fd.abs().max(1e-6); + assert!( + abs_err < 5e-3 || rel_err < tol, + "d_w_res[{idx}]: analytical={an:.6} fd={fd:.6} abs={abs_err:.6} rel={rel_err:.4}" + ); + } + + // 2) Probe d_bias_res at every horizon. + for h in 0..PHR_N_HORIZONS { + let mut bp = bias_host.clone(); bp[h] += eps; + let mut bm = bias_host.clone(); bm[h] -= eps; + let lp = forward_loss(&head, &stream, &context_host, &w_host, &bp); + let lm = forward_loss(&head, &stream, &context_host, &w_host, &bm); + let fd = (lp - lm) / (2.0 * eps); + let an = d_b_analytical[h]; + let abs_err = (fd - an).abs(); + let rel_err = abs_err / fd.abs().max(1e-6); + assert!( + abs_err < 5e-3 || rel_err < tol, + "d_bias_res[{h}]: analytical={an:.6} fd={fd:.6} abs={abs_err:.6} rel={rel_err:.4}" + ); + } + + // 3) Probe d_context_h at random indices. + for _ in 0..8 { + let idx = rng.gen_range(0..B * PHR_N_HORIZONS * PHR_HIDDEN_DIM); + let mut cp = context_host.clone(); cp[idx] += eps; + let mut cm = context_host.clone(); cm[idx] -= eps; + let lp = forward_loss(&head, &stream, &cp, &w_host, &bias_host); + let lm = forward_loss(&head, &stream, &cm, &w_host, &bias_host); + let fd = (lp - lm) / (2.0 * eps); + let an = d_ctx[idx]; + let abs_err = (fd - an).abs(); + let rel_err = abs_err / fd.abs().max(1e-6); + assert!( + abs_err < 5e-3 || rel_err < tol, + "d_context_h[{idx}]: analytical={an:.6} fd={fd:.6} abs={abs_err:.6} rel={rel_err:.4}" + ); + } + + Ok(()) +}