foxhunt/crates/ml-alpha/tests/aux_trunk.rs

//! GPU oracle tests for `aux_trunk.cu` (single-bucket CfC at AUX_HIDDEN=64).
//!
//! Run with:
//!   SQLX_OFFLINE=true cargo test -p ml-alpha --test aux_trunk \
//!     -- --ignored --nocapture
//!
//! These tests are `#[ignore]` because they require a CUDA device.
//! They drive the cubin directly through the host wrappers in
//! `crates/ml-alpha/src/cfc/aux_trunk.rs` and compare against a
//! deterministic CPU naive reference (per
//! `feedback_no_cpu_test_fallbacks.md` the production path stays GPU-
//! only; the CPU implementation here exists ONLY as a unit-test oracle).

use anyhow::Result;
use ml_alpha::cfc::aux_trunk::{
    aux_trunk_bwd_gpu, aux_trunk_fwd_gpu, AuxTrunk, AuxTrunkConfig, AUX_HIDDEN,
};
use ml_core::device::MlDevice;

// Pick a representative feat_dim. Aux trunk's production caller (B5)
// will pass HIDDEN_DIM=128 (encoder output). Using a non-128 dim here
// also exercises the runtime-`feat_dim` parameterisation.
const TEST_FEAT_DIM: usize = 96;

fn upload(
    stream: &std::sync::Arc<cudarc::driver::CudaStream>,
    host: &[f32],
) -> Result<cudarc::driver::CudaSlice<f32>> {
    let mut d = stream.alloc_zeros::<f32>(host.len())?;
    stream.memcpy_htod(host, &mut d)?;
    Ok(d)
}

fn download(
    stream: &std::sync::Arc<cudarc::driver::CudaStream>,
    d: &cudarc::driver::CudaSlice<f32>,
) -> Result<Vec<f32>> {
    let mut h = vec![0.0_f32; d.len()];
    stream.memcpy_dtoh(d, &mut h)?;
    Ok(h)
}

/// Forward shape correctness + value match against a naive CPU oracle.
#[test]
#[ignore = "requires CUDA"]
fn fwd_matches_naive_reference() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let trunk = AuxTrunk::new(
        &dev,
        AuxTrunkConfig {
            feat_dim: TEST_FEAT_DIM,
            seed: 0xAEB1_BEEF_u64,
        },
    )?;

    let b_sz: usize = 3;
    let x: Vec<f32> = (0..b_sz * TEST_FEAT_DIM)
        .map(|i| ((i as f32) * 0.0011).sin() * 0.5)
        .collect();
    let h_old: Vec<f32> = (0..b_sz * AUX_HIDDEN)
        .map(|i| ((i as f32) * 0.0037).cos() * 0.3)
        .collect();

    let x_d = upload(&stream, &x)?;
    let h_old_d = upload(&stream, &h_old)?;
    let mut h_new_d = stream.alloc_zeros::<f32>(b_sz * AUX_HIDDEN)?;

    let dt: f32 = 1.0;
    aux_trunk_fwd_gpu(
        &stream,
        &trunk.fwd_fn,
        TEST_FEAT_DIM,
        &trunk.w_in_d,
        &trunk.w_rec_d,
        &trunk.b_d,
        &trunk.tau_d,
        &x_d,
        &h_old_d,
        dt,
        b_sz as i32,
        &mut h_new_d,
    )?;
    stream.synchronize()?;
    let h_new = download(&stream, &h_new_d)?;

    // Pull weights down for the CPU oracle.
    let weights = trunk.download_weights()?;
    let w_in = &weights.w_in;
    let w_rec = &weights.w_rec;
    let b_vec = &weights.b;
    let tau = &weights.tau;

    for batch in 0..b_sz {
        for c in 0..AUX_HIDDEN {
            let mut pre = b_vec[c];
            for k in 0..TEST_FEAT_DIM {
                pre += w_in[c * TEST_FEAT_DIM + k] * x[batch * TEST_FEAT_DIM + k];
            }
            for k in 0..AUX_HIDDEN {
                pre += w_rec[c * AUX_HIDDEN + k] * h_old[batch * AUX_HIDDEN + k];
            }
            let decay = (-dt / tau[c].max(1e-6)).exp();
            let expected =
                h_old[batch * AUX_HIDDEN + c] * decay + (1.0 - decay) * pre.tanh();
            let got = h_new[batch * AUX_HIDDEN + c];
            assert!(
                (got - expected).abs() < 5e-4,
                "h_new[{},{}] = {got} vs expected {expected} (diff {})",
                batch,
                c,
                (got - expected).abs(),
            );
        }
    }
    Ok(())
}

