gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(ml-alpha): CfcTrunk v2 weight skeleton (no callers yet)

Browse Source 
Adds zero-initialised v2 weight tensors (VSN, LN_a/b, attn_q, GRN heads)
and Option<Mamba2Block> slots for stacks 1 and 2 to CfcTrunk. Extends
CfcConfig with mamba2_state_dim (default 16, matches
PerceptionTrainerConfig::default). Imports HEAD_MID_DIM for GRN head
sizing.

No forward path changes — fields allocated, not yet read. Existing
new_random init for V1 fields (CfC + simple heads + projection) is
preserved unchanged so the trunk's forward kernels still produce the
same output.

Skeleton for X2..X9 weight-group migrations.

Verification:
- ml-alpha lib tests: 34 pass
- perception_forward_golden bit-equivalence: PASS (max_diff = 0.000000)

Per spec 2026-05-19-ml-alpha-v2-trunk-grows-and-deployability-design.md
§1.1, §2.2 (X1).
This commit is contained in:
jgrusewski
2026-05-19 01:23:05 +02:00
parent b47b2fabfb
commit e338000eec
1 changed files with 62 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
crates/ml-alpha/src/cfc/trunk.rs
View File

@@ -19,7 +19,7 @@ use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha8Rng;
use crate::cfc::snap_features::{Mbp10RawInput, ES_TICK_SIZE, FEATURE_DIM, REGIME_DIM};
use crate::heads::{HIDDEN_DIM, N_HORIZONS, PROJ_DIM};
use crate::heads::{HEAD_MID_DIM, HIDDEN_DIM, N_HORIZONS, PROJ_DIM};
use serde::{Deserialize, Serialize};
/// On-disk checkpoint envelope for CfcTrunk weights. Bumped version on
@@ -51,11 +51,15 @@ const PROJ_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/projection.c
pub struct CfcConfig {
pub n_in: usize,
pub n_hid: usize,
/// Mamba2 state dimension — used by the v2 trunk weight skeleton
/// (X1) and the per-stack Mamba2Block constructions added in X3/X5.
/// Default matches `PerceptionTrainerConfig::default().mamba2_state_dim`.
pub mamba2_state_dim: usize,
}
impl Default for CfcConfig {
fn default() -> Self {
Self { n_in: FEATURE_DIM, n_hid: HIDDEN_DIM }
Self { n_in: FEATURE_DIM, n_hid: HIDDEN_DIM, mamba2_state_dim: 16 }
}
}
@@ -73,7 +77,9 @@ pub struct CfcTrunk {
heads_fn: CudaFunction,
proj_fn: CudaFunction,
// Device-resident weights
// Device-resident weights — V1 (CfC + simple heads + projection).
// Preserved unchanged until X8/X9 consolidate against the v2 fields
// added below.
w_in_d: CudaSlice<f32>,
w_rec_d: CudaSlice<f32>,
b_d: CudaSlice<f32>,
@@ -85,6 +91,29 @@ pub struct CfcTrunk {
proj_g_d: CudaSlice<f32>,
proj_n_d: CudaSlice<f32>,
// === V2 weight skeleton (X1) — zero-initialised slabs + None for
// Mamba2 stacks. Each is wired in a subsequent commit (X2..X9):
// X2 VSN, X3 Mamba2_l1, X4 LN_a, X5 Mamba2_l2, X6 LN_b,
// X7 attn_pool, X8 CfC consolidation, X9 GRN heads.
// Public so PerceptionTrainer can read them during migration —
// visibility tightens to crate-private after X10 hoists the
// forward kernels into CfcTrunk methods. ===
pub vsn_w_d: CudaSlice<f32>,
pub vsn_b_d: CudaSlice<f32>,
pub mamba2_stack_1: Option<crate::mamba2_block::Mamba2Block>,
pub mamba2_stack_2: Option<crate::mamba2_block::Mamba2Block>,
pub ln_a_gain_d: CudaSlice<f32>,
pub ln_a_bias_d: CudaSlice<f32>,
pub ln_b_gain_d: CudaSlice<f32>,
pub ln_b_bias_d: CudaSlice<f32>,
pub attn_q_d: CudaSlice<f32>,
pub heads_w_gate_d: CudaSlice<f32>,
pub heads_b_gate_d: CudaSlice<f32>,
pub heads_w_main_d: CudaSlice<f32>,
pub heads_b_main_d: CudaSlice<f32>,
pub heads_w_skip_d: CudaSlice<f32>,
pub heads_b_skip_d: CudaSlice<f32>,
// Ping-pong hidden state
h_ping: CudaSlice<f32>,
h_pong: CudaSlice<f32>,
@@ -191,6 +220,36 @@ impl CfcTrunk {
snap_feat_d: stream.alloc_zeros::<f32>(FEATURE_DIM).context("snap_feat alloc")?,
probs_d: stream.alloc_zeros::<f32>(N_HORIZONS).context("probs alloc")?,
proj_out_d: stream.alloc_zeros::<f32>(PROJ_DIM).context("proj_out alloc")?,
// === V2 skeleton fields (X1) — allocated BEFORE `stream` is
// moved into the struct below. Populated by X2..X9. ===
vsn_w_d: stream.alloc_zeros::<f32>(cfg.n_in * cfg.n_in)
.context("v2 vsn_w alloc")?,
vsn_b_d: stream.alloc_zeros::<f32>(cfg.n_in)
.context("v2 vsn_b alloc")?,
mamba2_stack_1: None,
mamba2_stack_2: None,
ln_a_gain_d: stream.alloc_zeros::<f32>(cfg.n_hid)
.context("v2 ln_a_gain alloc")?,
ln_a_bias_d: stream.alloc_zeros::<f32>(cfg.n_hid)
.context("v2 ln_a_bias alloc")?,
ln_b_gain_d: stream.alloc_zeros::<f32>(cfg.n_hid)
.context("v2 ln_b_gain alloc")?,
ln_b_bias_d: stream.alloc_zeros::<f32>(cfg.n_hid)
.context("v2 ln_b_bias alloc")?,
attn_q_d: stream.alloc_zeros::<f32>(cfg.n_hid)
.context("v2 attn_q alloc")?,
heads_w_gate_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM)
.context("v2 heads_w_gate alloc")?,
heads_b_gate_d: stream.alloc_zeros::<f32>(N_HORIZONS)
.context("v2 heads_b_gate alloc")?,
heads_w_main_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM)
.context("v2 heads_w_main alloc")?,
heads_b_main_d: stream.alloc_zeros::<f32>(N_HORIZONS)
.context("v2 heads_b_main alloc")?,
heads_w_skip_d: stream.alloc_zeros::<f32>(N_HORIZONS * cfg.n_hid)
.context("v2 heads_w_skip alloc")?,
heads_b_skip_d: stream.alloc_zeros::<f32>(N_HORIZONS)
.context("v2 heads_b_skip alloc")?,
stream,
_snap_module: snap_module,
_step_module: step_module,