From f50b466f7737365a845da76ec75bab2c5311b486 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 18 May 2026 11:02:21 +0200 Subject: [PATCH] feat(ml-alpha): PerceptionTrainer wires C21+C22 fwd+bwd+AdamW (C25) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- crates/ml-alpha/build.rs | 1 + .../ml-alpha/cuda/per_horizon_prob_blend.cu | 104 ++++++++++ .../ml-alpha/src/trainer/per_horizon_state.rs | 192 +++++++++++++++++- crates/ml-alpha/src/trainer/perception.rs | 38 ++++ 4 files changed, 334 insertions(+), 1 deletion(-) create mode 100644 crates/ml-alpha/cuda/per_horizon_prob_blend.cu diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 848541e35..0c5fbddd1 100644 --- a/crates/ml-alpha/build.rs +++ b/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 ]; diff --git a/crates/ml-alpha/cuda/per_horizon_prob_blend.cu b/crates/ml-alpha/cuda/per_horizon_prob_blend.cu new file mode 100644 index 000000000..e40659f93 --- /dev/null +++ b/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; +} diff --git a/crates/ml-alpha/src/trainer/per_horizon_state.rs b/crates/ml-alpha/src/trainer/per_horizon_state.rs index 6ade7ce87..e8e114811 100644 --- a/crates/ml-alpha/src/trainer/per_horizon_state.rs +++ b/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, + + // C25 prob-blend kernel handles + per-K scratch (baseline probs stash + + // ∂L/∂r_contrib per K position). + _prob_blend_module: Arc, + 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, } 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::(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, + logit_per_k: &CudaSlice, + probs_per_k: &mut CudaSlice, + ) -> 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, + grad_probs_per_k: &CudaSlice, + ln_b_out: &CudaSlice, + grad_ln_b_out: &mut CudaSlice, + ) -> 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; diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index fde699fcf..9771dc4ca 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/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