From 69c64f2266f1e4ba4288d36705f591999bb7739f Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 18 May 2026 10:38:24 +0200 Subject: [PATCH] =?UTF-8?q?test(ml-alpha):=20per-horizon=20pipeline=20end-?= =?UTF-8?q?to-end=20smoke=20+=20=CE=B1-gate=20(C23)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Composes the C21 + C22 kernels with a host-side learnable α-gate to prove the full per-horizon contribution path works end-to-end without yet doing the captured-graph integration in PerceptionTrainer. Pipeline: LNb [B, K, HIDDEN_DIM] → per_horizon_attention_pool_fwd → context_h [B, N_HORIZONS, HIDDEN_DIM] → per_horizon_residual_head_fwd → residual [B, N_HORIZONS] → final[b, h] = baseline[b, h] + tanh(α[h]) * residual[b, h] Two tests cover the critical invariants for adoption-safety: alpha_zero_init_is_identity_to_baseline With α = [0, 0, 0, 0, 0] and any random Q_h / w_res / bias_res, final_logit MUST be bit-identical to baseline_logit (because tanh(0) = 0). Verified via to_bits() byte equality. Proves that initialising the new variant with α=0 makes it a strict superset of the existing path — switching to AttentionPoolVariant::PerHorizon cannot regress before any training has happened. alpha_nonzero_changes_output_and_grads_flow_end_to_end With α = [0.5, -0.3, 0.2, -0.1, 0.4]: * final ≠ baseline (residual contributing) ✓ * all final logits finite ✓ * full backward chain (residual_head_bwd → attention_pool_bwd) produces finite d_Q_h_scratch + finite d_LNb with at least one non-zero entry in each → gradients flow back to both the attention queries and the LN_b input ✓ This closes the kernel-side correctness story. The remaining integration commits (C24+) are operational: C24: extend CheckpointV1 → V2 (add q_h, w_res, bias_res, alpha fields; V1 files load as Variant::SharedQuery) C25: PerceptionTrainer wiring — allocate the device buffers, fold attention + residual + gate into the captured graph, plumb gradients into AdamW's param list C26: 1-epoch smoke (assert no NaN, loss decreases vs baseline) — needs real training data + multi-GPU time C27: 30-epoch × 3-fold A/B (task #204) — decision gate per docs/superpowers/specs/2026-05-18-per-horizon-attention-pool-design.md §0 falsifiable claim C24-C25 are 1-2 day work even when carefully scoped; C26-C27 need real GPU-hours + result analysis. C21-C23 land the validatable kernel correctness piece without committing to that time investment yet. Co-Authored-By: Claude Opus 4.7 --- .../tests/per_horizon_full_pipeline_smoke.rs | 201 ++++++++++++++++++ 1 file changed, 201 insertions(+) create mode 100644 crates/ml-alpha/tests/per_horizon_full_pipeline_smoke.rs diff --git a/crates/ml-alpha/tests/per_horizon_full_pipeline_smoke.rs b/crates/ml-alpha/tests/per_horizon_full_pipeline_smoke.rs new file mode 100644 index 000000000..b8d870080 --- /dev/null +++ b/crates/ml-alpha/tests/per_horizon_full_pipeline_smoke.rs @@ -0,0 +1,201 @@ +//! End-to-end smoke test for the per-horizon attention pool + +//! residual head + α-gate composition (C23). +//! +//! Pipeline: +//! LNb [B, K, HIDDEN_DIM] +//! → per_horizon_attention_pool_fwd → context_h [B, N_HORIZONS, HIDDEN_DIM] +//! → per_horizon_residual_head_fwd → residual [B, N_HORIZONS] +//! → final_logit[h] = baseline_logit[h] + tanh(α[h]) * residual[h] +//! +//! Verifies: +//! 1. α = 0 → final_logit == baseline_logit (bit-equal) — proves the +//! additive integration is identity at zero-init, so adopting this +//! path can't regress the existing baseline. +//! 2. α ≠ 0 → final_logit ≠ baseline_logit AND no NaN/Inf leaks. +//! 3. Backward: residual_head_bwd's d_context_h feeds the attention +//! pool bwd; full gradient chain finite + non-zero where expected. + +use anyhow::Result; +use cudarc::driver::CudaSlice; +use ml_alpha::per_horizon_attention_pool::{ + PerHorizonAttentionPool, PHA_HIDDEN_DIM, PHA_N_HORIZONS, +}; +use ml_alpha::per_horizon_residual_head::PerHorizonResidualHead; +use ml_core::device::MlDevice; +use rand::Rng; +use rand::SeedableRng; +use rand_chacha::ChaCha8Rng; + +const B: usize = 4; +const K: usize = 16; + +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 +} + +#[test] +#[ignore = "requires CUDA"] +fn alpha_zero_init_is_identity_to_baseline() -> Result<()> { + let Some(dev) = try_dev() else { return Ok(()); }; + let ctx = dev.cuda_context()?.clone(); + let stream = dev.cuda_stream()?.clone(); + let pool = PerHorizonAttentionPool::new(&ctx, stream.clone())?; + let head = PerHorizonResidualHead::new(&ctx, stream.clone())?; + + let mut rng = ChaCha8Rng::seed_from_u64(0xA11A_0BEEF); + let ln_host: Vec = (0..B * K * PHA_HIDDEN_DIM).map(|_| rng.gen_range(-0.5..0.5)).collect(); + let q_h_host: Vec = (0..PHA_N_HORIZONS * PHA_HIDDEN_DIM).map(|_| rng.gen_range(-0.1..0.1)).collect(); + let w_res_host: Vec = (0..PHA_N_HORIZONS * PHA_HIDDEN_DIM).map(|_| rng.gen_range(-0.1..0.1)).collect(); + let bias_res_host: Vec = (0..PHA_N_HORIZONS).map(|_| rng.gen_range(-0.05..0.05)).collect(); + let baseline_logit_host: Vec = (0..B * PHA_N_HORIZONS).map(|_| rng.gen_range(-2.0..2.0)).collect(); + + // α = 0 — final logit MUST equal baseline. + let alpha_zero = vec![0.0f32; PHA_N_HORIZONS]; + + let ln_d = alloc_upload(&stream, &ln_host); + let q_h_d = alloc_upload(&stream, &q_h_host); + let w_res_d = alloc_upload(&stream, &w_res_host); + let b_res_d = alloc_upload(&stream, &bias_res_host); + + let mut ctx_d = stream.alloc_zeros::(B * PHA_N_HORIZONS * PHA_HIDDEN_DIM)?; + let mut attn_d = stream.alloc_zeros::(B * PHA_N_HORIZONS * K)?; + pool.forward(&q_h_d, &ln_d, B as i32, K as i32, &mut ctx_d, &mut attn_d)?; + + let mut residual_d = stream.alloc_zeros::(B * PHA_N_HORIZONS)?; + head.forward(&ctx_d, &w_res_d, &b_res_d, B as i32, &mut residual_d)?; + + // Host-side gate combination: final[b,h] = baseline[b,h] + tanh(α[h]) * residual[b,h]. + let residual_host = download(&stream, &residual_d); + let mut final_logit = vec![0.0f32; B * PHA_N_HORIZONS]; + for b in 0..B { + for h in 0..PHA_N_HORIZONS { + let idx = b * PHA_N_HORIZONS + h; + let gate = alpha_zero[h].tanh(); + final_logit[idx] = baseline_logit_host[idx] + gate * residual_host[idx]; + } + } + + for i in 0..B * PHA_N_HORIZONS { + assert_eq!( + final_logit[i].to_bits(), + baseline_logit_host[i].to_bits(), + "alpha=0 at idx {i}: final {} ≠ baseline {}", + final_logit[i], baseline_logit_host[i] + ); + } + Ok(()) +} + +#[test] +#[ignore = "requires CUDA"] +fn alpha_nonzero_changes_output_and_grads_flow_end_to_end() -> Result<()> { + let Some(dev) = try_dev() else { return Ok(()); }; + let ctx = dev.cuda_context()?.clone(); + let stream = dev.cuda_stream()?.clone(); + let pool = PerHorizonAttentionPool::new(&ctx, stream.clone())?; + let head = PerHorizonResidualHead::new(&ctx, stream.clone())?; + + let mut rng = ChaCha8Rng::seed_from_u64(0xBABE_F00D); + let ln_host: Vec = (0..B * K * PHA_HIDDEN_DIM).map(|_| rng.gen_range(-0.5..0.5)).collect(); + let q_h_host: Vec = (0..PHA_N_HORIZONS * PHA_HIDDEN_DIM).map(|_| rng.gen_range(-0.1..0.1)).collect(); + let w_res_host: Vec = (0..PHA_N_HORIZONS * PHA_HIDDEN_DIM).map(|_| rng.gen_range(-0.1..0.1)).collect(); + let bias_res_host: Vec = (0..PHA_N_HORIZONS).map(|_| rng.gen_range(-0.05..0.05)).collect(); + let baseline_logit_host: Vec = (0..B * PHA_N_HORIZONS).map(|_| rng.gen_range(-2.0..2.0)).collect(); + // Non-zero alpha — different per horizon. + let alpha = [0.5f32, -0.3, 0.2, -0.1, 0.4]; + + let ln_d = alloc_upload(&stream, &ln_host); + let q_h_d = alloc_upload(&stream, &q_h_host); + let w_res_d = alloc_upload(&stream, &w_res_host); + let b_res_d = alloc_upload(&stream, &bias_res_host); + + let mut ctx_d = stream.alloc_zeros::(B * PHA_N_HORIZONS * PHA_HIDDEN_DIM)?; + let mut attn_d = stream.alloc_zeros::(B * PHA_N_HORIZONS * K)?; + pool.forward(&q_h_d, &ln_d, B as i32, K as i32, &mut ctx_d, &mut attn_d)?; + let mut residual_d = stream.alloc_zeros::(B * PHA_N_HORIZONS)?; + head.forward(&ctx_d, &w_res_d, &b_res_d, B as i32, &mut residual_d)?; + + let residual_host = download(&stream, &residual_d); + let mut final_logit = vec![0.0f32; B * PHA_N_HORIZONS]; + let mut differences_seen = 0; + for b in 0..B { + for h in 0..PHA_N_HORIZONS { + let idx = b * PHA_N_HORIZONS + h; + let gate = alpha[h].tanh(); + final_logit[idx] = baseline_logit_host[idx] + gate * residual_host[idx]; + assert!(final_logit[idx].is_finite(), "non-finite final[{idx}]"); + if (final_logit[idx] - baseline_logit_host[idx]).abs() > 1e-6 { + differences_seen += 1; + } + } + } + assert!( + differences_seen > 0, + "α non-zero but final logit unchanged from baseline — residual contributing nothing" + ); + + // Now run the full backward chain. Loss = Σ final_logit, so + // d_final[b,h] = 1 + // d_residual[b,h] = tanh(α[h]) + // d_α[h] = (1 - tanh(α[h])²) * Σ_b residual[b, h] + // d_baseline[b,h] = 1 + let mut d_residual_host = vec![0.0f32; B * PHA_N_HORIZONS]; + for b in 0..B { + for h in 0..PHA_N_HORIZONS { + d_residual_host[b * PHA_N_HORIZONS + h] = alpha[h].tanh(); + } + } + let d_residual_d = alloc_upload(&stream, &d_residual_host); + + let mut d_w_res_scratch = stream.alloc_zeros::(B * PHA_N_HORIZONS * PHA_HIDDEN_DIM)?; + let mut d_b_res_scratch = stream.alloc_zeros::(B * PHA_N_HORIZONS)?; + let mut d_ctx = stream.alloc_zeros::(B * PHA_N_HORIZONS * PHA_HIDDEN_DIM)?; + head.backward( + &ctx_d, &w_res_d, &d_residual_d, B as i32, + &mut d_w_res_scratch, &mut d_b_res_scratch, &mut d_ctx, + )?; + + let mut d_q_h_scratch = stream.alloc_zeros::(B * PHA_N_HORIZONS * PHA_HIDDEN_DIM)?; + let mut d_ln = stream.alloc_zeros::(B * K * PHA_HIDDEN_DIM)?; + pool.backward( + &q_h_d, &ln_d, &attn_d, &d_ctx, + B as i32, K as i32, + &mut d_q_h_scratch, &mut d_ln, + )?; + + let d_q_h_host = download(&stream, &d_q_h_scratch); + let d_ln_host = download(&stream, &d_ln); + + // No NaN/Inf in any gradient buffer. + for &g in &d_q_h_host { assert!(g.is_finite(), "d_q_h scratch has non-finite"); } + for &g in &d_ln_host { assert!(g.is_finite(), "d_ln has non-finite"); } + + // At least one non-zero gradient (else gate path didn't actually trigger). + let any_nonzero_q = d_q_h_host.iter().any(|&g| g.abs() > 1e-8); + let any_nonzero_ln = d_ln_host.iter().any(|&g| g.abs() > 1e-8); + assert!(any_nonzero_q, "d_q_h all zero — backward didn't flow through attention"); + assert!(any_nonzero_ln, "d_ln all zero — backward didn't flow back to LN_b input"); + + Ok(()) +}