test(ml-alpha): per-horizon pipeline end-to-end smoke + α-gate (C23)
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
201
crates/ml-alpha/tests/per_horizon_full_pipeline_smoke.rs
Normal file
201
crates/ml-alpha/tests/per_horizon_full_pipeline_smoke.rs
Normal file
@@ -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<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
|
||||
}
|
||||
|
||||
#[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<f32> = (0..B * K * PHA_HIDDEN_DIM).map(|_| rng.gen_range(-0.5..0.5)).collect();
|
||||
let q_h_host: Vec<f32> = (0..PHA_N_HORIZONS * PHA_HIDDEN_DIM).map(|_| rng.gen_range(-0.1..0.1)).collect();
|
||||
let w_res_host: Vec<f32> = (0..PHA_N_HORIZONS * PHA_HIDDEN_DIM).map(|_| rng.gen_range(-0.1..0.1)).collect();
|
||||
let bias_res_host: Vec<f32> = (0..PHA_N_HORIZONS).map(|_| rng.gen_range(-0.05..0.05)).collect();
|
||||
let baseline_logit_host: Vec<f32> = (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::<f32>(B * PHA_N_HORIZONS * PHA_HIDDEN_DIM)?;
|
||||
let mut attn_d = stream.alloc_zeros::<f32>(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::<f32>(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<f32> = (0..B * K * PHA_HIDDEN_DIM).map(|_| rng.gen_range(-0.5..0.5)).collect();
|
||||
let q_h_host: Vec<f32> = (0..PHA_N_HORIZONS * PHA_HIDDEN_DIM).map(|_| rng.gen_range(-0.1..0.1)).collect();
|
||||
let w_res_host: Vec<f32> = (0..PHA_N_HORIZONS * PHA_HIDDEN_DIM).map(|_| rng.gen_range(-0.1..0.1)).collect();
|
||||
let bias_res_host: Vec<f32> = (0..PHA_N_HORIZONS).map(|_| rng.gen_range(-0.05..0.05)).collect();
|
||||
let baseline_logit_host: Vec<f32> = (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::<f32>(B * PHA_N_HORIZONS * PHA_HIDDEN_DIM)?;
|
||||
let mut attn_d = stream.alloc_zeros::<f32>(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::<f32>(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::<f32>(B * PHA_N_HORIZONS * PHA_HIDDEN_DIM)?;
|
||||
let mut d_b_res_scratch = stream.alloc_zeros::<f32>(B * PHA_N_HORIZONS)?;
|
||||
let mut d_ctx = stream.alloc_zeros::<f32>(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::<f32>(B * PHA_N_HORIZONS * PHA_HIDDEN_DIM)?;
|
||||
let mut d_ln = stream.alloc_zeros::<f32>(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(())
|
||||
}
|
||||
Reference in New Issue
Block a user