feat(ml-alpha): PerceptionTrainer wires C21+C22 fwd+bwd+AdamW (C25)

The per-horizon attention pool from C21+C22 is now fully integrated
into step_batched's hot loop. Forward and backward both flow; AdamW
updates all four new parameter groups (Q_h, w_res, bias_res, α) every
training step. At α=0 init the contribution is byte-identical to
baseline; training discovers whether α should grow.

New kernel: cuda/per_horizon_prob_blend.cu
  per_horizon_prob_blend_fwd
    Reads logit_per_k_d (already stored by GRN forward) +
    sigmoid(logit_baseline + tanh(α[h]) * residual[b, h]) →
    overwrites probs_per_k_d in place. At α=0: r_contrib=0, output
    == sigmoid(logit_baseline) == probs_baseline → bit-identical.

  per_horizon_prob_blend_reduce_alpha_residual
    Reads probs_per_k (= p_final, post-blend) + grad_probs_per_k
    (= ∂L/∂p_final from BCE) and computes:
      d_logit[k,b,h] = grad_probs[k,b,h] * p_final * (1 - p_final)
      d_residual[b,h] = tanh(α[h]) * Σ_k d_logit[k,b,h]
      d_alpha[h]      = sech²(α[h]) * Σ_{k,b} d_logit[k,b,h] * residual[b,h]
    No separate prob_blend_bwd needed — chain-rule equivalence
    ∂L/∂logit_baseline = ∂L/∂r_contrib (both flow through the same
    sigmoid derivative) means the existing GRN backward is UNCHANGED.

trainer/per_horizon_state.rs extensions:
  forward_with_blend(ln_b_out, logit_per_k, probs_per_k)
    Pool fwd → context_h; head fwd → residual; prob_blend fwd
    in-place rewrites probs_per_k.

  backward_through_blend(probs_per_k, grad_probs_per_k, ln_b_out,
                         grad_ln_b_out_target)
    Reduce kernel → d_residual + d_alpha. Then:
      head bwd  → d_w_res_scratch, d_bias_res_scratch, d_context.
      pool bwd  → d_q_h_scratch, += grad_h_enriched_seq_d.
    Per-batch scratches reduced to shared grads host-side
    (n_batch ≤ 64 → sub-millisecond on host).

  adamw_step()
    Steps the four optimizers using the shared grad buffers.

  zero_grads()
    Called once per step before forward to clear scratch.

trainer/perception.rs step_batched integration:
  ── 4.5 (after GRN K-loop, before BCE): zero_grads + forward_with_blend
       overwrites probs_per_k_d with p_final.
  ── 5  (existing BCE consumes probs_per_k_d as today; grad_probs is
        now ∂L/∂p_final automatically).
  ── 5a (after BCE, before ISV-lambda + heads bwd): backward_through_blend.
       Existing GRN bwd path is UNTOUCHED — the chain rule absorbs
       the bias.
  ── 9  (after existing 17 AdamW group steps): per_horizon.adamw_step
       updates Q_h, w_res, bias_res, α.

Verification:
  - 34 ml-alpha lib tests still green.
  - Per-horizon kernel numgrad parity (C21, C22) still green.
  - Per-horizon end-to-end pipeline smoke (C23, including the
    alpha=0 byte-identity invariant) still green.
  - Full workspace builds clean.

