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

//! P5 regression: CUDA Graph capture of `PerceptionTrainer::forward_only`.
//!
//! The captured-graph replay (3rd call) must agree with the eager
//! warmup dispatch (1st call) within 1e-5 relative tolerance per
//! element. Strict bit-identity is NOT asserted because CUDA Graph
//! capture can reorder kernel launches and flip f32 reductions by
//! 1 ULP harmlessly — that's why the spec specifies tolerance, not
//! equality.
//!
//! Runtime invocation (GPU required):
//!
//!     SQLX_OFFLINE=true cargo test -p ml-alpha \
//!         --test forward_graph_capture -- --ignored --nocapture
//!
//! The same `PerceptionTrainer` instance is used for all three calls;
//! the three-state machine inside `forward_only` is:
//!   call 1: eager warmup (sets `forward_warmed = true`)
//!   call 2: begin_capture → dispatch → end_capture (stores graph)
//!   call 3: graph.launch() — captured replay
//!
//! Asserting call 1 vs call 3 ensures the captured kernel chain is
//! semantically equivalent to the eager dispatch.

use anyhow::{Context, Result};
use ml_alpha::cfc::snap_features::{Mbp10RawInput, REGIME_DIM};
use ml_alpha::heads::N_HORIZONS;
use ml_alpha::trainer::perception::{PerceptionTrainer, PerceptionTrainerConfig};
use ml_core::device::MlDevice;

const REL_TOL: f32 = 1e-5;

/// Deterministic synthetic window: bid/ask prices decrement by 0.01 per
/// step, fixed sizes, monotonically increasing ts_ns. Same data on each
/// call — but the per-snapshot index is encoded into prices/timestamps
/// so a "shadow position 0 everywhere" bug would still produce
/// distinguishable output across positions.
fn make_test_window(n: usize) -> Vec<Mbp10RawInput> {
    let mut prev_ts_ns: u64 = 1_000_000_000;
    let mut prev_mid: f32 = 100.0;
    (0..n)
        .map(|i| {
            let mid = 100.0 - (i as f32) * 0.01;
            let ts_ns = prev_ts_ns + 1_000_000;
            let bid_px: [f32; 10] = std::array::from_fn(|j| mid - 0.125 - (j as f32) * 0.25);
            let ask_px: [f32; 10] = std::array::from_fn(|j| mid + 0.125 + (j as f32) * 0.25);
            let bid_sz: [f32; 10] = [10.0; 10];
            let ask_sz: [f32; 10] = [10.0; 10];
            let regime: [f32; REGIME_DIM] = [0.1; REGIME_DIM];
            let snap = Mbp10RawInput {
                bid_px,
                bid_sz,
                ask_px,
                ask_sz,
                prev_mid,
                trade_signed_vol: 0.0,
                trade_count: 0,
                ts_ns,
                prev_ts_ns,
                regime,
            };
            prev_mid = mid;
            prev_ts_ns = ts_ns;
            snap
        })
        .collect()
}

#[test]
#[ignore = "requires CUDA"]
fn forward_captured_matches_uncaptured() -> Result<()> {
    let dev = match MlDevice::cuda(0) {
        Ok(d) => d,
        Err(e) => {
            eprintln!("skipping: cuda device unavailable ({e})");
            return Ok(());
        }
    };

    let cfg = PerceptionTrainerConfig::default();
    let total = cfg.seq_len * cfg.n_batch;
    let window = make_test_window(total);

    let mut trainer = PerceptionTrainer::new(&dev, &cfg)
        .context("build PerceptionTrainer")?;

    // Call 1: eager warmup (no capture, no replay).
    let probs_eager = trainer
        .forward_only(&window)
        .context("forward_only call 1 (warmup)")?;

    // Call 2: capture pass (begin_capture → dispatch → end_capture).
    // The dispatch executes once; output is also valid here but not asserted.
    let _probs_capture = trainer
        .forward_only(&window)
        .context("forward_only call 2 (capture)")?;

    // Call 3: replay the captured graph.
    let probs_replay = trainer
        .forward_only(&window)
        .context("forward_only call 3 (replay)")?;

    assert_eq!(
        probs_eager.len(),
        probs_replay.len(),
        "probs length mismatch: eager={} replay={}",
        probs_eager.len(),
        probs_replay.len()
    );
    let expected_len = cfg.seq_len * cfg.n_batch * N_HORIZONS;
    assert_eq!(probs_eager.len(), expected_len, "probs len != K*B*N_HORIZONS");

    let mut max_rel: f32 = 0.0;
    let mut worst_idx: usize = 0;
    for (i, (a, b)) in probs_eager.iter().zip(probs_replay.iter()).enumerate() {
        let denom = a.abs().max(1e-9);
        let rel = (a - b).abs() / denom;
        if rel > max_rel {
            max_rel = rel;
            worst_idx = i;
        }
        assert!(
            rel < REL_TOL,
            "probs[{i}] eager={a} replay={b} rel_err={rel:e} > {REL_TOL:e}"
        );
    }
    eprintln!(
        "forward_captured_matches_uncaptured: {} elements, max rel_err = {:e} at index {}",
        probs_eager.len(),
        max_rel,
        worst_idx
    );
    Ok(())
}