diff --git a/crates/ml-alpha/src/cfc/step.rs b/crates/ml-alpha/src/cfc/step.rs index afc9f034c..a25e14669 100644 --- a/crates/ml-alpha/src/cfc/step.rs +++ b/crates/ml-alpha/src/cfc/step.rs @@ -89,6 +89,114 @@ pub fn cfc_step_backward_gpu( )) } +/// Fused per-branch CfC forward step (Phase 2 dispatch). +/// +/// Per spec §5.4 point 1 (docs/superpowers/specs/2026-05-21-per-horizon-cfc-inference-design.md): +/// single launch covers all 5 branches × n_batch with grid=(B, N_HORIZONS, 1) +/// and block=(MAX_BUCKET_DIM=28, 1, 1). Uniform-predicate handles uneven +/// bucket sizes [25, 25, 25, 25, 28] without warp divergence. +/// +/// Buffer contract (matches kernel signature in +/// `crates/ml-alpha/cuda/cfc_step_per_branch.cu`): +/// - `w_in_d`, `w_rec_d`: shared `[HIDDEN_DIM × HIDDEN_DIM]` body weights. +/// - `b_d`: `[HIDDEN_DIM]` bias. +/// - `tau_all_d`: `[HIDDEN_DIM]` bucket-grouped τ (populated by transition). +/// - `x_d`, `h_old_d`: `[B × HIDDEN_DIM]` per-batch inputs. +/// - `bucket_channel_offset_d`: `[N_HORIZONS + 1]` cumulative bucket starts. +/// - `bucket_dim_k_d`: `[N_HORIZONS]` per-bucket channel counts. +/// - `h_new_d`: `[B × HIDDEN_DIM]` output (overwrite). +/// +/// The cooperative-staging shared-memory layout (`2 × HIDDEN_DIM` floats) +/// stages `x` and `h_old` rows once per block per +/// `pearl_cooperative_staging_eliminates_redundant_reads`. +pub fn cfc_step_per_branch_fwd_gpu( + stream: &Arc, + func: &cudarc::driver::CudaFunction, + w_in_d: &CudaSlice, + w_rec_d: &CudaSlice, + b_d: &CudaSlice, + tau_all_d: &CudaSlice, + x_d: &CudaSlice, + h_old_d: &CudaSlice, + dt_s: f32, + b_sz: i32, + bucket_channel_offset_d: &CudaSlice, + bucket_dim_k_d: &CudaSlice, + h_new_d: &mut CudaSlice, +) -> Result<()> { + const N_HORIZONS: u32 = 5; + const MAX_BUCKET_DIM: u32 = 28; + const HIDDEN_DIM: u32 = 128; + let cfg = LaunchConfig { + grid_dim: (b_sz as u32, N_HORIZONS, 1), + block_dim: (MAX_BUCKET_DIM, 1, 1), + // 2 × HIDDEN_DIM floats: x_local + h_old_local cooperative staging. + shared_mem_bytes: HIDDEN_DIM * 4 * 2, + }; + let mut launch = stream.launch_builder(func); + launch + .arg(w_in_d).arg(w_rec_d).arg(b_d).arg(tau_all_d) + .arg(x_d).arg(h_old_d).arg(&dt_s).arg(&b_sz) + .arg(bucket_channel_offset_d).arg(bucket_dim_k_d) + .arg(h_new_d); + unsafe { launch.launch(cfg).context("cfc_step_per_branch_fwd launch")?; } + Ok(()) +} + +/// Fused per-branch CfC backward step (Phase 2 dispatch). +/// +/// Mirror of `cfc_step_per_branch_fwd_gpu`. Per-batch grad slices are +/// written by this kernel; cross-batch reduction is the caller's +/// responsibility (matches existing `reduce_axis0` pattern in +/// `cfc_step_backward_batched`). Each thread (batch, branch, c) is the +/// sole writer of its slice — no atomicAdd per +/// `feedback_no_atomicadd`. +/// +/// Note: this binding is staged for Task 11's backward dispatch wiring. +/// Task 10 only calls the forward path; Task 11 will land the backward +/// caller. The kernel itself is already validated by the GPU oracle +/// test in `crates/ml-alpha/tests/cfc_step_per_branch.rs`. +#[allow(clippy::too_many_arguments)] +pub fn cfc_step_per_branch_bwd_gpu( + stream: &Arc, + func: &cudarc::driver::CudaFunction, + w_in_d: &CudaSlice, + w_rec_d: &CudaSlice, + b_d: &CudaSlice, + tau_all_d: &CudaSlice, + x_d: &CudaSlice, + h_old_d: &CudaSlice, + grad_h_new_d: &CudaSlice, + dt_s: f32, + b_sz: i32, + bucket_channel_offset_d: &CudaSlice, + bucket_dim_k_d: &CudaSlice, + grad_w_in_d: &mut CudaSlice, + grad_w_rec_d: &mut CudaSlice, + grad_b_d: &mut CudaSlice, + grad_tau_all_d: &mut CudaSlice, + grad_h_old_d: &mut CudaSlice, +) -> Result<()> { + const N_HORIZONS: u32 = 5; + const MAX_BUCKET_DIM: u32 = 28; + const HIDDEN_DIM: u32 = 128; + let cfg = LaunchConfig { + grid_dim: (b_sz as u32, N_HORIZONS, 1), + block_dim: (MAX_BUCKET_DIM, 1, 1), + shared_mem_bytes: HIDDEN_DIM * 4 * 2, + }; + let mut launch = stream.launch_builder(func); + launch + .arg(w_in_d).arg(w_rec_d).arg(b_d).arg(tau_all_d) + .arg(x_d).arg(h_old_d).arg(grad_h_new_d) + .arg(&dt_s).arg(&b_sz) + .arg(bucket_channel_offset_d).arg(bucket_dim_k_d) + .arg(grad_w_in_d).arg(grad_w_rec_d).arg(grad_b_d) + .arg(grad_tau_all_d).arg(grad_h_old_d); + unsafe { launch.launch(cfg).context("cfc_step_per_branch_bwd launch")?; } + Ok(()) +} + pub fn cfc_step_gpu( dev: &MlDevice, w: &CfcWeights, diff --git a/crates/ml-alpha/src/cfc/trunk.rs b/crates/ml-alpha/src/cfc/trunk.rs index 250e775db..bc6a471fa 100644 --- a/crates/ml-alpha/src/cfc/trunk.rs +++ b/crates/ml-alpha/src/cfc/trunk.rs @@ -80,6 +80,11 @@ const HEADS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_horiz const LAYER_NORM_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/layer_norm.cubin")); const VARIABLE_SELECTION_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/variable_selection.cubin")); const ATTENTION_POOL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/attention_pool.cubin")); +// Per-horizon CfC Phase 2 dispatch (spec §5.4 points 1–2). Modules +// stay alive on the trunk so the cached `CudaFunction` handles below +// remain valid through the trainer's lifetime. +const CFC_STEP_PER_BRANCH_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cfc_step_per_branch.cubin")); +const HEADS_BLOCK_DIAGONAL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/heads_block_diagonal_fwd.cubin")); #[derive(Clone, Debug)] pub struct CfcConfig { @@ -126,6 +131,10 @@ pub struct CfcTrunk { _ln_module: Arc, _vsn_module: Arc, _attn_module: Arc, + // Per-horizon CfC Phase 2 cubin modules. Kept alive alongside the + // cached function handles below. + _cfc_step_per_branch_module: Arc, + _heads_block_diagonal_module: Arc, // Forward kernel handles — read by PerceptionTrainer's evaluate_batched. pub vsn_fwd_fn: CudaFunction, pub ln_fwd_fn: CudaFunction, @@ -134,6 +143,19 @@ pub struct CfcTrunk { pub step_batched_fn: CudaFunction, pub heads_grn_fwd_fn: CudaFunction, pub transpose_3d_fn: CudaFunction, + // Per-horizon CfC Phase 2 dispatch handles (spec §5.4 points 1–2). + // `cfc_step_per_branch_fwd_fn` is read by `dispatch_train_step`'s + // Phase 2 branch (Task 10) and the eventual `forward_step_into` + // Phase 2 path (Task 13). The `_bwd_fn` handle is staged here so + // Task 11's backward wiring lands without re-touching the trunk + // constructor. `heads_block_diagonal_fwd_fn` is loaded for the + // future compact-w_skip heads dispatch; per Task 10's scope (see + // perception.rs Phase 2 dispatch comment), it is NOT consumed yet + // — heads continue through the existing GRN kernel until a later + // task wires the compact-storage path through the GRN's skip term. + pub cfc_step_per_branch_fwd_fn: CudaFunction, + pub cfc_step_per_branch_bwd_fn: CudaFunction, + pub heads_block_diagonal_fwd_fn: CudaFunction, // CfC weights (cfc_n_in = HIDDEN_DIM). Public so PerceptionTrainer can // mutate during init + AdamW. @@ -192,6 +214,13 @@ impl CfcTrunk { let ln_module = ctx.load_cubin(LAYER_NORM_CUBIN.to_vec()).context("load ln cubin")?; let vsn_module = ctx.load_cubin(VARIABLE_SELECTION_CUBIN.to_vec()).context("load vsn cubin")?; let attn_module = ctx.load_cubin(ATTENTION_POOL_CUBIN.to_vec()).context("load attn cubin")?; + // Per-horizon CfC Phase 2 cubin modules. + let cfc_step_per_branch_module = ctx + .load_cubin(CFC_STEP_PER_BRANCH_CUBIN.to_vec()) + .context("load cfc_step_per_branch cubin")?; + let heads_block_diagonal_module = ctx + .load_cubin(HEADS_BLOCK_DIAGONAL_CUBIN.to_vec()) + .context("load heads_block_diagonal cubin")?; let vsn_fwd_fn = vsn_module.load_function("variable_selection_fwd").context("vsn_fwd_fn")?; let ln_fwd_fn = ln_module.load_function("layer_norm_fwd").context("ln_fwd_fn")?; let attn_fwd_fn = attn_module.load_function("attention_pool_fwd").context("attn_fwd_fn")?; @@ -199,6 +228,15 @@ impl CfcTrunk { let step_batched_fn = step_module.load_function("cfc_step_batched").context("step_batched_fn")?; let heads_grn_fwd_fn = heads_module.load_function("multi_horizon_heads_grn_fwd_batched").context("heads_grn_fwd_fn")?; let transpose_3d_fn = step_module.load_function("transpose_3d_swap_01").context("transpose_3d_fn")?; + let cfc_step_per_branch_fwd_fn = cfc_step_per_branch_module + .load_function("cfc_step_per_branch_fwd") + .context("cfc_step_per_branch_fwd_fn")?; + let cfc_step_per_branch_bwd_fn = cfc_step_per_branch_module + .load_function("cfc_step_per_branch_bwd") + .context("cfc_step_per_branch_bwd_fn")?; + let heads_block_diagonal_fwd_fn = heads_block_diagonal_module + .load_function("heads_block_diagonal_fwd") + .context("heads_block_diagonal_fwd_fn")?; // CfC weight init via a seeded ChaCha8Rng — the same chain feeds // every weight tensor so cfg.seed fully determines the initial @@ -317,6 +355,8 @@ impl CfcTrunk { _ln_module: ln_module, _vsn_module: vsn_module, _attn_module: attn_module, + _cfc_step_per_branch_module: cfc_step_per_branch_module, + _heads_block_diagonal_module: heads_block_diagonal_module, vsn_fwd_fn, ln_fwd_fn, attn_fwd_fn, @@ -324,6 +364,9 @@ impl CfcTrunk { step_batched_fn, heads_grn_fwd_fn, transpose_3d_fn, + cfc_step_per_branch_fwd_fn, + cfc_step_per_branch_bwd_fn, + heads_block_diagonal_fwd_fn, }) } diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index b74bb353b..fcc853504 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -2412,19 +2412,72 @@ impl PerceptionTrainer { let main_k_ptr = main_base + (k * kb_nh_bytes) as u64; let logit_k_ptr = logit_base + (k * kb_nh_bytes) as u64; - // cfc_step_batched(W, x[B*n_in], h_old[B*n_hid], …, → h_new[B*n_hid]). + // CfC forward dispatch, two-phase per spec §2.3 + Task 10: + // + // Phase 1 (warmup): existing single shared `cfc_step_batched` + // kernel — block-per-batch over all HIDDEN_DIM channels. + // + // Phase 2 (routed): fused per-branch kernel + // `cfc_step_per_branch_fwd` — grid=(B, N_HORIZONS=5, 1), + // block=(MAX_BUCKET_DIM=28, 1, 1) uniform predicate, with + // `bucket_channel_offset_d` / `bucket_dim_k_d` routing each + // block to its quintile of HIDDEN_DIM. Spec §5.4 point 1. + // + // The branch on `self.phase` lives OUTSIDE the captured graph's + // hot kernel scope: the captured train graph is recaptured at + // the Phase 1→2 transition (cf. `pearl_no_host_branches_in_captured_graph`). + // Once Phase 2 has been entered, every subsequent capture replays + // the routed path; the if-branch here is evaluated by host code + // building the graph, not inside the device-side kernel stream. 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(&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); - launch.launch(cfg_cfc).context("cfc fwd k batched")?; + match self.phase { + TrainingPhase::Phase1Warmup => { + 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(&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); + launch.launch(cfg_cfc).context("cfc fwd k batched")?; + } + TrainingPhase::Phase2Routed => { + // grid=(B, N_HORIZONS, 1), block=(MAX_BUCKET_DIM, 1, 1). + // Shared mem = 2 × HIDDEN_DIM floats (x_local + h_old_local + // cooperative staging per pearl_cooperative_staging_eliminates_redundant_reads). + const N_HORIZONS_U32: u32 = N_HORIZONS as u32; + const MAX_BUCKET_DIM: u32 = 28; + let cfg_cfc_per_branch = LaunchConfig { + grid_dim: (b_sz as u32, N_HORIZONS_U32, 1), + block_dim: (MAX_BUCKET_DIM, 1, 1), + shared_mem_bytes: (2 * HIDDEN_DIM * std::mem::size_of::()) as u32, + }; + let mut launch = self.stream.launch_builder(&self.trunk.cfc_step_per_branch_fwd_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(&x_k_ptr).arg(&h_old_k_ptr) + .arg(&dt_s).arg(&n_batch_i) + .arg(&self.trunk.bucket_channel_offset_d).arg(&self.trunk.bucket_dim_k_d) + .arg(&h_new_k_ptr); + launch.launch(cfg_cfc_per_branch).context("cfc_step_per_branch_fwd k")?; + } + } } // multi_horizon_heads_grn_fwd_batched(10 params + h[B*128] + n_batch // → probs[B*5] + z1/a1/z2 [B,5,HEAD_MID] + gate_logit/main/logit [B,5]). + // + // Task 10 scope note (Option B): the heads dispatch stays at the + // existing full-`heads_w_skip_d` GRN kernel for BOTH phases. The + // Phase 2 compact-w_skip path (`heads_block_diagonal_fwd`) and + // its 128-float `heads_w_skip_compact_d` storage are populated + // by the transition kernels (Task 9) but NOT yet consumed here: + // the GRN kernel computes the full per-horizon logit including + // skip + gate + main sub-components, and routing only the skip + // term through the compact path requires integration with + // GRN's main/gate computations (out of Task 10 scope per + // `feedback_no_partial_refactor`). A follow-up task wires the + // compact heads dispatch through GRN; the trunk handle + // `heads_block_diagonal_fwd_fn` is already loaded for that. unsafe { let mut launch = self.stream.launch_builder(&self.trunk.heads_grn_fwd_fn); launch