feat(ml-alpha): per-horizon residual head kernel + numgrad parity (C22)

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 <noreply@anthropic.com>

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
 ];
 

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;
+        }
+    }
+}

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;

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<CudaModule>,
+    fwd_fn: CudaFunction,
+    bwd_fn: CudaFunction,
+    stream: Arc<CudaStream>,
+}
+
+impl PerHorizonResidualHead {
+    pub fn new(ctx: &Arc<CudaContext>, stream: Arc<CudaStream>) -> Result<Self> {
+        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<f32>,
+        w_res: &CudaSlice<f32>,
+        bias_res: &CudaSlice<f32>,
+        n_batch: i32,
+        residual_out: &mut CudaSlice<f32>,
+    ) -> 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<f32>,
+        w_res: &CudaSlice<f32>,
+        grad_residual: &CudaSlice<f32>,
+        n_batch: i32,
+        d_w_res_scratch: &mut CudaSlice<f32>,
+        d_bias_res_scratch: &mut CudaSlice<f32>,
+        d_context_h: &mut CudaSlice<f32>,
+    ) -> 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(())
+    }
+}

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<MlDevice> {
+    match MlDevice::cuda(0) {
+        Ok(d) => Some(d),
+        Err(e) => {
+            eprintln!("skipping: cuda device unavailable ({e})");
+            None
+        }
+    }
+}
+
+fn alloc_upload(stream: &std::sync::Arc<cudarc::driver::CudaStream>, host: &[f32]) -> CudaSlice<f32> {
+    let mut buf = stream.alloc_zeros::<f32>(host.len()).expect("alloc");
+    stream.memcpy_htod(host, &mut buf).expect("htod");
+    buf
+}
+
+fn download(stream: &std::sync::Arc<cudarc::driver::CudaStream>, src: &CudaSlice<f32>) -> Vec<f32> {
+    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<cudarc::driver::CudaStream>,
+    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::<f32>(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<f32> = (0..B * PHR_N_HORIZONS * PHR_HIDDEN_DIM)
+        .map(|_| rng.gen_range(-0.5..0.5)).collect();
+    let w_host: Vec<f32> = (0..PHR_N_HORIZONS * PHR_HIDDEN_DIM)
+        .map(|_| rng.gen_range(-0.1..0.1)).collect();
+    let bias_host: Vec<f32> = (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::<f32>(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::<f32>(B * PHR_N_HORIZONS * PHR_HIDDEN_DIM)?;
+    let mut d_b_scratch_d  = stream.alloc_zeros::<f32>(B * PHR_N_HORIZONS)?;
+    let mut d_ctx_d        = stream.alloc_zeros::<f32>(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(())
+}