Closes the kernels+wiring portion of #203 (per-horizon attention pool
kernels + wiring). What remains (#204): 30-epoch × 3-fold A/B vs
single-Q baseline. The branch is ready for that sweep when GPU time
is budgeted; the implementation is structurally adoption-safe
(α=0 → identity to baseline) so it can be merged before the A/B if
desired.

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

crates/ml-alpha/build.rs

@@ -22,6 +22,7 @@ const KERNELS: &[&str] = &[
     "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
+    "per_horizon_prob_blend",        // C25: closed-form logit-bias applied in probability space
     "reduce_axis0",                  // Phase B: cross-batch param-grad reducer
 ];
 

crates/ml-alpha/cuda/per_horizon_prob_blend.cu

@@ -0,0 +1,104 @@
+// per_horizon_prob_blend.cu — logit-space additive bias over per-K probs.
+//
+// Existing pipeline (perception.rs::step_batched) stores per-K BASELINE
+// logits in `logit_per_k_d[K, B, N_HORIZONS]` AND their sigmoid in
+// `probs_per_k_d[K, B, N_HORIZONS]`. The per-horizon attention pool's
+// learnable α-gated residual adds a logit-space shift:
+//
+//   r_contrib[b, h]   = tanh(α[h]) * residual[b, h]
+//   probs_final[k,b,h] = sigmoid(logit_baseline[k,b,h] + r_contrib[b,h])
+//
+// At α=0 ⇒ r_contrib=0 ⇒ probs_final = sigmoid(logit_baseline) =
+// probs_baseline → bit-identical to baseline (C23 invariant extended
+// into trainer).
+//
+// Forward (fwd kernel): overwrites probs_per_k_d in place with
+// probs_final. BCE consumes probs_per_k_d as-is afterward.
+//
+// Backward: NO separate kernel needed for d_logit_baseline. The existing
+// heads_grn_bwd already computes d_logit_baseline = grad_probs * p *
+// (1 - p), and since p_final = sigmoid(logit_baseline + r_contrib),
+// d_logit_baseline = d_p_final * sigmoid_deriv(logit_baseline + r_contrib)
+// = d_p_final * p_final * (1 - p_final) — which is exactly what the
+// existing GRN backward kernel computes. So the GRN backward path is
+// UNCHANGED — the chain rule absorbs the bias automatically.
+//
+// What we DO need: compute d_r_contrib_per_k[k, b, h] = the same
+// quantity (= d_logit_baseline by chain-rule equivalence) and reduce
+// over k + b to produce d_residual[b, h] and d_alpha[h].
+
+#define PHB_N_HORIZONS 5
+
+extern "C" __global__ void per_horizon_prob_blend_fwd(
+    const float* __restrict__ alpha,         // [N_HORIZONS]
+    const float* __restrict__ residual,      // [B, N_HORIZONS]
+    const float* __restrict__ logit_per_k,   // [K, B, N_HORIZONS] from GRN forward
+    float*       __restrict__ probs_per_k,   // [K, B, N_HORIZONS] OVERWRITTEN with probs_final
+    int n_batch,
+    int k_seq
+) {
+    int b = blockIdx.x;
+    int h = threadIdx.x;
+    if (b >= n_batch || h >= PHB_N_HORIZONS) return;
+
+    const float r_contrib = tanhf(alpha[h]) * residual[b * PHB_N_HORIZONS + h];
+
+    for (int k = 0; k < k_seq; ++k) {
+        const long long idx =
+            (long long)k * n_batch * PHB_N_HORIZONS
+          + (long long)b * PHB_N_HORIZONS
+          + (long long)h;
+        const float z = logit_per_k[idx] + r_contrib;
+        probs_per_k[idx] = 1.0f / (1.0f + expf(-z));
+    }
+}
+
+// Reduce step. Reads grad_probs_per_k (= ∂L/∂p_final, from BCE) and
+// probs_per_k (= p_final, in-place overwrite from fwd kernel) to derive:
+//
+//   d_logit_per_k[k, b, h] = grad_probs[k, b, h] * p_final * (1 - p_final)
+//
+// then sums over the appropriate axes:
+//
+//   d_residual[b, h] = tanh(α[h]) * Σ_k d_logit_per_k[k, b, h]
+//   d_alpha[h]       = (1 - tanh²(α[h])) * Σ_{k, b} d_logit_per_k[k, b, h] * residual[b, h]
+//
+// One block per horizon. Single-thread block (tid==0) does both
+// reductions sequentially — sizes are tiny (K up to a few hundred, B
+// up to 64) so a parallel reduce isn't worth the complexity.
+extern "C" __global__ void per_horizon_prob_blend_reduce_alpha_residual(
+    const float* __restrict__ alpha,         // [N_HORIZONS]
+    const float* __restrict__ residual,      // [B, N_HORIZONS]
+    const float* __restrict__ probs_per_k,   // [K, B, N_HORIZONS] (= probs_final)
+    const float* __restrict__ grad_probs,    // [K, B, N_HORIZONS] (= ∂L/∂p_final from BCE)
+    int n_batch,
+    int k_seq,
+    float*       __restrict__ d_residual,    // [B, N_HORIZONS] (written)
+    float*       __restrict__ d_alpha        // [N_HORIZONS] (written)
+) {
+    int h = blockIdx.x;
+    if (h >= PHB_N_HORIZONS || threadIdx.x != 0) return;
+
+    const float a = alpha[h];
+    const float ta = tanhf(a);
+    const float sech2 = 1.0f - ta * ta;
+
+    float alpha_sum = 0.0f;
+    for (int b = 0; b < n_batch; ++b) {
+        const float resid_bh = residual[b * PHB_N_HORIZONS + h];
+        float k_sum = 0.0f;
+        for (int k = 0; k < k_seq; ++k) {
+            const long long idx =
+                (long long)k * n_batch * PHB_N_HORIZONS
+              + (long long)b * PHB_N_HORIZONS
+              + (long long)h;
+            const float p = probs_per_k[idx];
+            const float gp = grad_probs[idx];
+            const float d_logit = gp * p * (1.0f - p);
+            k_sum += d_logit;
+            alpha_sum += d_logit * resid_bh;
+        }
+        d_residual[b * PHB_N_HORIZONS + h] = ta * k_sum;
+    }
+    d_alpha[h] = sech2 * alpha_sum;
+}

crates/ml-alpha/src/trainer/per_horizon_state.rs

@@ -13,7 +13,9 @@
 //! the residual signal correlates with loss gradient.
 
 use anyhow::{Context, Result};
-use cudarc::driver::{CudaSlice, CudaStream};
+use cudarc::driver::{
+    CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg,
+};
 use ml_core::device::MlDevice;
 use rand::{Rng, SeedableRng};
 use rand_chacha::ChaCha8Rng;
@@ -24,6 +26,9 @@ use crate::per_horizon_attention_pool::PerHorizonAttentionPool;
 use crate::per_horizon_residual_head::PerHorizonResidualHead;
 use crate::trainer::optim::AdamW;
 
+const PROB_BLEND_CUBIN: &[u8] =
+    include_bytes!(concat!(env!("OUT_DIR"), "/per_horizon_prob_blend.cubin"));
+
 pub struct PerHorizonTrainState {
     // Kernel bindings.
     pub pool: PerHorizonAttentionPool,
@@ -63,6 +68,15 @@ pub struct PerHorizonTrainState {
     pub n_batch: usize,
     pub k_seq: usize,
     stream: Arc<CudaStream>,
+
+    // C25 prob-blend kernel handles + per-K scratch (baseline probs stash +
+    // ∂L/∂r_contrib per K position).
+    _prob_blend_module: Arc<CudaModule>,
+    prob_blend_fwd_fn: CudaFunction,
+    prob_blend_reduce_fn: CudaFunction,
+    /// d_α [N_HORIZONS] — what AdamW updates for α. Written by the
+    /// reduce kernel after BCE backward has populated grad_probs_per_k.
+    pub d_alpha_reduced_d: CudaSlice<f32>,
 }
 
 impl PerHorizonTrainState {
@@ -78,6 +92,16 @@ impl PerHorizonTrainState {
         let head = PerHorizonResidualHead::new(ctx, stream.clone())
             .context("PerHorizonResidualHead")?;
 
+        let prob_blend_module = ctx
+            .load_cubin(PROB_BLEND_CUBIN.to_vec())
+            .context("load per_horizon_prob_blend cubin")?;
+        let prob_blend_fwd_fn = prob_blend_module
+            .load_function("per_horizon_prob_blend_fwd")
+            .context("per_horizon_prob_blend_fwd")?;
+        let prob_blend_reduce_fn = prob_blend_module
+            .load_function("per_horizon_prob_blend_reduce_alpha_residual")
+            .context("per_horizon_prob_blend_reduce_alpha_residual")?;
+
         let mut rng = ChaCha8Rng::seed_from_u64(seed);
 
         // Q_h: Xavier-style init at 1/sqrt(HIDDEN_DIM). Per the spec
@@ -130,6 +154,10 @@ impl PerHorizonTrainState {
         // default; can be tuned per spec §5 open question 2.
         let opt_alpha = AdamW::new(dev, N_HORIZONS, lr * 0.25).context("opt_alpha")?;
 
+        let d_alpha_reduced_d = stream
+            .alloc_zeros::<f32>(N_HORIZONS)
+            .context("d_alpha_reduced alloc")?;
+
         Ok(Self {
             pool, head,
             q_h_d, w_res_d, bias_res_d, alpha_d,
@@ -140,9 +168,171 @@ impl PerHorizonTrainState {
             opt_q_h, opt_w_res, opt_bias_res, opt_alpha,
             n_batch, k_seq,
             stream,
+            _prob_blend_module: prob_blend_module,
+            prob_blend_fwd_fn,
+            prob_blend_reduce_fn,
+            d_alpha_reduced_d,
         })
     }
 
+    /// C25 forward: per-horizon attention pool → residual head → in-place
+    /// logit-bias rewrite of `probs_per_k` using `logit_per_k` as the
+    /// baseline (saved by the GRN forward).
+    /// At α=0 (init), r_contrib = 0 → probs unchanged → bit-identical
+    /// to baseline (extension of the C23 identity invariant into the
+    /// trainer hot loop).
+    ///
+    /// `ln_b_out` is the perception trainer's LN_b output `[B, K, HIDDEN_DIM]`
+    /// (same tensor that feeds the existing single-Q attention pool).
+    /// `logit_per_k` is `[K, B, N_HORIZONS]` from GRN forward.
+    /// `probs_per_k` is `[K, B, N_HORIZONS]` — overwritten in-place
+    /// with `sigmoid(logit_per_k + tanh(α) * residual)`.
+    pub fn forward_with_blend(
+        &mut self,
+        ln_b_out: &CudaSlice<f32>,
+        logit_per_k: &CudaSlice<f32>,
+        probs_per_k: &mut CudaSlice<f32>,
+    ) -> Result<()> {
+        let n = self.n_batch as i32;
+        let k = self.k_seq as i32;
+
+        // 1. Per-horizon attention pool → context_h, attn_weights.
+        self.pool.forward(&self.q_h_d, ln_b_out, n, k,
+            &mut self.context_d, &mut self.attn_weights_d)?;
+
+        // 2. Residual head → residual[B, N_HORIZONS].
+        self.head.forward(&self.context_d, &self.w_res_d, &self.bias_res_d, n,
+            &mut self.residual_d)?;
+
+        // 3. In-place logit-bias rewrite: probs_per_k ← sigmoid(logit_per_k + r).
+        let cfg = LaunchConfig {
+            grid_dim: (self.n_batch as u32, 1, 1),
+            block_dim: (N_HORIZONS as u32, 1, 1),
+            shared_mem_bytes: 0,
+        };
+        let mut launch = self.stream.launch_builder(&self.prob_blend_fwd_fn);
+        unsafe {
+            launch
+                .arg(&self.alpha_d)
+                .arg(&self.residual_d)
+                .arg(logit_per_k)
+                .arg(probs_per_k)
+                .arg(&n)
+                .arg(&k)
+                .launch(cfg)
+                .context("per_horizon_prob_blend_fwd")?;
+        }
+        self.stream.synchronize()?;
+        Ok(())
+    }
+
+    /// C25 backward. Called AFTER BCE has populated `grad_probs_per_k`
+    /// (= ∂L/∂p_final).
+    ///
+    /// Key insight: because the bias is additive INSIDE the sigmoid,
+    /// ∂L/∂logit_baseline = ∂L/∂r_contrib = ∂L/∂p_final * p_final * (1 - p_final).
+    /// The existing GRN backward already computes the LHS via the same
+    /// formula (using probs_per_k = p_final), so the GRN backward path
+    /// is UNTOUCHED. We just read the same gradient ourselves and reduce
+    /// to get d_residual + d_alpha.
+    ///
+    /// `probs_per_k` is the trainer's probs buffer (post-blend, == p_final).
+    /// `grad_probs_per_k` is from BCE backward.
+    /// `ln_b_out` + `grad_ln_b_out` are LN_b's fwd output and accumulating
+    /// gradient — the attention pool bwd `+=`'s onto grad_ln_b_out.
+    pub fn backward_through_blend(
+        &mut self,
+        probs_per_k: &CudaSlice<f32>,
+        grad_probs_per_k: &CudaSlice<f32>,
+        ln_b_out: &CudaSlice<f32>,
+        grad_ln_b_out: &mut CudaSlice<f32>,
+    ) -> Result<()> {
+        let n = self.n_batch as i32;
+        let k = self.k_seq as i32;
+
+        // 1. Reduce: compute d_logit_baseline = grad_probs * p * (1-p)
+        //    inline + reduce over k for d_residual, over k+b for d_alpha.
+        let cfg_per_h = LaunchConfig {
+            grid_dim: (N_HORIZONS as u32, 1, 1),
+            block_dim: (1, 1, 1),
+            shared_mem_bytes: 0,
+        };
+        let mut launch = self.stream.launch_builder(&self.prob_blend_reduce_fn);
+        unsafe {
+            launch
+                .arg(&self.alpha_d)
+                .arg(&self.residual_d)
+                .arg(probs_per_k)
+                .arg(grad_probs_per_k)
+                .arg(&n)
+                .arg(&k)
+                .arg(&mut self.grad_residual_d)
+                .arg(&mut self.d_alpha_reduced_d)
+                .launch(cfg_per_h)
+                .context("per_horizon_prob_blend_reduce_alpha_residual")?;
+        }
+        self.stream.synchronize()?;
+
+        // 2. Residual-head bwd: consumes d_residual → produces
+        //    d_w_res_scratch, d_bias_res_scratch, d_context.
+        self.head.backward(
+            &self.context_d, &self.w_res_d, &self.grad_residual_d, n,
+            &mut self.grad_w_res_scratch_d,
+            &mut self.grad_bias_res_scratch_d,
+            &mut self.grad_context_d,
+        )?;
+
+        // 3. Attention-pool bwd: consumes d_context → produces
+        //    d_q_h_scratch + adds to grad_ln_b_out.
+        self.pool.backward(
+            &self.q_h_d, ln_b_out, &self.attn_weights_d, &self.grad_context_d,
+            n, k,
+            &mut self.grad_q_h_scratch_d, grad_ln_b_out,
+        )?;
+
+        // 4. Reduce per-batch scratches to shared grads.
+        self.reduce_per_batch_scratches_to_shared()?;
+
+        Ok(())
+    }
+
+    fn reduce_per_batch_scratches_to_shared(&mut self) -> Result<()> {
+        // Read per-batch scratches; sum over batch axis; push reduced
+        // shared grads back. For typical n_batch ≤ 64 this is
+        // negligible compared to the kernel launches.
+        let n_qh = N_HORIZONS * HIDDEN_DIM;
+        let n_bres = N_HORIZONS;
+        let mut qh_scratch = vec![0.0_f32; self.n_batch * n_qh];
+        let mut wres_scratch = vec![0.0_f32; self.n_batch * n_qh];
+        let mut bres_scratch = vec![0.0_f32; self.n_batch * n_bres];
+        self.stream.memcpy_dtoh(&self.grad_q_h_scratch_d,       qh_scratch.as_mut_slice())?;
+        self.stream.memcpy_dtoh(&self.grad_w_res_scratch_d,     wres_scratch.as_mut_slice())?;
+        self.stream.memcpy_dtoh(&self.grad_bias_res_scratch_d,  bres_scratch.as_mut_slice())?;
+
+        let mut qh_reduced   = vec![0.0_f32; n_qh];
+        let mut wres_reduced = vec![0.0_f32; n_qh];
+        let mut bres_reduced = vec![0.0_f32; n_bres];
+        for b in 0..self.n_batch {
+            for i in 0..n_qh   { qh_reduced[i]   += qh_scratch  [b * n_qh   + i]; }
+            for i in 0..n_qh   { wres_reduced[i] += wres_scratch[b * n_qh   + i]; }
+            for i in 0..n_bres { bres_reduced[i] += bres_scratch[b * n_bres + i]; }
+        }
+        self.stream.memcpy_htod(&qh_reduced,   &mut self.grad_q_h_d)?;
+        self.stream.memcpy_htod(&wres_reduced, &mut self.grad_w_res_d)?;
+        self.stream.memcpy_htod(&bres_reduced, &mut self.grad_bias_res_d)?;
+        Ok(())
+    }
+
+    /// AdamW step for all four per-horizon param groups. Called after
+    /// `backward_through_blend` populates the shared grad buffers.
+    pub fn adamw_step(&mut self) -> Result<()> {
+        self.opt_q_h.step(&mut self.q_h_d, &self.grad_q_h_d)?;
+        self.opt_w_res.step(&mut self.w_res_d, &self.grad_w_res_d)?;
+        self.opt_bias_res.step(&mut self.bias_res_d, &self.grad_bias_res_d)?;
+        self.opt_alpha.step(&mut self.alpha_d, &self.d_alpha_reduced_d)?;
+        Ok(())
+    }
+
     /// Zero all gradient scratch buffers between training steps.
     pub fn zero_grads(&mut self) -> Result<()> {
         let stream = &self.stream;

crates/ml-alpha/src/trainer/perception.rs

@@ -1665,6 +1665,21 @@ impl PerceptionTrainer {
         }
         drop((_g_hpk, _g_probs, _g_z1, _g_a1, _g_z2, _g_gate, _g_main, _g_logit));
 
+        // ── 4.5. C25 per-horizon attention pool forward + prob-blend.
+        //         Reads ln_out_d + logit_per_k_d (baseline GRN logits)
+        //         and overwrites probs_per_k_d in place with
+        //         sigmoid(logit_baseline + tanh(α) * residual).
+        //         At α=0 (init): r_contrib=0, probs unchanged → baseline
+        //         training is bit-identical to pre-C25 (per the C23
+        //         alpha_zero_init_is_identity invariant extended into
+        //         the trainer hot loop).
+        self.per_horizon.zero_grads()?;
+        self.per_horizon.forward_with_blend(
+            &self.ln_out_d,
+            &self.logit_per_k_d,
+            &mut self.probs_per_k_d,
+        )?;
+
         // ── 5. Fused multi-horizon BCE over the full [K*B, N_HORIZONS]
         //       grid. The kernel doesn't distinguish position-vs-batch
         //       since the per-horizon weight is identified by `i % n_horizons`.
@@ -1688,6 +1703,23 @@ impl PerceptionTrainer {
             unsafe { launch.launch(bce_cfg).context("bce launch")?; }
         }
 
+        // ── 5a. C25 per-horizon attention pool backward.
+        //         Reads probs_per_k_d (now post-blend, = p_final) +
+        //         grad_probs_per_k_d (= ∂L/∂p_final from BCE) and:
+        //           1. Computes d_residual + d_α via per-horizon reduce
+        //              (chain rule equivalence makes the existing GRN
+        //              backward unmodified — see kernel header).
+        //           2. Runs residual_head bwd → d_w_res, d_bias_res, d_context.
+        //           3. Runs attention_pool bwd → d_q_h, += grad_h_enriched_seq_d.
+        //         AdamW updates for the four per-horizon param groups
+        //         land in section 9 alongside the existing 17 groups.
+        self.per_horizon.backward_through_blend(
+            &self.probs_per_k_d,
+            &self.grad_probs_per_k_d,
+            &self.ln_out_d,
+            self.grad_h_enriched_seq_d.data_mut(),
+        )?;
+
         // ── 5b. ISV-driven per-horizon EMA + lambda. Updates the EMA
         //       of unweighted per-horizon BCE and emits a clamped
         //       multiplier `lambda_d[h]` per
@@ -2200,6 +2232,12 @@ impl PerceptionTrainer {
         self.opt_vsn_w.step(&mut self.vsn_w_d, &self.grad_vsn_w_d)?;
         self.opt_vsn_b.step(&mut self.vsn_b_d, &self.grad_vsn_b_d)?;
         self.opt_attn_q.step(&mut self.attn_q_d, &self.grad_attn_q_d)?;
+
+        // C25: per-horizon attention pool AdamW step (4 param groups:
+        // q_h, w_res, bias_res, α). At α=0 init + with backward-flow
+        // proving residual signal correlates with loss gradient, α
+        // will move away from 0 if the per-horizon path is useful.
+        self.per_horizon.adamw_step()?;
         // GPU-resident grad-clip + AdamW (zero host roundtrips).
         // Replaces step_from_buffers which did 9× memcpy_dtoh per step.
         self.mamba2_adamw