diff --git a/crates/ml-alpha/src/cfc/trunk.rs b/crates/ml-alpha/src/cfc/trunk.rs index 9bc3d2940..250e775db 100644 --- a/crates/ml-alpha/src/cfc/trunk.rs +++ b/crates/ml-alpha/src/cfc/trunk.rs @@ -57,7 +57,12 @@ struct Checkpoint { // Attention pool attn_q: Vec, // CfC - w_in: Vec, w_rec: Vec, b: Vec, tau: Vec, + w_in: Vec, w_rec: Vec, b: Vec, tau_all: Vec, + // Bucket metadata (populated at Phase 1→2 transition; zero-initialized in new_random) + bucket_id_per_channel: Vec, // [HIDDEN_DIM] + bucket_channel_offset: Vec, // [N_HORIZONS + 1] + bucket_dim_k: Vec, // [N_HORIZONS] + heads_w_skip_offset: Vec, // [N_HORIZONS + 1] // Heads (GRN structure × 5 horizons) heads_w1: Vec, heads_b1: Vec, heads_w2: Vec, heads_b2: Vec, @@ -135,7 +140,14 @@ pub struct CfcTrunk { pub w_in_d: CudaSlice, pub w_rec_d: CudaSlice, pub b_d: CudaSlice, - pub tau_d: CudaSlice, + pub tau_all_d: CudaSlice, + + // Bucket metadata for per-horizon CfC routing (Phase 1→2 transition + // populates these; zero-initialized at `new_random`). See spec §2.2. + pub bucket_id_per_channel_d: CudaSlice, + pub bucket_channel_offset_d: CudaSlice, + pub bucket_dim_k_d: CudaSlice, + pub heads_w_skip_offset_d: CudaSlice, // Inference weights for the full perception chain: // snap → vsn → mamba2_l1 → ln_a → mamba2_l2 → ln_b → attn_pool → @@ -217,12 +229,34 @@ impl CfcTrunk { Ok(buf) }; + // Bucket metadata: zero-initialised here; populated by the Phase + // 1→2 transition orchestrator in `cfc/bucket_routing.rs` once + // Controller A fires. Allocated unconditionally so that + // checkpoint save/load round-trips include them even before the + // transition has run. + let bucket_id_per_channel_d = stream + .alloc_zeros::(HIDDEN_DIM) + .context("bucket_id_per_channel alloc")?; + let bucket_channel_offset_d = stream + .alloc_zeros::(N_HORIZONS + 1) + .context("bucket_channel_offset alloc")?; + let bucket_dim_k_d = stream + .alloc_zeros::(N_HORIZONS) + .context("bucket_dim_k alloc")?; + let heads_w_skip_offset_d = stream + .alloc_zeros::(N_HORIZONS + 1) + .context("heads_w_skip_offset alloc")?; + Ok(Self { cfg: cfg.clone(), w_in_d: upload(&w_in)?, w_rec_d: upload(&w_rec)?, b_d: upload(&b)?, - tau_d: upload(&tau)?, + tau_all_d: upload(&tau)?, + bucket_id_per_channel_d, + bucket_channel_offset_d, + bucket_dim_k_d, + heads_w_skip_offset_d, // Inference layers — zero-initialised here; PerceptionTrainer's // construction overwrites them with its own seeded init via // memcpy_htod. Standalone trunk callers (tests, fxt-backtest's @@ -329,6 +363,16 @@ impl CfcTrunk { self.stream.memcpy_dtoh(slice, v.as_mut_slice())?; Ok(v) }; + let download_u8 = |slice: &CudaSlice| -> Result> { + let mut v = vec![0u8; slice.len()]; + self.stream.memcpy_dtoh(slice, v.as_mut_slice())?; + Ok(v) + }; + let download_u32 = |slice: &CudaSlice| -> Result> { + let mut v = vec![0u32; slice.len()]; + self.stream.memcpy_dtoh(slice, v.as_mut_slice())?; + Ok(v) + }; let m1 = self.mamba2_l1(); let m2 = self.mamba2_l2(); let ckpt = Checkpoint { @@ -367,7 +411,11 @@ impl CfcTrunk { w_in: download(&self.w_in_d)?, w_rec: download(&self.w_rec_d)?, b: download(&self.b_d)?, - tau: download(&self.tau_d)?, + tau_all: download(&self.tau_all_d)?, + bucket_id_per_channel: download_u8(&self.bucket_id_per_channel_d)?, + bucket_channel_offset: download_u32(&self.bucket_channel_offset_d)?, + bucket_dim_k: download_u32(&self.bucket_dim_k_d)?, + heads_w_skip_offset: download_u32(&self.heads_w_skip_offset_d)?, heads_w1: download(&self.heads_w1_d)?, heads_b1: download(&self.heads_b1_d)?, heads_w2: download(&self.heads_w2_d)?, @@ -420,6 +468,22 @@ impl CfcTrunk { stream.memcpy_htod(src, dst)?; Ok(()) }; + let upload_u8 = |dst: &mut CudaSlice, src: &[u8], name: &str| -> Result<()> { + anyhow::ensure!( + dst.len() == src.len(), + "size mismatch loading {name}: device {} != file {}", dst.len(), src.len() + ); + stream.memcpy_htod(src, dst)?; + Ok(()) + }; + let upload_u32 = |dst: &mut CudaSlice, src: &[u32], name: &str| -> Result<()> { + anyhow::ensure!( + dst.len() == src.len(), + "size mismatch loading {name}: device {} != file {}", dst.len(), src.len() + ); + stream.memcpy_htod(src, dst)?; + Ok(()) + }; upload(&mut trunk.vsn_w_d, &ckpt.vsn_w, "vsn_w")?; upload(&mut trunk.vsn_b_d, &ckpt.vsn_b, "vsn_b")?; { @@ -454,7 +518,11 @@ impl CfcTrunk { upload(&mut trunk.w_in_d, &ckpt.w_in, "cfc.w_in")?; upload(&mut trunk.w_rec_d, &ckpt.w_rec, "cfc.w_rec")?; upload(&mut trunk.b_d, &ckpt.b, "cfc.b")?; - upload(&mut trunk.tau_d, &ckpt.tau, "cfc.tau")?; + upload(&mut trunk.tau_all_d, &ckpt.tau_all, "cfc.tau_all")?; + upload_u8(&mut trunk.bucket_id_per_channel_d, &ckpt.bucket_id_per_channel, "bucket_id_per_channel")?; + upload_u32(&mut trunk.bucket_channel_offset_d, &ckpt.bucket_channel_offset, "bucket_channel_offset")?; + upload_u32(&mut trunk.bucket_dim_k_d, &ckpt.bucket_dim_k, "bucket_dim_k")?; + upload_u32(&mut trunk.heads_w_skip_offset_d, &ckpt.heads_w_skip_offset, "heads_w_skip_offset")?; upload(&mut trunk.heads_w1_d, &ckpt.heads_w1, "heads_w1")?; upload(&mut trunk.heads_b1_d, &ckpt.heads_b1, "heads_b1")?; upload(&mut trunk.heads_w2_d, &ckpt.heads_w2, "heads_w2")?; @@ -509,10 +577,27 @@ mod checkpoint_tests { trunk.stream.memcpy_dtoh(slice, v.as_mut_slice()).unwrap(); v }; + let read_u8 = |trunk: &CfcTrunk, slice: &CudaSlice| -> Vec { + let mut v = vec![0u8; slice.len()]; + trunk.stream.memcpy_dtoh(slice, v.as_mut_slice()).unwrap(); + v + }; + let read_u32 = |trunk: &CfcTrunk, slice: &CudaSlice| -> Vec { + let mut v = vec![0u32; slice.len()]; + trunk.stream.memcpy_dtoh(slice, v.as_mut_slice()).unwrap(); + v + }; assert_eq!(read(&trunk_a, &trunk_a.w_in_d), read(&trunk_b, &trunk_b.w_in_d)); assert_eq!(read(&trunk_a, &trunk_a.w_rec_d), read(&trunk_b, &trunk_b.w_rec_d)); assert_eq!(read(&trunk_a, &trunk_a.b_d), read(&trunk_b, &trunk_b.b_d)); - assert_eq!(read(&trunk_a, &trunk_a.tau_d), read(&trunk_b, &trunk_b.tau_d)); + assert_eq!(read(&trunk_a, &trunk_a.tau_all_d), read(&trunk_b, &trunk_b.tau_all_d)); + // Bucket metadata: zero-initialised in new_random; round-trips must + // preserve them bit-identically (trivially zero == zero pre-transition, + // but the producer/consumer contract is verified here). + assert_eq!(read_u8(&trunk_a, &trunk_a.bucket_id_per_channel_d), read_u8(&trunk_b, &trunk_b.bucket_id_per_channel_d)); + assert_eq!(read_u32(&trunk_a, &trunk_a.bucket_channel_offset_d), read_u32(&trunk_b, &trunk_b.bucket_channel_offset_d)); + assert_eq!(read_u32(&trunk_a, &trunk_a.bucket_dim_k_d), read_u32(&trunk_b, &trunk_b.bucket_dim_k_d)); + assert_eq!(read_u32(&trunk_a, &trunk_a.heads_w_skip_offset_d), read_u32(&trunk_b, &trunk_b.heads_w_skip_offset_d)); assert_eq!(read(&trunk_a, &trunk_a.vsn_w_d), read(&trunk_b, &trunk_b.vsn_w_d)); assert_eq!(read(&trunk_a, &trunk_a.ln_a_gain_d), read(&trunk_b, &trunk_b.ln_a_gain_d)); assert_eq!(read(&trunk_a, &trunk_a.attn_q_d), read(&trunk_b, &trunk_b.attn_q_d)); diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 0604b3b95..107e29546 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -213,7 +213,7 @@ pub struct PerceptionTrainer { // CfC + heads weights + their AdamWs (6 groups — tau is trained now). // 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 + // `self.trunk.w_in_d` / `w_rec_d` / `b_d` / `tau_all_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 @@ -840,7 +840,7 @@ impl PerceptionTrainer { 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")?; + stream.memcpy_htod(&tau, &mut trunk.tau_all_d).context("trunk.tau_all_d upload")?; // X9: upload trainer's GRN heads into the trunk's slots. stream.memcpy_htod(&heads_w1, &mut trunk.heads_w1_d).context("trunk.heads_w1 upload")?; stream.memcpy_htod(&heads_b1, &mut trunk.heads_b1_d).context("trunk.heads_b1 upload")?; @@ -2095,7 +2095,7 @@ impl PerceptionTrainer { 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_d) + .arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_all_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); @@ -2345,7 +2345,7 @@ impl PerceptionTrainer { unsafe { let mut launch = self.stream.launch_builder(&self.step_bwd_batched_fn); launch - .arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_d) + .arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_all_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) @@ -2736,7 +2736,7 @@ impl PerceptionTrainer { 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_tau.step(&mut self.trunk.tau_all_d, &self.grad_tau_d)?; self.opt_heads_w1.step(&mut self.trunk.heads_w1_d, &self.grad_heads_w1_d)?; self.opt_heads_b1.step(&mut self.trunk.heads_b1_d, &self.grad_heads_b1_d)?; self.opt_heads_w2.step(&mut self.trunk.heads_w2_d, &self.grad_heads_w2_d)?; @@ -3205,7 +3205,7 @@ impl PerceptionTrainer { 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_d) + .arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_all_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); @@ -3480,7 +3480,7 @@ impl PerceptionTrainer { 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_d) + .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); @@ -4118,7 +4118,7 @@ impl PerceptionTrainer { 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_d) + .arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_all_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);