From ed8b53b474b07c5691aea48163667b0ca07a58c6 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 21 May 2026 19:10:22 +0200 Subject: [PATCH] feat(per-horizon-cfc): forward_step_into uses Phase 2 fused dispatch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per spec §2.4 and Task 13 of the plan. Replace the single-CfC `cfc_step_batched` dispatch in `forward_step_into` with `cfc_step_per_branch_fwd_gpu`. Production checkpoints have Phase 2 routing frozen at the training-time Phase 1→2 transition, so inference unconditionally takes the per-branch path. The fused kernel reads `bucket_channel_offset_d` / `bucket_dim_k_d` from the trunk; in deployment these are populated via `CfcTrunk::load_checkpoint`. For freshly-constructed trainers without a checkpoint load, the trunk fields are zero — the kernel's uniform predicate (`tid >= bucket_dim_k`) then early-returns every thread and h_new is left untouched. That matches the "Phase 2 only at inference" contract. Heads dispatch unchanged: the existing GRN kernel reads from `trunk.heads_w_skip_d` which has off-bucket entries zeroed at the transition (sparsification per the prior block-diagonal-grad-mask follow-up commit). Mathematically the full-buffer read is equivalent to a compact-only per-bucket read; no inference-time kernel change required. forward_step_golden's convergence test is now architecturally incompatible with the new dispatch — `forward_only` runs Phase 1 single-CfC math while `forward_step_into` requires populated Phase 2 routing. Annotated `#[ignore]` with a clear divergence note; the deterministic + reset semantics tests remain valid invariants per `pearl_training_smoothness_does_not_transfer_to_inference`. cargo check workspace clean (excluding pre-existing unrelated cupti + insert_batch errors in vendor/cudarc and ml/tests). ml-alpha + ml- backtesting lib tests pass (33 + 33). Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml-alpha/src/trainer/perception.rs | 66 +++++++++++------ crates/ml-alpha/tests/forward_step_golden.rs | 77 ++++++++++---------- 2 files changed, 83 insertions(+), 60 deletions(-) diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 289a1d4c3..b0f9bd2a4 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -4625,29 +4625,43 @@ impl PerceptionTrainer { unsafe { launch.launch(cfg_ln).context("step LN_b fwd")?; } } - // ── 9. CfC single step. h_old is the persistent cfc_h_state_step; - // h_new is written to cfc_h_new_step (then copied back into - // cfc_h_state_step). - // A1: decision_stride removed; event-rate dt = 1.0 per CfC step. + // ── 9. Phase 2 fused per-branch CfC step. + // + // Per spec §2.4 (per-horizon CfC inference) and Task 13: production + // checkpoints always have Phase 2 routing frozen at the training-time + // Phase 1→2 transition. Inference therefore unconditionally dispatches + // the per-branch fused kernel (`cfc_step_per_branch_fwd`) instead of + // the single-CfC `cfc_step_batched` used during Phase 1 warmup. + // + // The fused kernel reads `bucket_channel_offset_d` / `bucket_dim_k_d` + // from the trunk; in deployment these are populated via + // `CfcTrunk::load_checkpoint` (cfc/trunk.rs:566-567). For freshly- + // constructed trainers (no checkpoint), the trunk fields are zero- + // initialised — the kernel's uniform predicate (`tid >= bucket_dim_k`) + // then early-returns every thread and h_new is left untouched. That + // matches the "Phase 2 only at inference" contract: callers without + // populated routing metadata are NOT on the production deployment + // path. + // + // dt = 1.0 per event-rate inference (A1: decision_stride removed). + // n_batch_i kept here because the heads GRN kernel below also reads it. let dt_s: f32 = 1.0; - let n_in_i: i32 = HIDDEN_DIM as i32; - let n_hid_i: i32 = HIDDEN_DIM as i32; let n_batch_i: i32 = b_sz as i32; - let cfc_fwd_smem = (2 * HIDDEN_DIM * std::mem::size_of::()) as u32; - let cfg_cfc = LaunchConfig { - grid_dim: (b_sz as u32, 1, 1), - block_dim: (HIDDEN_DIM as u32, 1, 1), - shared_mem_bytes: cfc_fwd_smem, - }; - unsafe { - let mut launch = self.stream.launch_builder(&self.trunk.step_batched_fn); - launch - .arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_all_d) - .arg(&self.ln_b_step_out_d).arg(&self.cfc_h_state_step_d) - .arg(&dt_s).arg(&n_in_i).arg(&n_hid_i).arg(&n_batch_i) - .arg(&mut self.cfc_h_new_step_d); - launch.launch(cfg_cfc).context("step cfc")?; - } + crate::cfc::step::cfc_step_per_branch_fwd_gpu( + &self.stream, + &self.trunk.cfc_step_per_branch_fwd_fn, + &self.trunk.w_in_d, + &self.trunk.w_rec_d, + &self.trunk.b_d, + &self.trunk.tau_all_d, + &self.ln_b_step_out_d, + &self.cfc_h_state_step_d, + dt_s, + n_batch_i, + &self.trunk.bucket_channel_offset_d, + &self.trunk.bucket_dim_k_d, + &mut self.cfc_h_new_step_d, + ).context("forward_step_into cfc_step_per_branch_fwd")?; // Carry-forward: h_new → h_state (in-place via DtoD). unsafe { @@ -4660,6 +4674,16 @@ impl PerceptionTrainer { } // ── 10. Heads GRN fwd on the new h. + // + // The heads dispatch stays at the existing GRN kernel for both phases. + // Per the Task 10 follow-up commit (block-diagonal heads via grad-mask), + // `trunk.heads_w_skip_d` is sparsified at the Phase 1→2 transition: + // off-bucket entries are zeroed in place, and the per-step grad mask + // (applied in `step_batched` after the optimizer step) keeps those + // zeros frozen. Mathematically the GRN kernel's read of the full + // [N_HORIZONS × HIDDEN_DIM] `heads_w_skip_d` is equivalent to a + // compact-only read of each horizon's bucket slice — no inference- + // time kernel change required. let cfg_grn_fwd = LaunchConfig { grid_dim: (b_sz as u32, 1, 1), block_dim: (HEAD_MID_DIM as u32, 1, 1), diff --git a/crates/ml-alpha/tests/forward_step_golden.rs b/crates/ml-alpha/tests/forward_step_golden.rs index e28284501..dae6f97f3 100644 --- a/crates/ml-alpha/tests/forward_step_golden.rs +++ b/crates/ml-alpha/tests/forward_step_golden.rs @@ -1,46 +1,32 @@ //! CRT Phase A0.5 — forward_step structural integrity gate. //! //! Validates the `PerceptionTrainer::forward_step` event-rate inference -//! path lands on a stable, deterministic, reset-able recurrent state and -//! produces predictions that are STRUCTURALLY CONSISTENT with -//! `forward_only` after sufficient warmup. +//! path lands on a stable, deterministic, reset-able recurrent state. //! -//! Architecture note (informs the test tolerance below): +//! Architecture note — per-horizon CfC inference (Task 13, 2026-05-21): //! -//! forward_step is NOT bit-identical to forward_only over a K-window. -//! The two paths diverge at trainer init by design: +//! `forward_step_into` now dispatches the Phase 2 fused per-branch CfC +//! kernel (`cfc_step_per_branch_fwd_gpu`) per spec §2.4. The kernel reads +//! `bucket_channel_offset_d` / `bucket_dim_k_d` from the trunk; in +//! deployment these are populated via `CfcTrunk::load_checkpoint` +//! (cfc/trunk.rs:566-567). For freshly-constructed trainers WITHOUT a +//! checkpoint load (this test's setup) the trunk fields are zero, so the +//! kernel's uniform predicate early-returns every thread and h_new is +//! left untouched. //! -//! * forward_only initialises CfC's `h_old` at k=0 from the attention -//! pool over the K-window of LN_b outputs (a learned content-summary). -//! * forward_step has no attention pool — CfC carries its own hidden -//! state across calls; after `reset_step_state()`, that state is zero. +//! This is by design: production inference uses checkpoints with frozen +//! Phase 2 routing. The convergence test (`forward_step_converges_to_ +//! forward_only_at_end_of_window`) is therefore architecturally +//! incompatible with the new dispatch — `forward_only` runs the Phase 1 +//! single-CfC math while `forward_step_into` requires populated Phase 2 +//! routing. The test is preserved with `#[ignore]` documenting this +//! divergence; the deterministic + reset semantics tests remain valid +//! invariants that do not depend on convergence between the two paths. //! -//! The attention pool was dropped from the per-event path because it -//! requires K LN_b rows on every call, defeating the O(1)/event target -//! that motivated A0.5 in the first place. The trade-off: CfC's natural -//! decay (`exp(-dt/tau)`) absorbs the initial-state discrepancy as the -//! sequence grows. For τ < N · dt the influence of the initial h -//! dampens to float-noise; for τ ≫ N · dt the steady-state difference -//! remains. -//! -//! Test design: -//! 1. Generate N=320 deterministic snapshots from a fixed PRNG seed. -//! 2. Way A: call forward_only ONCE on the last seq_len=64 snapshots -//! of the prefix. -//! 3. Way B: call forward_step on a FRESHLY-RESET trainer over ALL N -//! snapshots, take the final-step probs. -//! 4. Compare last-position probs (Way A) to final-step probs (Way B). -//! -//! Tolerance: 0.15 — covers the attn_context vs zero initial-state -//! contribution to CfC after ~K iterations of decay. The kernel-level -//! correctness invariants (determinism across runs; reset returns to -//! clean state) are checked in separate strict tests below. -//! -//! Bit-identical equivalence requires either (a) attention-pool the -//! step path's LN_b history (defeats the per-event O(1) target), or -//! (b) extract Mamba2 + CfC terminal state from a one-shot forward_only -//! and seed forward_step from it (A0 memo §4.5 option (a)). Both are -//! deferred to future tasks; A0.5's scope is the structural path. +//! Pearl: `pearl_training_smoothness_does_not_transfer_to_inference` +//! already established that K-window training output and sequential +//! `forward_step_into` output are different geometries; per-horizon CfC +//! makes that divergence structural rather than incidental. use anyhow::{Context, Result}; use ml_alpha::cfc::snap_features::{Mbp10RawInput, REGIME_DIM}; @@ -103,10 +89,23 @@ fn build_trainer(dev: &MlDevice) -> Result { /// End-to-end convergence test: forward_step over a long warmup prefix /// reaches the same per-horizon probs as forward_only on the trailing -/// K-window. Tolerance is loose (1e-2) — see docs at module head for -/// the two divergence sources that don't fully vanish at finite N. +/// K-window. +/// +/// **Architectural divergence (Task 13, 2026-05-21):** This test is +/// preserved for historical context but is no longer expected to pass +/// against a freshly-constructed trainer. `forward_step_into` now +/// dispatches the Phase 2 fused per-branch CfC kernel which requires +/// `trunk.bucket_channel_offset_d` / `trunk.bucket_dim_k_d` to be +/// populated (via `load_checkpoint` in production deployment). A trainer +/// built from `PerceptionTrainer::new` carries zero-initialised bucket +/// metadata, so the per-branch kernel's uniform predicate early-returns +/// every thread and `forward_step_into` leaves h_new untouched. Meanwhile +/// `forward_only` runs the Phase 1 single-CfC math and produces real +/// h_new values — the two paths cannot converge by construction. See the +/// module-level doc and `pearl_training_smoothness_does_not_transfer_to_ +/// inference` for the broader context. #[test] -#[ignore = "requires CUDA"] +#[ignore = "architectural divergence: forward_step uses Phase 2 per-branch CfC requiring populated bucket metadata (production checkpoint); fresh trainer has zero-init metadata. See module doc."] fn forward_step_converges_to_forward_only_at_end_of_window() -> Result<()> { let dev = MlDevice::cuda(0).context("init MlDevice")?;