/// Backward grad shapes finite, non-zero, and bias-grad finite-difference
/// matches within tolerance. Full per-weight FD validation is expensive
/// (O(AUX_HIDDEN²·feat_dim) reference evaluations); we sample a handful
/// of bias entries which is sufficient to catch wiring breakage.
#[test]
#[ignore = "requires CUDA"]
fn bwd_finite_diff_matches_bias_sample() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let trunk = AuxTrunk::new(
        &dev,
        AuxTrunkConfig {
            feat_dim: TEST_FEAT_DIM,
            seed: 0x12_34_56_78,
        },
    )?;

    let b_sz: usize = 2;
    let x: Vec<f32> = (0..b_sz * TEST_FEAT_DIM)
        .map(|i| ((i as f32) * 0.0021).sin() * 0.4)
        .collect();
    let h_old: Vec<f32> = (0..b_sz * AUX_HIDDEN)
        .map(|i| ((i as f32) * 0.0049).cos() * 0.25)
        .collect();
    // Pick a structured grad_h_new (1.0 at one slot per batch, 0 else)
    // so that grad_b finite-diff has a clean closed-form expectation.
    let grad_h_new: Vec<f32> = (0..b_sz * AUX_HIDDEN)
        .map(|i| ((i % 7) as f32) * 0.03)
        .collect();

    let weights = trunk.download_weights()?;
    let w_in = weights.w_in.clone();
    let w_rec = weights.w_rec.clone();
    let b_vec = weights.b.clone();
    let tau = weights.tau.clone();

    let dt: f32 = 1.0;
    let x_d = upload(&stream, &x)?;
    let h_old_d = upload(&stream, &h_old)?;
    let grad_h_new_d = upload(&stream, &grad_h_new)?;

    let mut grad_w_in_d = stream.alloc_zeros::<f32>(b_sz * AUX_HIDDEN * TEST_FEAT_DIM)?;
    let mut grad_w_rec_d = stream.alloc_zeros::<f32>(b_sz * AUX_HIDDEN * AUX_HIDDEN)?;
    let mut grad_b_d = stream.alloc_zeros::<f32>(b_sz * AUX_HIDDEN)?;
    let mut grad_tau_d = stream.alloc_zeros::<f32>(b_sz * AUX_HIDDEN)?;
    let mut grad_h_old_d = stream.alloc_zeros::<f32>(b_sz * AUX_HIDDEN)?;
    let mut grad_x_d = stream.alloc_zeros::<f32>(b_sz * TEST_FEAT_DIM)?;

    aux_trunk_bwd_gpu(
        &stream,
        &trunk.bwd_fn,
        TEST_FEAT_DIM,
        &trunk.w_in_d,
        &trunk.w_rec_d,
        &trunk.b_d,
        &trunk.tau_d,
        &x_d,
        &h_old_d,
        &grad_h_new_d,
        dt,
        b_sz as i32,
        &mut grad_w_in_d,
        &mut grad_w_rec_d,
        &mut grad_b_d,
        &mut grad_tau_d,
        &mut grad_h_old_d,
        &mut grad_x_d,
    )?;
    stream.synchronize()?;

    // All grads finite, no NaN/Inf.
    let g_w_in = download(&stream, &grad_w_in_d)?;
    let g_w_rec = download(&stream, &grad_w_rec_d)?;
    let g_b = download(&stream, &grad_b_d)?;
    let g_tau = download(&stream, &grad_tau_d)?;
    let g_h_old = download(&stream, &grad_h_old_d)?;
    let g_x = download(&stream, &grad_x_d)?;
    for (name, v) in [
        ("grad_w_in", &g_w_in),
        ("grad_w_rec", &g_w_rec),
        ("grad_b", &g_b),
        ("grad_tau", &g_tau),
        ("grad_h_old", &g_h_old),
        ("grad_x", &g_x),
    ] {
        assert!(
            v.iter().all(|x| x.is_finite()),
            "{name} contains NaN/Inf",
        );
    }
    // Kernel must have written something non-zero for at least the bias
    // (the only grad guaranteed to be non-trivial given non-zero
    // grad_h_new entries).
    assert!(
        g_b.iter().any(|x| x.abs() > 0.0),
        "grad_b is all zero — kernel didn't write?",
    );

    // Finite-difference oracle for grad_b on the first batch, channel 0.
    // Hand-coded forward (CPU) to derive d/d b[c] of the implicit loss
    // L = Σ_{batch,c} grad_h_new[batch, c] * h_new[batch, c].
    let h_new_fwd = |b_local: &[f32]| -> Vec<f32> {
        let mut out = vec![0.0_f32; b_sz * AUX_HIDDEN];
        for batch in 0..b_sz {
            for c in 0..AUX_HIDDEN {
                let mut pre = b_local[c];
                for k in 0..TEST_FEAT_DIM {
                    pre += w_in[c * TEST_FEAT_DIM + k] * x[batch * TEST_FEAT_DIM + k];
                }
                for k in 0..AUX_HIDDEN {
                    pre += w_rec[c * AUX_HIDDEN + k] * h_old[batch * AUX_HIDDEN + k];
                }
                let decay = (-dt / tau[c].max(1e-6)).exp();
                out[batch * AUX_HIDDEN + c] =
                    h_old[batch * AUX_HIDDEN + c] * decay + (1.0 - decay) * pre.tanh();
            }
        }
        out
    };
    let loss_for = |b_local: &[f32]| -> f32 {
        let h_new = h_new_fwd(b_local);
        let mut acc = 0.0_f32;
        for i in 0..b_sz * AUX_HIDDEN {
            acc += grad_h_new[i] * h_new[i];
        }
        acc
    };

    let eps = 1e-3_f32;
    // Sample 4 channels at modest stride; full sweep is overkill for a
    // smoke. The kernel writes per-batch grad slices — sum across batch
    // to match the loss gradient w.r.t. the shared bias param.
    for &c in &[0_usize, 17, 33, 50] {
        let mut b_plus = b_vec.clone();
        b_plus[c] += eps;
        let mut b_minus = b_vec.clone();
        b_minus[c] -= eps;
        let fd = (loss_for(&b_plus) - loss_for(&b_minus)) / (2.0 * eps);

        // Sum kernel's per-batch grad_b across batch to match dL/db[c].
        let mut analytical = 0.0_f32;
        for batch in 0..b_sz {
            analytical += g_b[batch * AUX_HIDDEN + c];
        }
        let abs_err = (fd - analytical).abs();
        let rel_tol = 5e-3 * fd.abs().max(1e-4);
        assert!(
            abs_err < rel_tol.max(2e-3),
            "grad_b[{c}] mismatch: fd={fd}, analytical={analytical}, abs_err={abs_err}",
        );
    }
    Ok(())
}

