From ef29e13c7448af0325c3c3ab707cafaa296d9100 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 19 May 2026 08:20:33 +0200 Subject: [PATCH] refactor(ml-alpha): move CfC weights from PerceptionTrainer to CfcTrunk (X8) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds CfcConfig.cfc_n_in field (default HIDDEN_DIM) so the trunk's CfC layer is sized for v2 usage (CfC input = LN_b output = HIDDEN_DIM) rather than V1 usage (CfC input = snap features = FEATURE_DIM). The previous CfcConfig.n_in field stays as "raw snap feature dim" for any V1 callers still in the tree (fxt-backtest, trunk_forward.rs, graph_a_replay.rs); their forward paths will be cleaned up in X10/X11 when the v2 forward graph lives natively on the trunk. Trainer's w_in_d / w_rec_d / b_d / tau_d fields are removed. Trainer init still draws CfC weights from its ChaCha8Rng chain at the same call position, then memcpy_htod's them into the trunk's now-v2-shaped slots. Forward, backward, and AdamW access sites redirect to self.trunk.{w_in,w_rec,b,tau}_d. Verification: - perception_forward_golden: PASS (max_diff = 0.000000) - ml-alpha lib tests: 34 pass Per spec §2.2 (X8). --- crates/ml-alpha/src/cfc/trunk.rs | 36 ++++++++++++++++------- crates/ml-alpha/src/trainer/perception.rs | 36 +++++++++++------------ 2 files changed, 43 insertions(+), 29 deletions(-) diff --git a/crates/ml-alpha/src/cfc/trunk.rs b/crates/ml-alpha/src/cfc/trunk.rs index 51b968176..8a31c4140 100644 --- a/crates/ml-alpha/src/cfc/trunk.rs +++ b/crates/ml-alpha/src/cfc/trunk.rs @@ -49,8 +49,14 @@ const PROJ_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/projection.c #[derive(Clone, Debug)] pub struct CfcConfig { + /// Raw snap feature dimensionality (V1 forward path uses this). pub n_in: usize, pub n_hid: usize, + /// CfC layer input dimensionality. In V1 this was conflated with + /// `n_in` (CfC consumed snap features directly). In V2 the CfC sits + /// AFTER LN_b, consuming HIDDEN_DIM-dim tokens. Default matches V2 + /// usage. + pub cfc_n_in: usize, /// Mamba2 state dimension — used by the v2 trunk weight skeleton /// (X1) and the per-stack Mamba2Block constructions added in X3/X5. /// Default matches `PerceptionTrainerConfig::default().mamba2_state_dim`. @@ -59,7 +65,12 @@ pub struct CfcConfig { impl Default for CfcConfig { fn default() -> Self { - Self { n_in: FEATURE_DIM, n_hid: HIDDEN_DIM, mamba2_state_dim: 16 } + Self { + n_in: FEATURE_DIM, + n_hid: HIDDEN_DIM, + cfc_n_in: HIDDEN_DIM, + mamba2_state_dim: 16, + } } } @@ -77,13 +88,15 @@ pub struct CfcTrunk { heads_fn: CudaFunction, proj_fn: CudaFunction, - // Device-resident weights — V1 (CfC + simple heads + projection). - // Preserved unchanged until X8/X9 consolidate against the v2 fields - // added below. - w_in_d: CudaSlice, - w_rec_d: CudaSlice, - b_d: CudaSlice, - tau_d: CudaSlice, + // CfC weights — now v2-shaped (cfc_n_in = HIDDEN_DIM). Public so + // PerceptionTrainer can mutate during init + AdamW. + pub w_in_d: CudaSlice, + pub w_rec_d: CudaSlice, + pub b_d: CudaSlice, + pub tau_d: CudaSlice, + // V1 legacy fields — only used by the now-stub V1 forward path + // (dispatch_perception). Kept allocated so the trunk struct lays out + // the same in tests; will be removed in a follow-up cleanup commit. heads_w_d: CudaSlice, heads_b_d: CudaSlice, proj_w_d: CudaSlice, @@ -165,9 +178,12 @@ impl CfcTrunk { // Random init. let mut r = ChaCha8Rng::seed_from_u64(seed); - let scale_in = (1.0_f32 / cfg.n_in as f32).sqrt(); + // X8: CfC layer weights now use cfc_n_in (V2 = HIDDEN_DIM), not + // n_in (which still means snap feature dimensionality for any + // residual V1 callers). This re-shapes w_in_d to v2 layout. + let scale_in = (1.0_f32 / cfg.cfc_n_in as f32).sqrt(); let scale_rec = (1.0_f32 / cfg.n_hid as f32).sqrt(); - let w_in: Vec = (0..cfg.n_hid * cfg.n_in).map(|_| r.gen_range(-scale_in..scale_in)).collect(); + let w_in: Vec = (0..cfg.n_hid * cfg.cfc_n_in).map(|_| r.gen_range(-scale_in..scale_in)).collect(); let w_rec: Vec = (0..cfg.n_hid * cfg.n_hid).map(|_| r.gen_range(-scale_rec..scale_rec)).collect(); let b: Vec = vec![0.0; cfg.n_hid]; // tau log-uniform over [10ms, 1000s] per spec Section 2 (ln(0.01)=-4.6, ln(1000)=6.9). diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 805d5dcac..adcf7ced9 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -206,10 +206,9 @@ pub struct PerceptionTrainer { grad_h_enriched_seq_d: GpuTensor, // CfC + heads weights + their AdamWs (6 groups — tau is trained now). - pub w_in_d: CudaSlice, - pub w_rec_d: CudaSlice, - pub b_d: CudaSlice, - pub tau_d: CudaSlice, + // X8: CfC weights (w_in, w_rec, b, tau) moved to + // `self.trunk.w_in_d` / `w_rec_d` / `b_d` / `tau_d` (now v2-shaped + // via CfcConfig.cfc_n_in = HIDDEN_DIM). // ── TFT GRN heads (Phase 1.7) ── // Per-horizon Gated Residual Network: 2-layer GELU MLP body // (eta_2 → eta_1) + GLU gate + skip-projection. Gives per-horizon @@ -521,6 +520,7 @@ impl PerceptionTrainer { &crate::cfc::trunk::CfcConfig { n_in: FEATURE_DIM, n_hid: HIDDEN_DIM, + cfc_n_in: HIDDEN_DIM, mamba2_state_dim: cfg.mamba2_state_dim, }, cfg.seed, @@ -714,10 +714,12 @@ impl PerceptionTrainer { .map(|_| r.gen_range(-scale_h_skip..scale_h_skip)).collect(); let heads_b_skip: Vec = vec![0.0; N_HORIZONS]; - let w_in_d = upload(&stream, &w_in)?; - let w_rec_d = upload(&stream, &w_rec)?; - let b_d = upload(&stream, &b)?; - let tau_d = upload(&stream, &tau)?; + // X8: upload trainer's PRNG-derived CfC weights into the trunk's + // (now v2-shaped) CfC slots. + stream.memcpy_htod(&w_in, &mut trunk.w_in_d).context("trunk.w_in_d upload")?; + stream.memcpy_htod(&w_rec, &mut trunk.w_rec_d).context("trunk.w_rec_d upload")?; + stream.memcpy_htod(&b, &mut trunk.b_d).context("trunk.b_d upload")?; + stream.memcpy_htod(&tau, &mut trunk.tau_d).context("trunk.tau_d upload")?; let heads_w1_d = upload(&stream, &heads_w1)?; let heads_b1_d = upload(&stream, &heads_b1)?; let heads_w2_d = upload(&stream, &heads_w2)?; @@ -1043,10 +1045,6 @@ impl PerceptionTrainer { ts_ns_all_d: stream.alloc_zeros::(cfg.n_batch * k)?, prev_ts_ns_all_d: stream.alloc_zeros::(cfg.n_batch * k)?, - w_in_d, - w_rec_d, - b_d, - tau_d, heads_w1_d, heads_b1_d, heads_w2_d, @@ -1717,7 +1715,7 @@ impl PerceptionTrainer { unsafe { let mut launch = self.stream.launch_builder(&self.step_batched_fn); launch - .arg(&self.w_in_d).arg(&self.w_rec_d).arg(&self.b_d).arg(&self.tau_d) + .arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_d) .arg(&x_k_ptr).arg(&h_old_k_ptr) .arg(&dt_s).arg(&n_in_i).arg(&n_hid_i).arg(&n_batch_i) .arg(&h_new_k_ptr); @@ -1885,7 +1883,7 @@ impl PerceptionTrainer { unsafe { let mut launch = self.stream.launch_builder(&self.step_bwd_batched_fn); launch - .arg(&self.w_in_d).arg(&self.w_rec_d).arg(&self.b_d).arg(&self.tau_d) + .arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_d) .arg(&x_k_ptr).arg(&h_old_k_ptr).arg(&self.grad_h_new_d) .arg(&dt_s).arg(&n_in_i).arg(&n_hid_i).arg(&n_batch_i) .arg(&mut self.cfc_grad_w_in_scratch_d) @@ -2273,10 +2271,10 @@ impl PerceptionTrainer { // ── 9. Apply AdamW updates on all 17 param groups: CfC×4 + // GRN heads×10 + LN×2 + Mamba2 grouped. - self.opt_w_in.step(&mut self.w_in_d, &self.grad_w_in_d)?; - self.opt_w_rec.step(&mut self.w_rec_d, &self.grad_w_rec_d)?; - self.opt_b.step(&mut self.b_d, &self.grad_b_d)?; - self.opt_tau.step(&mut self.tau_d, &self.grad_tau_d)?; + self.opt_w_in.step(&mut self.trunk.w_in_d, &self.grad_w_in_d)?; + self.opt_w_rec.step(&mut self.trunk.w_rec_d, &self.grad_w_rec_d)?; + self.opt_b.step(&mut self.trunk.b_d, &self.grad_b_d)?; + self.opt_tau.step(&mut self.trunk.tau_d, &self.grad_tau_d)?; self.opt_heads_w1.step(&mut self.heads_w1_d, &self.grad_heads_w1_d)?; self.opt_heads_b1.step(&mut self.heads_b1_d, &self.grad_heads_b1_d)?; self.opt_heads_w2.step(&mut self.heads_w2_d, &self.grad_heads_w2_d)?; @@ -2661,7 +2659,7 @@ impl PerceptionTrainer { unsafe { let mut launch = self.stream.launch_builder(&self.step_batched_fn); launch - .arg(&self.w_in_d).arg(&self.w_rec_d).arg(&self.b_d).arg(&self.tau_d) + .arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_d) .arg(&x_k_ptr).arg(&h_old_k_ptr) .arg(&dt_s).arg(&n_in_i).arg(&n_hid_i).arg(&n_batch_i) .arg(&h_new_k_ptr);