diff --git a/crates/ml-alpha/src/cfc/trunk.rs b/crates/ml-alpha/src/cfc/trunk.rs index 3292580fe..5abc32f3a 100644 --- a/crates/ml-alpha/src/cfc/trunk.rs +++ b/crates/ml-alpha/src/cfc/trunk.rs @@ -46,6 +46,12 @@ const SNAP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/snap_feature const STEP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cfc_step.cubin")); const HEADS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_horizon_heads.cubin")); const PROJ_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/projection.cubin")); +// X10: v2 forward kernel cubins — loaded into the trunk so it owns +// every kernel handle the v2 forward chain needs. Trainer will read +// these handles via `self.trunk.` in a subsequent commit (X10b). +const LAYER_NORM_CUBIN_V2: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/layer_norm.cubin")); +const VARIABLE_SELECTION_CUBIN_V2: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/variable_selection.cubin")); +const ATTENTION_POOL_CUBIN_V2: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/attention_pool.cubin")); #[derive(Clone, Debug)] pub struct CfcConfig { @@ -87,6 +93,19 @@ pub struct CfcTrunk { step_fn: CudaFunction, heads_fn: CudaFunction, proj_fn: CudaFunction, + // X10: v2 forward function handles. Loaded from corresponding cubins + // in new_random; consumed by X10b (PerceptionTrainer reads through + // self.trunk.) and X11 (full v2 capture_graph_a). + _ln_module_v2: Arc, + _vsn_module_v2: Arc, + _attn_module_v2: Arc, + pub vsn_fwd_fn: CudaFunction, + pub ln_fwd_fn: CudaFunction, + pub attn_fwd_fn: CudaFunction, + pub snap_batched_fn: CudaFunction, + pub step_batched_fn: CudaFunction, + pub heads_grn_fwd_fn: CudaFunction, + pub transpose_3d_fn: CudaFunction, // CfC weights — now v2-shaped (cfc_n_in = HIDDEN_DIM). Public so // PerceptionTrainer can mutate during init + AdamW. @@ -183,6 +202,18 @@ impl CfcTrunk { let heads_fn = heads_module.load_function("multi_horizon_heads").context("heads_fn")?; let proj_fn = proj_module.load_function("projection_kernel").context("proj_fn")?; + // X10: load v2 forward kernels onto the trunk. + let ln_module_v2 = ctx.load_cubin(LAYER_NORM_CUBIN_V2.to_vec()).context("load ln cubin")?; + let vsn_module_v2 = ctx.load_cubin(VARIABLE_SELECTION_CUBIN_V2.to_vec()).context("load vsn cubin")?; + let attn_module_v2 = ctx.load_cubin(ATTENTION_POOL_CUBIN_V2.to_vec()).context("load attn cubin")?; + let vsn_fwd_fn = vsn_module_v2.load_function("variable_selection_fwd").context("vsn_fwd_fn")?; + let ln_fwd_fn = ln_module_v2.load_function("layer_norm_fwd").context("ln_fwd_fn")?; + let attn_fwd_fn = attn_module_v2.load_function("attention_pool_fwd").context("attn_fwd_fn")?; + let snap_batched_fn = snap_module.load_function("snap_feature_assemble_batched").context("snap_batched_fn")?; + 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")?; + // Random init. let mut r = ChaCha8Rng::seed_from_u64(seed); // X8: CfC layer weights now use cfc_n_in (V2 = HIDDEN_DIM), not @@ -290,6 +321,16 @@ impl CfcTrunk { step_fn, heads_fn, proj_fn, + _ln_module_v2: ln_module_v2, + _vsn_module_v2: vsn_module_v2, + _attn_module_v2: attn_module_v2, + vsn_fwd_fn, + ln_fwd_fn, + attn_fwd_fn, + snap_batched_fn, + step_batched_fn, + heads_grn_fwd_fn, + transpose_3d_fn, stg_bid_px, stg_bid_sz, stg_ask_px,