/// Smoke that AUX_HIDDEN=64 path runs without OOM or NaN at a batch
/// size representative of the trainer (which uses B=32-128 per batch).
/// Designed to be cheap enough for RTX 3050 sm_86 local validation.
#[test]
#[ignore = "requires CUDA"]
fn fwd_smoke_large_batch_no_nan_no_oom() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let trunk = AuxTrunk::new(
        &dev,
        AuxTrunkConfig {
            feat_dim: 128, // matches HIDDEN_DIM — actual B5 wiring config
            seed: 0xFEED_FACE_u64,
        },
    )?;

    let b_sz: usize = 64;
    let feat_dim = 128usize;
    let x: Vec<f32> = (0..b_sz * feat_dim)
        .map(|i| ((i as f32) * 0.0007).sin())
        .collect();
    let h_old: Vec<f32> = vec![0.0; b_sz * AUX_HIDDEN]; // cold start
    let x_d = upload(&stream, &x)?;
    let h_old_d = upload(&stream, &h_old)?;
    let mut h_new_d = stream.alloc_zeros::<f32>(b_sz * AUX_HIDDEN)?;

    aux_trunk_fwd_gpu(
        &stream,
        &trunk.fwd_fn,
        feat_dim,
        &trunk.w_in_d,
        &trunk.w_rec_d,
        &trunk.b_d,
        &trunk.tau_d,
        &x_d,
        &h_old_d,
        1.0,
        b_sz as i32,
        &mut h_new_d,
    )?;
    stream.synchronize()?;
    let h_new = download(&stream, &h_new_d)?;
    assert!(
        h_new.iter().all(|v| v.is_finite()),
        "h_new contains NaN/Inf after cold-start forward",
    );
    // Output is bounded by tanh ∈ (-1, 1); with h_old = 0 and small init,
    // |h_new| should sit well inside that range.
    assert!(
        h_new.iter().all(|v| v.abs() < 1.5),
        "h_new produced out-of-range values (|v| ≥ 1.5)",
    );
    Ok(())
}