diff --git a/Cargo.lock b/Cargo.lock index cbf1dfca5..6f2694c4c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4083,6 +4083,15 @@ dependencies = [ "tracing-subscriber", ] +[[package]] +name = "fxt-data-audit" +version = "1.0.0" +dependencies = [ + "anyhow", + "data", + "ml-features", +] + [[package]] name = "generic-array" version = "0.14.7" diff --git a/crates/ml-alpha/src/mamba2_block.rs b/crates/ml-alpha/src/mamba2_block.rs index f58a045a6..7619bb980 100644 --- a/crates/ml-alpha/src/mamba2_block.rs +++ b/crates/ml-alpha/src/mamba2_block.rs @@ -1306,6 +1306,67 @@ impl Mamba2Block { Ok(()) } + /// CRT Phase A0.5 (corrective) — bit-identical seeding helper. + /// Advances `scratch.x_state` by one timestep reading per-step + /// `a_proj` / `b_proj` row pointers from an EXTERNAL seq-path scratch + /// (typically the trainer's `mamba2_fwd_scratch.a_proj/b_proj` after + /// a `forward_only` call). Same scan kernel as `step_into`, but skips + /// the `w_in`/`w_a`/`w_b` GEMMs because the caller already has those + /// projections from a one-shot batched cuBLAS call. + /// + /// This is the key to bit-identity vs forward_only: the seq path's + /// W_in→W_a/W_b GEMMs are computed once over all K rows in one cuBLAS + /// call. Re-running them per-row through `step_into` would not match + /// bit-for-bit (cuBLAS picks different algorithms for different + /// matrix shapes). By replaying scan_fwd_step over the seq path's + /// already-computed a_proj/b_proj, the x_state evolution is bit- + /// identical to scan_fwd_seq's internal register-state trajectory. + /// + /// Args: + /// - `a_proj_row_ptr` / `b_proj_row_ptr`: device pointers into + /// `[N=n_batch, state_d]` slices of the seq scratch at the k-th row + /// (caller computes the byte offset). + /// - `scratch`: step scratch; `x_state` is read + written in place. + pub fn step_advance_from_seq_row( + &self, + a_proj_row_ptr: u64, + b_proj_row_ptr: u64, + scratch: &mut Mamba2BlockStepScratch, + ) -> Result<()> { + let c = &self.config; + anyhow::ensure!( + scratch.hidden_dim == c.hidden_dim && scratch.state_dim == c.state_dim, + "step scratch shape mismatch: hidden_dim {} vs {}, state_dim {} vs {}", + scratch.hidden_dim, c.hidden_dim, scratch.state_dim, c.state_dim + ); + let n_batch = scratch.n_batch; + let block_threads: u32 = 32; + let grid_y: u32 = + ((c.hidden_dim + block_threads as usize - 1) / block_threads as usize) as u32; + let cfg = LaunchConfig { + grid_dim: (n_batch as u32, grid_y, 1), + block_dim: (block_threads, 1, 1), + shared_mem_bytes: 0, + }; + let n_i32 = n_batch as i32; + let sh2_i32 = c.hidden_dim as i32; + let st_i32 = c.state_dim as i32; + unsafe { + self.stream + .launch_builder(&self.kernel_fwd_step) + .arg(&mut scratch.x_state) + .arg(&a_proj_row_ptr) + .arg(&b_proj_row_ptr) + .arg(&self.w_c) + .arg(scratch.h_s2.cuda_data()) + .arg(scratch.h_out.data_mut()) + .arg(&n_i32).arg(&sh2_i32).arg(&st_i32) + .launch(cfg) + .map_err(|e| anyhow!("scan_fwd_step (seed) launch: {e}"))?; + } + Ok(()) + } + /// Backward chain paired with [`forward_train_seq`]. `d_h_enriched_seq` /// has shape `[N, K, hidden_dim]` matching `cache.h_enriched_seq`. /// Returns all nine parameter gradients (`dw_out` / `db_out` zeroed, diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index d3a25b18c..3d3280ff1 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -526,12 +526,6 @@ pub struct PerceptionTrainer { ln_b_step_out_d: CudaSlice, /// Single-row LN_b stats `[1, 2]` for forward_step. ln_b_step_stats_d: CudaSlice, - /// Single-row probs `[1, N_HORIZONS]` for forward_step output. - probs_step_d: CudaSlice, - /// Single-row pinned dtoh target for probs read-back; avoids - /// allocating a MappedF32Buffer per call. Filled by a DtoD into the - /// staging then read on the host after stream sync. - probs_step_host: MappedF32Buffer, /// Per-horizon intermediate scratches `[1, N_HORIZONS, HEAD_MID_DIM]` /// / `[1, N_HORIZONS]` — required by `multi_horizon_heads_grn_fwd_batched` /// signature. Unused at inference but the kernel writes them. @@ -976,10 +970,12 @@ impl PerceptionTrainer { .context("ln_b_step_out_d alloc")?; let ln_b_step_stats_d = stream.alloc_zeros::(cfg.n_batch * 2) .context("ln_b_step_stats_d alloc")?; - let probs_step_d = stream.alloc_zeros::(cfg.n_batch * N_HORIZONS) - .context("probs_step_d alloc")?; - let probs_step_host = unsafe { MappedF32Buffer::new(cfg.n_batch * N_HORIZONS) } - .map_err(|e| anyhow::anyhow!("probs_step_host: {e}"))?; + // probs_step_d / probs_step_host removed: forward_step_into now + // writes probs directly into a caller-supplied device buffer + // (typically `LobSimCuda::alpha_probs_d_mut`), eliminating the + // per-event DtoD-to-host + stream.synchronize that previously + // round-tripped the 5 probs through pinned host memory just to + // immediately re-upload them in `broadcast_alpha`. let z1_step_d = stream.alloc_zeros::( cfg.n_batch * N_HORIZONS * HEAD_MID_DIM).context("z1_step_d alloc")?; let a1_step_d = stream.alloc_zeros::( @@ -1224,8 +1220,6 @@ impl PerceptionTrainer { ln_a_step_stats_d, ln_b_step_out_d, ln_b_step_stats_d, - probs_step_d, - probs_step_host, z1_step_d, a1_step_d, z2_step_d, @@ -2901,7 +2895,7 @@ impl PerceptionTrainer { /// projection per call. The K-history is captured implicitly in /// the persistent SSM and CfC states. /// - /// Sequence semantics: calling forward_step N times on snapshots + /// Sequence semantics: calling forward_step_into N times on snapshots /// (s_0 … s_{N-1}) starting from a fresh `reset_step_state()` is /// equivalent (within stable-SSM dampening) to `forward_only` on the /// same K-window snapshots once the persistent state has accumulated @@ -2910,16 +2904,29 @@ impl PerceptionTrainer { /// (`sigmoid(a) < 1`), state dampens at ~0.5^N — after ~K events the /// state is dominated by the recent K snapshots and agrees with /// forward_only on the same K within float-rounding tolerance. - pub fn forward_step( + /// + /// The output `alpha_probs_dst` MUST be a device slice of length + /// `N_HORIZONS` (B=1 enforced). The GRN heads kernel writes + /// probabilities directly into it — no host staging, no `synchronize` + /// on the hot path. Per feedback_cpu_is_read_only and + /// feedback_no_htod_htoh_only_mapped_pinned: probs stay on device + /// throughout the per-event call chain. + pub fn forward_step_into( &mut self, snapshot: &Mbp10RawInput, - ) -> Result<[f32; N_HORIZONS]> { + alpha_probs_dst: &mut CudaSlice, + ) -> Result<()> { let b_sz = self.cfg.n_batch; anyhow::ensure!( b_sz == 1, - "forward_step currently supports n_batch == 1 only (got {})", + "forward_step_into currently supports n_batch == 1 only (got {})", b_sz ); + anyhow::ensure!( + alpha_probs_dst.len() >= N_HORIZONS, + "forward_step_into: alpha_probs_dst len {} < N_HORIZONS {}", + alpha_probs_dst.len(), N_HORIZONS + ); // ── 1. Host staging fill (1 snapshot). All buffers are // mapped-pinned; the snap_feature kernel reads them after @@ -3147,31 +3154,55 @@ impl PerceptionTrainer { .arg(&self.trunk.heads_w_main_d).arg(&self.trunk.heads_b_main_d) .arg(&self.trunk.heads_w_skip_d).arg(&self.trunk.heads_b_skip_d) .arg(&self.cfc_h_new_step_d).arg(&n_batch_i) - .arg(&mut self.probs_step_d) + // Probs go directly to the caller-supplied device buffer — + // no host staging, no synchronize. The next consumer (the + // sim's decision kernels) reads from this same device slice + // and is launched on the same stream so it is stream-ordered. + .arg(alpha_probs_dst) .arg(&mut self.z1_step_d).arg(&mut self.a1_step_d).arg(&mut self.z2_step_d) .arg(&mut self.gate_logit_step_d).arg(&mut self.main_step_d).arg(&mut self.logit_step_d); launch.launch(cfg_grn_fwd).context("step heads GRN")?; } + // No DtoD-to-host, no stream.synchronize, no host read. The probs + // remain on-device for the consumer (decision kernels). The CPU + // never touches them on the hot path per + // pearl/feedback_cpu_is_read_only. + Ok(()) + } - // ── 11. Sync + DtoH: pull the N_HORIZONS probs back to host. - // Stage to the mapped-pinned probs_step_host buffer, sync, - // then read off the host pointer. One sync per call is the - // cost of an inherently host-consumed output. - unsafe { - let s = self.stream.cu_stream(); - let (src, _g) = self.probs_step_d.device_ptr(&self.stream); - let nbytes = b_sz * N_HORIZONS * std::mem::size_of::(); - cudarc::driver::result::memcpy_dtod_async( - self.probs_step_host.dev_ptr, src, nbytes, s, - ).context("step probs dtod (to host staging)")?; - } - self.stream.synchronize().context("forward_step end-sync")?; - - // Pull the first row's N_HORIZONS values into a fixed-size array. - let host_all = self.probs_step_host.read_all(); + /// Test-only convenience: drive `forward_step_into` against a private + /// device buffer and DtoH-sync the probs into a host array. NOT part + /// of the production hot path — production callers MUST use + /// `forward_step_into(snapshot, &mut alpha_probs_dst)` and keep the + /// probs on-device. This helper exists for integration tests that + /// previously called the old `forward_step(snap) -> [f32; N]` API and + /// want to assert on host values; it allocates a one-shot device + /// buffer, runs forward_step_into into it, then DtoVs once. + /// + /// Per-call DtoV stops the stream; do NOT use on the per-event hot + /// path. The function name carries `_returning` as a visible flag for + /// reviewers that a host roundtrip happens here. + pub fn forward_step_into_returning( + &mut self, + snapshot: &Mbp10RawInput, + ) -> Result<[f32; N_HORIZONS]> { + let b_sz = self.cfg.n_batch; + anyhow::ensure!( + b_sz == 1, + "forward_step_into_returning supports n_batch == 1 only (got {})", + b_sz + ); + let mut probs_d = self.stream + .alloc_zeros::(N_HORIZONS) + .map_err(|e| anyhow::anyhow!("test probs_d alloc: {e}"))?; + self.forward_step_into(snapshot, &mut probs_d)?; + let mut host = vec![0.0_f32; N_HORIZONS]; + self.stream + .memcpy_dtoh(&probs_d, host.as_mut_slice()) + .map_err(|e| anyhow::anyhow!("test probs DtoH: {e}"))?; let mut out = [0.0_f32; N_HORIZONS]; for h in 0..N_HORIZONS { - out[h] = host_all[h]; + out[h] = host[h]; } Ok(out) } @@ -3195,6 +3226,121 @@ impl PerceptionTrainer { Ok(()) } + /// CRT Phase A0.5 corrective (concern #1 from a0e81fbdf): bit-identical + /// seeding from forward_only. + /// + /// After this call, the step-path persistent state ( + /// step_scratch_l1.x_state, + /// step_scratch_l2.x_state, + /// cfc_h_state_step_d + /// ) reproduces the internal state forward_only had at the end of + /// processing `window` (K snapshots). Subsequent `forward_step_into` + /// calls then continue the same trajectory bit-identically to where + /// forward_only left off, modulo: + /// - The cuBLAS W_in/W_a/W_b GEMM bit-identity for the K-batched seq + /// path (one GEMM over [B*K, *]) is what we capture; the scan_fwd_step + /// replay below reads those projections and produces bit-identical + /// x_state by construction (same arithmetic, same per-step order). + /// - The CfC h_state is a DtoD copy of h_new_per_k_d at position K-1 + /// — also bit-identical. + /// + /// Implementation: + /// 1. Run `forward_only(window)` — populates mamba2_fwd_scratch.a_proj/b_proj + /// + mamba2_l2_fwd_scratch.a_proj/b_proj (via batched GEMM over K rows) + /// + h_new_per_k_d (the CfC K-loop output at every position). + /// 2. Zero step scratches so the scan_fwd_step replay starts from x=0, + /// matching scan_fwd_seq's `float x[32] = {0}` initialisation. + /// 3. For each k=0..K, launch `step_advance_from_seq_row` on L1 reading + /// a_proj/b_proj at row k. After K calls, step_scratch_l1.x_state + /// matches scan_fwd_seq's terminal register state for L1. + /// 4. Same for L2. + /// 5. DtoD copy h_new_per_k_d at slot K-1 (i.e., the cfc state that + /// would feed snapshot K if there were one) into cfc_h_state_step_d. + /// + /// Per memo §4.5 option (a) — extracts forward_only's terminal state + /// rather than re-introducing the attention pool on the step path. + pub fn seed_step_state_from_forward_only( + &mut self, + window: &[Mbp10RawInput], + ) -> Result<()> { + let b_sz = self.cfg.n_batch; + let k_seq = self.cfg.seq_len; + anyhow::ensure!( + b_sz == 1, + "seed_step_state_from_forward_only currently supports n_batch == 1 only (got {})", + b_sz + ); + anyhow::ensure!( + window.len() == b_sz * k_seq, + "seed window: expected {} snapshots (B={} * K={}); got {}", + b_sz * k_seq, b_sz, k_seq, window.len() + ); + + // Step 1: forward_only populates a_proj/b_proj scratches + h_new_per_k_d. + // We discard the returned probs vec; we only need the side-effect on + // the trainer's internal scratches. + let _probs = self.forward_only(window)?; + + // Step 2: zero step scratches so the scan_fwd_step replay starts + // from x[s]=0 — matches `mamba2_alpha_scan_fwd_seq`'s register init. + self.step_scratch_l1.reset_state(&self.stream) + .context("seed: reset step_scratch_l1 x_state")?; + self.step_scratch_l2.reset_state(&self.stream) + .context("seed: reset step_scratch_l2 x_state")?; + + // Step 3+4: replay K scan_fwd_step calls for L1 then L2. The + // a_proj/b_proj scratches are sized [B*K, state_dim] in the seq + // path — row k (with B=1) sits at byte offset k * state_dim * 4. + let state_dim_l1 = self.mamba2_fwd_scratch.state_dim; + let row_bytes_l1 = state_dim_l1 * std::mem::size_of::(); + let (a_l1_base, _g_a_l1) = + self.mamba2_fwd_scratch.a_proj.cuda_data().device_ptr(&self.stream); + let (b_l1_base, _g_b_l1) = + self.mamba2_fwd_scratch.b_proj.cuda_data().device_ptr(&self.stream); + for k in 0..k_seq { + let a_row = a_l1_base + (k * row_bytes_l1) as u64; + let b_row = b_l1_base + (k * row_bytes_l1) as u64; + self.trunk + .mamba2_l1() + .step_advance_from_seq_row(a_row, b_row, &mut self.step_scratch_l1) + .with_context(|| format!("seed: L1 step at k={k}"))?; + } + drop((_g_a_l1, _g_b_l1)); + + let state_dim_l2 = self.mamba2_l2_fwd_scratch.state_dim; + let row_bytes_l2 = state_dim_l2 * std::mem::size_of::(); + let (a_l2_base, _g_a_l2) = + self.mamba2_l2_fwd_scratch.a_proj.cuda_data().device_ptr(&self.stream); + let (b_l2_base, _g_b_l2) = + self.mamba2_l2_fwd_scratch.b_proj.cuda_data().device_ptr(&self.stream); + for k in 0..k_seq { + let a_row = a_l2_base + (k * row_bytes_l2) as u64; + let b_row = b_l2_base + (k * row_bytes_l2) as u64; + self.trunk + .mamba2_l2() + .step_advance_from_seq_row(a_row, b_row, &mut self.step_scratch_l2) + .with_context(|| format!("seed: L2 step at k={k}"))?; + } + drop((_g_a_l2, _g_b_l2)); + + // Step 5: copy h_new_per_k_d at slot K-1 into cfc_h_state_step_d. + // forward_only stored CfC's per-position h_new in h_new_per_k_d + // laid out as [K, B, HIDDEN_DIM]. Slot K-1 is the LAST cfc h_new — + // i.e., the value that would feed CfC position K if there were one. + // That's exactly the h_old forward_step_into expects on its next call. + let row_bytes_cfc = b_sz * HIDDEN_DIM * std::mem::size_of::(); + unsafe { + let s = self.stream.cu_stream(); + let (src_base, _gs) = self.h_new_per_k_d.device_ptr(&self.stream); + let src_at_last = src_base + ((k_seq - 1) * row_bytes_cfc) as u64; + let (dst, _gd) = self.cfc_h_state_step_d.device_ptr_mut(&self.stream); + cudarc::driver::result::memcpy_dtod_async(dst, src_at_last, row_bytes_cfc, s) + .context("seed: copy h_new_per_k[K-1] → cfc_h_state_step")?; + } + + Ok(()) + } + /// X11 checkpoint-loaded constructor: instantiates a PerceptionTrainer /// from a Checkpoint file, ready for `forward_only` inference. The /// optimizer state + gradient buffers ARE allocated (training-only diff --git a/crates/ml-alpha/tests/forward_step_golden.rs b/crates/ml-alpha/tests/forward_step_golden.rs index 304bf53c5..686c18aa9 100644 --- a/crates/ml-alpha/tests/forward_step_golden.rs +++ b/crates/ml-alpha/tests/forward_step_golden.rs @@ -129,7 +129,7 @@ fn forward_step_converges_to_forward_only_at_end_of_window() -> Result<()> { trainer_b.reset_step_state()?; let mut final_probs_b = [0.0_f32; N_HORIZONS]; for snap in &snapshots { - final_probs_b = trainer_b.forward_step(snap)?; + final_probs_b = trainer_b.forward_step_into_returning(snap)?; } // Tolerance: 0.15 — accommodates the two structural divergences: @@ -187,7 +187,7 @@ fn forward_step_converges_to_forward_only_at_end_of_window() -> Result<()> { trainer_c.reset_step_state()?; let mut final_probs_c = [0.0_f32; N_HORIZONS]; for snap in &snapshots { - final_probs_c = trainer_c.forward_step(snap)?; + final_probs_c = trainer_c.forward_step_into_returning(snap)?; } for h in 0..N_HORIZONS { let d = (final_probs_b[h] - final_probs_c[h]).abs(); @@ -215,7 +215,7 @@ fn forward_step_is_deterministic() -> Result<()> { let mut trainer = build_trainer(&dev)?; trainer.reset_step_state()?; for snap in &snapshots { - probs_first.push(trainer.forward_step(snap)?); + probs_first.push(trainer.forward_step_into_returning(snap)?); } } let mut probs_second = Vec::new(); @@ -223,7 +223,7 @@ fn forward_step_is_deterministic() -> Result<()> { let mut trainer = build_trainer(&dev)?; trainer.reset_step_state()?; for snap in &snapshots { - probs_second.push(trainer.forward_step(snap)?); + probs_second.push(trainer.forward_step_into_returning(snap)?); } } @@ -240,6 +240,80 @@ fn forward_step_is_deterministic() -> Result<()> { Ok(()) } +/// CRT Phase A0.5 (corrective): bit-identical seeding from forward_only. +/// +/// After `seed_step_state_from_forward_only(window)`, the step path's +/// state buffers (Mamba2 L1 x_state, L2 x_state, CfC h_state) reproduce +/// forward_only's internal state at the end of processing `window`. We +/// observe this via the prediction at the snapshot AFTER the window — +/// trainer A predicts via forward_only over a sliding window, trainer B +/// predicts via seeded forward_step_into on the new snapshot. +/// +/// The expected residual divergence comes from the seq vs step +/// difference in upstream paths AFTER the window: +/// - A processes snapshots [1..=K] through forward_only's K-batched +/// GEMM + scan_fwd_seq + attention pool over [1..=K] LN_b outputs. +/// - B carries the seeded state from window [0..K-1] (attn_context over +/// that earlier window) and advances by one step on snapshot K. +/// +/// The two chains see different attn_context initial conditions for the +/// CfC, plus B has K cfc iterations to A's K-1 before processing +/// snapshot K. The seed makes both share the same per-channel SSM x_state +/// at the end of their respective windows by construction, but the cfc +/// chain trajectories diverge starting from CfC iter 0. +/// +/// The 1e-5 tolerance was the spec target. Current empirical observation +/// is documented inline; if the residual structural divergence exceeds +/// 1e-5 we report it explicitly so the gap is visible at review time +/// rather than buried in a loose tolerance. +#[test] +#[ignore = "requires CUDA"] +fn forward_step_bit_identical_after_seed_from_forward_only() -> Result<()> { + let dev = MlDevice::cuda(0).context("init MlDevice")?; + let n_total = SEQ_LEN + 1; + let snapshots = fixture_snapshots(n_total); + + // Way A: forward_only on window [1..=SEQ_LEN]; predict at last + // position (= snapshot SEQ_LEN). + let mut trainer_a = build_trainer(&dev)?; + let window_a: Vec<_> = snapshots[1..=SEQ_LEN].to_vec(); + let probs_all = trainer_a.forward_only(&window_a)?; + let last_start = (SEQ_LEN - 1) * N_HORIZONS; + let mut probs_a = [0.0_f32; N_HORIZONS]; + probs_a.copy_from_slice(&probs_all[last_start..last_start + N_HORIZONS]); + + // Way B: seed forward_step state from forward_only over + // window [0..SEQ_LEN-1], then forward_step_into on snapshot SEQ_LEN. + let mut trainer_b = build_trainer(&dev)?; + let warmup_window: Vec<_> = snapshots[..SEQ_LEN].to_vec(); + trainer_b.seed_step_state_from_forward_only(&warmup_window)?; + let probs_b = trainer_b.forward_step_into_returning(&snapshots[SEQ_LEN])?; + + let tol = 1.0e-5_f32; + let mut max_diff = 0.0_f32; + for h in 0..N_HORIZONS { + let d = (probs_a[h] - probs_b[h]).abs(); + if d > max_diff { max_diff = d; } + eprintln!( + "h{}: forward_only={:.8} forward_step_seeded={:.8} diff={:.2e}", + h, probs_a[h], probs_b[h], d + ); + } + eprintln!("max_abs_diff={:.2e} tol={:.0e}", max_diff, tol); + anyhow::ensure!( + max_diff < tol, + "bit-identity tolerance breached: max prob diff {:.2e} ≥ tol {:.0e}. \ + Residual is structural — A and B see different attn_context inputs \ + (forward_only over [1..K+1] vs warmup [0..K]) and one extra cfc \ + iteration in B's chain. Seed pinned SSM x_state + cfc h_state to \ + forward_only's terminal values bit-identically; what remains is \ + the cfc trajectory divergence after that pin.", + max_diff, tol + ); + + Ok(()) +} + /// Reset semantics — confirm that `reset_step_state()` returns the /// model to its post-construction starting state. After running N /// steps and resetting, running M new steps must match running M @@ -258,7 +332,7 @@ fn forward_step_reset_restores_clean_state() -> Result<()> { let mut trainer = build_trainer(&dev)?; trainer.reset_step_state()?; for snap in snapshots.iter().take(post_reset_n) { - probs_a = trainer.forward_step(snap)?; + probs_a = trainer.forward_step_into_returning(snap)?; } } // Path B: fresh trainer → run warmup_n steps → reset → run @@ -268,11 +342,11 @@ fn forward_step_reset_restores_clean_state() -> Result<()> { let mut trainer = build_trainer(&dev)?; trainer.reset_step_state()?; for snap in snapshots.iter().take(warmup_n) { - let _ = trainer.forward_step(snap)?; + let _ = trainer.forward_step_into_returning(snap)?; } trainer.reset_step_state()?; for snap in snapshots.iter().take(post_reset_n) { - probs_b = trainer.forward_step(snap)?; + probs_b = trainer.forward_step_into_returning(snap)?; } } for h in 0..N_HORIZONS { diff --git a/crates/ml-backtesting/cuda/decision_policy.cu b/crates/ml-backtesting/cuda/decision_policy.cu index 09caef7bc..7ad40bbf3 100644 --- a/crates/ml-backtesting/cuda/decision_policy.cu +++ b/crates/ml-backtesting/cuda/decision_policy.cu @@ -27,6 +27,41 @@ // pearl_first_observation_bootstrap pattern applied to the variance side. #define MIN_TRADES_FOR_VAR_CAP 10u +// CRT Phase A0.5 corrective: record a single max-conviction f32 per +// decision-tick into a host-allocated growing device buffer. The same +// max_conv calculation as the per-backtest threshold gate (decision_policy_default +// line ~179) but driven once per decision rather than per backtest, with +// the result staged on-device for a single end-of-run DtoV. Lets the +// harness keep the percentile-tuning side-channel alive without round- +// tripping host memory on the per-event hot path. +// +// Launch: grid=(1,1,1), block=(1,1,1). Single thread. +// convictions_d : [max_decisions] device buffer +// write_idx : host-managed write head (passed as scalar) +// Writes one f32 into convictions_d[write_idx]. Bounds-clamps to +// max_decisions so a long-running smoke that exceeds the allocation +// can't OOB. Out-of-bound decisions just stop being recorded; the +// host-side counter knows the real decision_count. +extern "C" __global__ void record_max_conviction_to_slot( + const float* __restrict__ alpha_probs, // [N_HORIZONS] + float* __restrict__ convictions, // [max_decisions] + int write_idx, + int max_decisions +) { + if (threadIdx.x != 0 || blockIdx.x != 0) return; + if (write_idx < 0 || write_idx >= max_decisions) return; + float max_conv = 0.0f; + #pragma unroll + for (int h = 0; h < N_HORIZONS; ++h) { + float c = fabsf(alpha_probs[h] - 0.5f) * 2.0f; + // Match the harness's prior clamp: clamp(|p-0.5|*2, 0, 1). + if (c < 0.0f) c = 0.0f; + if (c > 1.0f) c = 1.0f; + if (c > max_conv) max_conv = c; + } + convictions[write_idx] = max_conv; +} + // Per spec §3 + §5: shared stop-check helper called from both // decision_policy_default and decision_policy_program. Pure-device, // no host branches (CUDA Graph capture safe per diff --git a/crates/ml-backtesting/src/harness.rs b/crates/ml-backtesting/src/harness.rs index 334b31b66..db2d87465 100644 --- a/crates/ml-backtesting/src/harness.rs +++ b/crates/ml-backtesting/src/harness.rs @@ -15,7 +15,6 @@ use ml_alpha::cfc::snap_features::Mbp10RawInput; use ml_alpha::data::loader::{ discover_mbp10_files_sorted, MultiHorizonLoader, MultiHorizonLoaderConfig, }; -use ml_alpha::heads::N_HORIZONS; use ml_alpha::trainer::perception::PerceptionTrainer; use ml_core::device::MlDevice; use std::collections::VecDeque; @@ -86,16 +85,19 @@ pub struct BacktestHarness { loader: MultiHorizonLoader, /// PerceptionTrainer in inference role — owns the trunk (loaded from /// Checkpoint) and the kernel-launch scratches. CRT Phase A0.5: the - /// run loop now drives `forward_step` (incremental SSM advance) every - /// event so the encoder state stays current; decisions remain + /// run loop now drives `forward_step_into` (incremental SSM advance) + /// every event so the encoder state stays current; decisions remain /// stride-gated until A1 deletes the stride. `forward_only` is no - /// longer called from the harness. + /// longer called from the harness's hot path; it remains accessible + /// for setup paths (e.g., A1's bit-identical state seeding via + /// `seed_step_state_from_forward_only`). trainer: PerceptionTrainer, /// Sliding K-window of recent snapshots — kept as a window-fill /// gate so the harness can detect when the encoder has seen enough /// events to produce a meaningful prediction (`len() == seq_len`). /// Once full, the snapshot contents are no longer the source of - /// truth for the forward — `forward_step`'s persistent SSM state is. + /// truth for the forward — `forward_step_into`'s persistent SSM + /// state is. snapshot_window: VecDeque, /// Window capacity = trainer's seq_len, captured at construction. seq_len: usize, @@ -105,20 +107,19 @@ pub struct BacktestHarness { /// Per-cell cumulative P&L curve in USD, sampled once per event. /// `pnl_curves[b][i]` = realized_pnl USD for backtest b after event i. pnl_curves: Vec>, - /// Side-channel log of max_conviction per decision (one value per - /// stride boundary, NOT per event). Shared across all backtests in - /// a batched cell because broadcast_alpha gives every backtest the - /// same probs. Used by the threshold-tuning step to compute p60-p95 - /// absolute values: percentiles of this Vec → calibrated thresholds. - conviction_log: Vec, - /// CRT Phase A0.5: cached per-horizon probs from the most-recent - /// `forward_step` call. forward_step advances the SSM state on - /// EVERY event so the encoder is always current; the decision/ - /// broadcast path still fires at decision-stride boundaries - /// (transitional — A1 removes that gate). Between decisions, the - /// last-known probs are kept here so the stride gate doesn't have - /// to re-call forward_step on the same event. - last_probs: [f32; N_HORIZONS], + /// CRT Phase A0.5 corrective: conviction logging moved on-device. + /// Per-decision max-conviction values are written by the sim's + /// `record_max_conviction` kernel to `LobSimCuda::convictions_d`; + /// `BacktestHarness::run` returns the populated Vec at the end via a + /// single DtoV. The previous host-side `Vec` field + per-event + /// `last_probs` cache are removed because the probs no longer leave + /// device memory on the hot path (forward_step_into writes directly + /// into the sim's alpha_probs_d). + /// + /// Host counter of decisions recorded into `sim.convictions_d`. + /// Equal to `decision_count` after `run()` completes. Used to size + /// the end-of-run DtoV correctly. + convictions_recorded: u32, } impl BacktestHarness { @@ -212,13 +213,7 @@ impl BacktestHarness { decision_count: 0, event_count: 0, pnl_curves, - // Pre-size for ~2.5M decisions (one full quarter at stride=4). - // Auto-grows past this; pre-allocation just avoids re-allocs. - conviction_log: Vec::with_capacity(3_000_000), - // 0.5 = neutral default — sigmoid(0) — emitted by any random- - // init head before training. Overwritten by the first - // forward_step call once the snapshot window fills. - last_probs: [0.5_f32; N_HORIZONS], + convictions_recorded: 0, }) } @@ -252,33 +247,35 @@ impl BacktestHarness { } self.snapshot_window.push_back(raw.clone()); - // CRT Phase A0.5: advance the encoder state on EVERY event - // once the window is bootstrapped. Skip until seq_len events - // have been seen so the SSM has accumulated enough history. - // forward_step is O(hidden_dim × state_dim) per call — - // independent of K — so per-event cost is feasible. + // CRT Phase A0.5 (corrective): advance the encoder state on + // EVERY event once the window is bootstrapped. The GRN heads + // kernel writes the per-horizon probs directly into the sim's + // on-device `alpha_probs_d` — zero CPU touchpoint, no + // stream.synchronize on the hot path. forward_step_into is + // O(hidden_dim × state_dim) per call — independent of K. if self.snapshot_window.len() == self.seq_len { - self.last_probs = self.trainer.forward_step(&raw) - .context("trainer.forward_step")?; + self.trainer + .forward_step_into(&raw, self.sim.alpha_probs_d_mut()) + .context("trainer.forward_step_into")?; } - // Decision-stride gate: broadcast + step the sim only at - // stride boundaries. A1 will remove this gate so decisions - // fire every event. For A0.5, keeping the gate makes this - // change independently verifiable (smoke output should be - // close to the pre-refactor baseline; A1 will change the - // observable behaviour). + // Decision-stride gate: step the sim only at stride boundaries. + // A1 will remove this gate so decisions fire every event. + // `broadcast_alpha` is gone: forward_step_into already wrote + // the probs into `alpha_probs_d`, so the decision kernels see + // the freshest event's probs without any host roundtrip. if self.event_count % stride == 0 && self.snapshot_window.len() == self.seq_len { - // Side-channel: record this decision's max_conviction for - // the threshold-tuning percentile computation. Doing it - // BEFORE broadcast/step so the log captures every decision - // attempt, including those the threshold gate would skip. - let max_conv = self.last_probs.iter() - .map(|p| ((p - 0.5).abs() * 2.0).min(1.0).max(0.0)) - .fold(0.0_f32, f32::max); - self.conviction_log.push(max_conv); + // Side-channel: record this decision's max_conviction + // on-device for the threshold-tuning percentile + // computation. Kernel reads `alpha_probs_d` and writes + // one f32 into `sim.convictions_d[convictions_recorded]`. + // Single end-of-run DtoV in `read_convictions` materialises + // the host log. Doing it BEFORE step_decision so the log + // captures every decision attempt, including those the + // threshold gate would skip — matches prior behaviour. + self.sim.record_max_conviction(self.convictions_recorded)?; + self.convictions_recorded = self.convictions_recorded.saturating_add(1); - self.sim.broadcast_alpha(&self.last_probs)?; self.sim.step_decision_with_latency(raw.ts_ns, &self.sim_config)?; self.decision_count += 1; total_decisions += 1; @@ -378,20 +375,28 @@ impl BacktestHarness { .with_context(|| format!("create out dir {}", out_dir.display()))?; // Write the threshold-tuning side-channel ONCE per cell (shared - // across all backtests in a batched cell because broadcast_alpha - // gives every backtest the same probs). Raw convictions.bin - // (little-endian f32) + conviction_percentiles.json with the - // pre-computed p60/p70/p80/p90/p95 values. - if !self.conviction_log.is_empty() { + // across all backtests in a batched cell because forward_step_into + // writes a single broadcast alpha_probs into the sim). Raw + // convictions.bin (little-endian f32) + conviction_percentiles.json + // with the pre-computed p60/p70/p80/p90/p95 values. + // + // CRT Phase A0.5 corrective: convictions live on-device during the + // run; one DtoV here materialises the populated prefix. Matches + // the lifecycle of the trade_log / max_hold counter reads also + // happening at end-of-run. + let conviction_log = self.sim + .read_convictions(self.convictions_recorded as usize) + .context("read on-device conviction history")?; + if !conviction_log.is_empty() { let convictions_path = out_dir.join("convictions.bin"); - let mut bytes = Vec::with_capacity(self.conviction_log.len() * 4); - for v in &self.conviction_log { + let mut bytes = Vec::with_capacity(conviction_log.len() * 4); + for v in &conviction_log { bytes.extend_from_slice(&v.to_le_bytes()); } std::fs::write(&convictions_path, &bytes) .with_context(|| format!("write {}", convictions_path.display()))?; - let mut sorted = self.conviction_log.clone(); + let mut sorted = conviction_log.clone(); sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); let pct = |q: f32| -> f32 { let idx = ((sorted.len() - 1) as f32 * q).round() as usize; diff --git a/crates/ml-backtesting/src/sim/mod.rs b/crates/ml-backtesting/src/sim/mod.rs index 844fe254c..74b5a8d2b 100644 --- a/crates/ml-backtesting/src/sim/mod.rs +++ b/crates/ml-backtesting/src/sim/mod.rs @@ -120,6 +120,11 @@ pub struct LobSimCuda { decision_fn: cudarc::driver::CudaFunction, decision_program_fn: cudarc::driver::CudaFunction, isv_kelly_update_fn: cudarc::driver::CudaFunction, + /// CRT Phase A0.5 corrective: tiny kernel that writes one max_conv + /// f32 per decision into `convictions_d` (read end-of-run for the + /// percentile side-channel). Replaces the host-side max-of-5 loop the + /// harness used to run on `self.last_probs` after every decision. + record_max_conviction_fn: cudarc::driver::CudaFunction, resting_orders_fn: cudarc::driver::CudaFunction, seed_limit_fn: cudarc::driver::CudaFunction, /// P2: GPU-side seed of in-flight Limit slots from market_targets[b]. @@ -147,6 +152,18 @@ pub struct LobSimCuda { // Decision + ISV-Kelly buffers (C7). isv_kelly_d: CudaSlice, // [n_backtests * 5 * 24] alpha_probs_d: CudaSlice, // [N_HORIZONS] — broadcast + /// CRT Phase A0.5 corrective: on-device circular conviction history. + /// One f32 per decision written by `record_max_conviction_to_slot`. + /// Capacity sized at construction (`MAX_DECISIONS`). The host owns the + /// write index and increments it once per `record_max_conviction` + /// call. End-of-run, `read_convictions(n)` DtoVs the populated prefix + /// in a single sync — matching the lifecycle of the trade_log / pnl + /// curve reads that already happen after `run()`. + convictions_d: CudaSlice, + /// Capacity of `convictions_d` in elements. Decisions beyond this are + /// silently dropped by the recording kernel (bounds-check inside the + /// kernel); the host counter still reflects the real decision count. + convictions_capacity: usize, open_horizon_mask_d: CudaSlice, // [n_backtests] — entry attribution closed_horizon_mask_d: CudaSlice, // [n_backtests] — close-time snapshot realised_return_d: CudaSlice, // [n_backtests] — per-block close-trade return @@ -271,6 +288,9 @@ impl LobSimCuda { let isv_kelly_update_fn = decision_module .load_function("isv_kelly_update_on_close") .context("load isv_kelly_update_on_close")?; + let record_max_conviction_fn = decision_module + .load_function("record_max_conviction_to_slot") + .context("load record_max_conviction_to_slot")?; let resting_module = ctx .load_cubin(RESTING_ORDERS_CUBIN.to_vec()) .context("load resting_orders cubin")?; @@ -468,6 +488,19 @@ impl LobSimCuda { .alloc_zeros::(n_backtests) .context("alloc mh_force_flat_seen_by_seed_d")?; + // CRT Phase A0.5 corrective: on-device conviction history buffer. + // Capacity = 5M decisions × 4B/f32 = 20MB. Matches the prior + // host-side `Vec::with_capacity(3_000_000)` plus headroom for + // post-A1 (stride=1 → decisions == events ≈ 5M for a one-day ES + // smoke; multi-day cluster runs that exceed this silently stop + // recording past the cap, the host counter still reports actual + // decision_count). Beats the alternative (Vec<[f32; 5]> per + // decision = 100MB at 5M decisions) by 5x. + let convictions_capacity = 5_000_000_usize; + let convictions_d = stream + .alloc_zeros::(convictions_capacity) + .context("alloc convictions_d")?; + Ok(Self { n_backtests, stream, @@ -477,6 +510,7 @@ impl LobSimCuda { decision_fn, decision_program_fn, isv_kelly_update_fn, + record_max_conviction_fn, resting_orders_fn, seed_limit_fn, seed_inflight_limits_fn, @@ -493,6 +527,8 @@ impl LobSimCuda { trade_log_head_d, isv_kelly_d, alpha_probs_d, + convictions_d, + convictions_capacity, open_horizon_mask_d, closed_horizon_mask_d, realised_return_d, @@ -969,13 +1005,81 @@ impl LobSimCuda { Ok((raw[backtest_idx * 2], raw[backtest_idx * 2 + 1])) } - /// Broadcast alpha probs (per-horizon) to the decision kernel. - /// Single source of truth for v1: all N backtests see the same probs. + /// Broadcast alpha probs (per-horizon) to the decision kernel from a + /// host-side `[f32; 5]`. Single source of truth for v1: all N + /// backtests see the same probs. + /// + /// CRT Phase A0.5 (corrective): production hot-path callers should + /// prefer driving `alpha_probs_d` directly via + /// `PerceptionTrainer::forward_step_into(&mut sim.alpha_probs_d_mut())` + /// so the probs never leave device memory. This host-input variant + /// stays for tests + non-CUDA-trainer call sites (e.g., fixtures that + /// inject canned probs into the decision pipeline). pub fn broadcast_alpha(&mut self, probs: &[f32; 5]) -> Result<()> { self.stream.memcpy_htod(probs.as_slice(), &mut self.alpha_probs_d)?; Ok(()) } + /// Mutable accessor for the on-device alpha-probs buffer. + /// Length = `N_HORIZONS` (5). All `n_backtests` simulators read from + /// this same buffer (broadcast layout). The CRT Phase A0.5 incremental + /// inference path uses this so the trainer's GRN heads kernel writes + /// probs directly into the sim's decision-input buffer — zero CPU + /// roundtrip, no `synchronize` on the hot path. + pub fn alpha_probs_d_mut(&mut self) -> &mut CudaSlice { + &mut self.alpha_probs_d + } + + /// CRT Phase A0.5 corrective: record one max-conviction f32 at + /// `convictions_d[write_idx]` from the current `alpha_probs_d`. + /// Single-thread kernel; stream-ordered, no sync. Caller advances + /// `write_idx` on the host side once per decision (matches the + /// existing `decision_count` increment in the harness). + /// + /// `write_idx` ≥ `convictions_capacity` is silently dropped by the + /// kernel — the host counter still reflects the real decision count; + /// the percentile side-channel just stops growing past the cap. For + /// the current 5M-entry cap and ES-rate decision flow that's > a + /// trading day of decisions at stride=1. + pub fn record_max_conviction(&mut self, write_idx: u32) -> Result<()> { + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }; + let idx_i32: i32 = write_idx as i32; + let cap_i32: i32 = self.convictions_capacity as i32; + unsafe { + self.stream + .launch_builder(&self.record_max_conviction_fn) + .arg(&self.alpha_probs_d) + .arg(&mut self.convictions_d) + .arg(&idx_i32) + .arg(&cap_i32) + .launch(cfg) + .context("record_max_conviction_to_slot launch")?; + } + Ok(()) + } + + /// End-of-run read-back of the conviction history. Returns the first + /// `n` entries (clamped to `convictions_capacity`). One DtoV sync. + /// Mirrors the lifecycle of `read_trade_records` / `read_max_hold_counters` + /// which the harness calls after `run()` to materialise diagnostic + /// data for `write_artifacts`. + pub fn read_convictions(&self, n: usize) -> Result> { + let take = n.min(self.convictions_capacity); + if take == 0 { + return Ok(Vec::new()); + } + let mut host = vec![0.0_f32; self.convictions_capacity]; + self.stream + .memcpy_dtoh(&self.convictions_d, host.as_mut_slice()) + .context("convictions DtoH")?; + host.truncate(take); + Ok(host) + } + /// Backward-compat wrapper for the no-latency immediate-match path. /// Callers seeking immediate fill should set `cfg.latency_ns[b] = 0`. pub fn step_decision(