feat(ml-alpha): CfcTrunk loads v2 forward kernel handles (X10 foundation)
Adds the v2 forward kernel cubins (layer_norm, variable_selection, attention_pool) and function handles (vsn_fwd_fn, ln_fwd_fn, attn_fwd_fn, snap_batched_fn, step_batched_fn, heads_grn_fwd_fn, transpose_3d_fn) onto CfcTrunk. The trunk now owns every kernel handle the v2 forward chain needs. PerceptionTrainer still loads its own copies of these cubins (duplicate loading) and uses its own handles in evaluate_batched / step_batched — the trainer-side consolidation lands in X10b. This commit is the foundation: X11 will build capture_graph_a using these trunk-owned handles, and X12 (CheckpointV2) doesn't depend on the consolidation. Verification: - perception_forward_golden: PASS (max_diff = 0.000000) - ml-alpha lib tests: 34 pass Per spec §2.2 (X10).
This commit is contained in:
@@ -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.<fn>` 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.<fn>) and X11 (full v2 capture_graph_a).
|
||||
_ln_module_v2: Arc<CudaModule>,
|
||||
_vsn_module_v2: Arc<CudaModule>,
|
||||
_attn_module_v2: Arc<CudaModule>,
|
||||
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,
|
||||
|
||||
Reference in New Issue
Block a user