diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index b0f9bd2a4..b83a4aeb6 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -2196,12 +2196,76 @@ impl PerceptionTrainer { } } - // 5. Persist routing metadata in trainer state. Tasks + // 5. Sync trunk's bucket metadata fields from the trainer-owned + // BucketRoutingMetadata. Without this, training-time Phase 2 + // dispatch sites that read `self.trunk.bucket_*_d` + // (cfc_step_per_branch_fwd at the per-branch dispatch and + // h_mag_per_bucket_kernel for Controller D) would see the + // zero-initialized trunk buffers — the per-branch kernel's + // uniform predicate `tid >= bucket_dim_k[branch]` would + // early-return ALL threads → silent Phase 2 no-op. + // + // Trunk fields are the single source of truth for BOTH + // training-time dispatch AND checkpoint serialization. Cost: + // 4 DtoD memcpys, ~256 bytes total at the one-shot + // transition (off the hot path). Sync MUST happen BEFORE + // the `Some(metadata)` move below (borrow of `metadata.*_d` + // requires `metadata` to still be live), and BEFORE the + // heads_w_skip_mask_init kernel at step 5.5 (which still + // reads `metadata.bucket_id_per_channel_d` directly through + // `self.bucket_routing_metadata.as_ref()` after the move). + // Per `pearl_canary_input_freshness_launch_order`: producer + // (sync) precedes every consumer. + unsafe { + let s = self.stream.cu_stream(); + // bucket_channel_offset_d: [N_HORIZONS + 1] × u32 + let (src, _gs) = metadata.bucket_channel_offset_d.device_ptr(&self.stream); + let (dst, _gd) = self.trunk.bucket_channel_offset_d.device_ptr_mut(&self.stream); + cudarc::driver::result::memcpy_dtod_async( + dst, + src, + (N_HORIZONS + 1) * std::mem::size_of::(), + s, + ) + .context("sync bucket_channel_offset to trunk")?; + // bucket_dim_k_d: [N_HORIZONS] × u32 + let (src, _gs) = metadata.bucket_dim_k_d.device_ptr(&self.stream); + let (dst, _gd) = self.trunk.bucket_dim_k_d.device_ptr_mut(&self.stream); + cudarc::driver::result::memcpy_dtod_async( + dst, + src, + N_HORIZONS * std::mem::size_of::(), + s, + ) + .context("sync bucket_dim_k to trunk")?; + // bucket_id_per_channel_d: [HIDDEN_DIM] × u8 + let (src, _gs) = metadata.bucket_id_per_channel_d.device_ptr(&self.stream); + let (dst, _gd) = self.trunk.bucket_id_per_channel_d.device_ptr_mut(&self.stream); + cudarc::driver::result::memcpy_dtod_async( + dst, + src, + HIDDEN_DIM * std::mem::size_of::(), + s, + ) + .context("sync bucket_id_per_channel to trunk")?; + // heads_w_skip_offset_d: [N_HORIZONS + 1] × u32 + let (src, _gs) = metadata.heads_w_skip_offset_d.device_ptr(&self.stream); + let (dst, _gd) = self.trunk.heads_w_skip_offset_d.device_ptr_mut(&self.stream); + cudarc::driver::result::memcpy_dtod_async( + dst, + src, + (N_HORIZONS + 1) * std::mem::size_of::(), + s, + ) + .context("sync heads_w_skip_offset to trunk")?; + } + + // 6. Persist routing metadata in trainer state. Tasks // 10–12 read its slices to drive Phase 2 dispatch + // Controllers B / C / D. self.bucket_routing_metadata = Some(metadata); - // 5.5. Block-diagonal heads grad-mask init (follow-up to + // 7. Block-diagonal heads grad-mask init (follow-up to // Task 10 Option B). One-shot: build the routing // mask AND zero-multiply off-bucket entries of // `heads_w_skip_d` in place. From here on, the per- @@ -2236,7 +2300,7 @@ impl PerceptionTrainer { } } - // 6. Latch trainer into Phase 2. + // 8. Latch trainer into Phase 2. // // NOTE: Task 9 leaves the per-step dispatch path // unchanged (Phase 1 single-CfC, 640-float heads_w_skip).