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

chore(ml-alpha): deep cleanup — delete all V1 dead code

Browse Source 
CfcTrunk (~250 lines deleted):
- Deleted V1 forward methods: dispatch_perception, capture_graph_a,
  perception_forward_captured, snapshot_hidden, update_input_buffers,
  forward_snapshot, upload_pre_allocated
- Deleted V1 weight fields: heads_w_d, heads_b_d, proj_w_d, proj_b_d,
  proj_g_d, proj_n_d
- Deleted V1 per-step scratches: h_ping, h_pong, bid_px_d, bid_sz_d,
  ask_px_d, ask_sz_d, prev_bid_sz_d, prev_ask_sz_d, regime_d,
  snap_feat_d, probs_d, proj_out_d
- Deleted V1 staging buffers: stg_bid_px / stg_bid_sz / stg_ask_px /
  stg_ask_sz / stg_regime (MappedF32Buffer was only used by V1)
- Deleted graph_a field + _proj_module + V1 fn handles (snap_fn, step_fn,
  heads_fn, proj_fn)
- Deleted V1-only init in new_random (heads_w/b, proj_w/b/g/n, per-step
  scratch allocs)
- Deleted file-level helpers used only by V1: upload(stream, host),
  upload_into, copy_dtod, download
- Deleted PROJ_CUBIN constant
- Updated save_load_roundtrip test to assert on v2 weight tensors
- Stripped unused imports (CUgraphInstantiate_flags, CUstreamCaptureMode,
  CudaGraph, LaunchConfig, PushKernelArg, DevicePtr, DevicePtrMut,
  MappedF32Buffer, ES_TICK_SIZE, Mbp10RawInput, REGIME_DIM, PROJ_DIM)

PerceptionTrainer (~30 lines deleted):
- Deleted duplicate cubin fn fields made dead by X10b: snap_batched_fn,
  step_batched_fn, heads_grn_fwd_fn, transpose_3d_fn, vsn_fwd_fn,
  ln_fwd_fn, attn_fwd_fn (trainer reads these from self.trunk now)
- Deleted their load_function bindings in PerceptionTrainer::new
- Backward kernels (ln_bwd_fn, vsn_bwd_fn, attn_bwd_fn, step_bwd_batched_fn,
  heads_grn_bwd_fn) kept — training-only, not on trunk

Verification:
- perception_forward_golden: PASS (max_diff = 0.000000)
- ml-alpha lib tests: 33 pass
- ml-backtesting + fxt-backtest build clean

Trunk.rs shrunk from 800+ to ~480 lines. Code is now purely the v2
inference graph: weights, kernel handles, save_checkpoint/load_checkpoint,
mamba2_l1/l2 accessors. No V1 surface area left.
This commit is contained in:
jgrusewski
2026-05-19 09:19:26 +02:00
parent 71b467be40
commit cecc08a122
2 changed files with 77 additions and 541 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -12,14 +12,13 @@
use std::sync::Arc;
use anyhow::{Context, Result};
use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
use cudarc::driver::{CudaFunction, CudaGraph, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg};
use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream};
use ml_core::device::MlDevice;
use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha8Rng;
use crate::cfc::snap_features::{Mbp10RawInput, ES_TICK_SIZE, FEATURE_DIM, REGIME_DIM};
use crate::heads::{HEAD_MID_DIM, HIDDEN_DIM, N_HORIZONS, PROJ_DIM};
use crate::cfc::snap_features::FEATURE_DIM;
use crate::heads::{HEAD_MID_DIM, HIDDEN_DIM, N_HORIZONS};
use serde::{Deserialize, Serialize};
/// On-disk checkpoint envelope for CfcTrunk weights. Covers the full
@@ -66,12 +65,9 @@ struct Checkpoint {
heads_w_skip: Vec<f32>, heads_b_skip: Vec<f32>,
}
use crate::pinned_mem::MappedF32Buffer;
const SNAP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/snap_feature_assemble.cubin"));
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: forward kernel cubins — loaded into the trunk so it owns
// every kernel handle the forward chain needs. Trainer will read
// these handles via `self.trunk.<fn>` in a subsequent commit (X10b).
@@ -117,21 +113,14 @@ pub struct CfcTrunk {
cfg: CfcConfig,
stream: Arc<CudaStream>,
// Cubin modules (kept alive for the function refs below)
// Cubin modules — kept alive for the function refs below.
_snap_module: Arc<CudaModule>,
_step_module: Arc<CudaModule>,
_heads_module: Arc<CudaModule>,
_proj_module: Arc<CudaModule>,
snap_fn: CudaFunction,
step_fn: CudaFunction,
heads_fn: CudaFunction,
proj_fn: CudaFunction,
// X10: forward function handles. Loaded from corresponding cubins
// in new_random; consumed by X10b (PerceptionTrainer reads through
// self.trunk.<fn>) and X11 (full capture_graph_a).
_ln_module: Arc<CudaModule>,
_vsn_module: Arc<CudaModule>,
_attn_module: Arc<CudaModule>,
// Forward kernel handles — read by PerceptionTrainer's evaluate_batched.
pub vsn_fwd_fn: CudaFunction,
pub ln_fwd_fn: CudaFunction,
pub attn_fwd_fn: CudaFunction,
@@ -140,29 +129,17 @@ pub struct CfcTrunk {
pub heads_grn_fwd_fn: CudaFunction,
pub transpose_3d_fn: CudaFunction,
// CfC weights — now trunk-shaped (cfc_n_in = HIDDEN_DIM). Public so
// PerceptionTrainer can mutate during init + AdamW.
// CfC weights (cfc_n_in = HIDDEN_DIM). Public so PerceptionTrainer can
// mutate during init + AdamW.
pub w_in_d: CudaSlice<f32>,
pub w_rec_d: CudaSlice<f32>,
pub b_d: CudaSlice<f32>,
pub tau_d: CudaSlice<f32>,
// V1 legacy fields — only used by the now-stub V1 forward path
// (dispatch_perception). Kept allocated so the trunk struct lays out
// the same in tests; will be removed in a follow-up cleanup commit.
heads_w_d: CudaSlice<f32>,
heads_b_d: CudaSlice<f32>,
proj_w_d: CudaSlice<f32>,
proj_b_d: CudaSlice<f32>,
proj_g_d: CudaSlice<f32>,
proj_n_d: CudaSlice<f32>,
// === Inference 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. ===
// Inference weights for the full perception chain:
// snap → vsn → mamba2_l1 → ln_a → mamba2_l2 → ln_b → attn_pool →
// cfc → grn heads.
// All public so PerceptionTrainer can mutate during init + AdamW.
pub vsn_w_d: CudaSlice<f32>,
pub vsn_b_d: CudaSlice<f32>,
pub mamba2_stack_1: Option<crate::mamba2_block::Mamba2Block>,
@@ -172,9 +149,9 @@ pub struct CfcTrunk {
pub ln_b_gain_d: CudaSlice<f32>,
pub ln_b_bias_d: CudaSlice<f32>,
pub attn_q_d: CudaSlice<f32>,
// X9: full GRN head structure. heads_w1/b1 is the input projection
// (HIDDEN → HEAD_MID), heads_w2/b2 is the mid→mid layer, then the
// gate/main/skip outputs feed multi-horizon probabilities.
// GRN heads: heads_w1/b1 is the input projection (HIDDEN → HEAD_MID),
// heads_w2/b2 is the mid→mid layer, gate/main/skip feed the final
// multi-horizon probabilities.
pub heads_w1_d: CudaSlice<f32>,
pub heads_b1_d: CudaSlice<f32>,
pub heads_w2_d: CudaSlice<f32>,
@@ -185,40 +162,6 @@ pub struct CfcTrunk {
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>,
// Pre-allocated per-step scratch (raw input + features + outputs)
bid_px_d: CudaSlice<f32>,
bid_sz_d: CudaSlice<f32>,
ask_px_d: CudaSlice<f32>,
ask_sz_d: CudaSlice<f32>,
prev_bid_sz_d: CudaSlice<f32>,
prev_ask_sz_d: CudaSlice<f32>,
regime_d: CudaSlice<f32>,
snap_feat_d: CudaSlice<f32>,
probs_d: CudaSlice<f32>,
proj_out_d: CudaSlice<f32>,
// Pre-allocated mapped-pinned staging slots for raw input. Reused
// across upload_into calls so no host-malloc happens during or
// around CUDA Graph capture (allocations on the capturing thread
// cause CUDA_ERROR_STREAM_CAPTURE_ISOLATION).
stg_bid_px: MappedF32Buffer,
stg_bid_sz: MappedF32Buffer,
stg_ask_px: MappedF32Buffer,
stg_ask_sz: MappedF32Buffer,
stg_regime: MappedF32Buffer,
// CUDA Graph A — captured perception forward. Scalars (dt_s, ts_ns,
// prev_mid, ...) are frozen at capture time per cudarc 0.19 semantics
// (`launch_builder.arg(&scalar)` records the value, not the pointer).
// The trunk re-captures via `capture_graph_a` whenever those scalars
// would change meaningfully. A follow-up task moves scalars into a
// device-resident buffer so the graph stays valid across changes.
graph_a: Option<CudaGraph>,
}
impl CfcTrunk {
@@ -226,16 +169,13 @@ impl CfcTrunk {
let stream: Arc<CudaStream> = dev.cuda_stream().context("trunk stream")?.clone();
let ctx = dev.cuda_context().context("trunk ctx")?;
// Load forward kernels onto the trunk. snap/step/heads cubins are
// each shared between V1 single-step kernels (no longer used) and
// the batched variants we actually launch; we keep the modules
// alive and only resolve the batched function handles.
let snap_module = ctx.load_cubin(SNAP_CUBIN.to_vec()).context("load snap cubin")?;
let step_module = ctx.load_cubin(STEP_CUBIN.to_vec()).context("load step cubin")?;
let heads_module = ctx.load_cubin(HEADS_CUBIN.to_vec()).context("load heads cubin")?;
let proj_module = ctx.load_cubin(PROJ_CUBIN.to_vec()).context("load proj cubin")?;
let snap_fn = snap_module.load_function("snap_feature_assemble").context("snap_fn")?;
let step_fn = step_module.load_function("cfc_step").context("step_fn")?;
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 forward kernels onto the trunk.
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")?;
@@ -247,113 +187,70 @@ impl CfcTrunk {
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.
// CfC weight init via a seeded ChaCha8Rng — the same chain feeds
// every weight tensor so cfg.seed fully determines the initial
// state. CfC's input dim is `cfg.cfc_n_in` (HIDDEN_DIM in the
// current architecture), separate from the snap-feature `cfg.n_in`.
let mut r = ChaCha8Rng::seed_from_u64(seed);
// X8: CfC layer weights use cfc_n_in (HIDDEN_DIM). The separate
// `n_in` config field still means snap feature dimensionality
// for any residual snap-feature consumers.
let scale_in = (1.0_f32 / cfg.cfc_n_in as f32).sqrt();
let scale_rec = (1.0_f32 / cfg.n_hid as f32).sqrt();
let w_in: Vec<f32> = (0..cfg.n_hid * cfg.cfc_n_in).map(|_| r.gen_range(-scale_in..scale_in)).collect();
let w_rec: Vec<f32> = (0..cfg.n_hid * cfg.n_hid).map(|_| r.gen_range(-scale_rec..scale_rec)).collect();
let w_in: Vec<f32> = (0..cfg.n_hid * cfg.cfc_n_in)
.map(|_| r.gen_range(-scale_in..scale_in)).collect();
let w_rec: Vec<f32> = (0..cfg.n_hid * cfg.n_hid)
.map(|_| r.gen_range(-scale_rec..scale_rec)).collect();
let b: Vec<f32> = vec![0.0; cfg.n_hid];
// tau log-uniform over [10ms, 1000s] per spec Section 2 (ln(0.01)=-4.6, ln(1000)=6.9).
// tau log-uniform over [10ms, 1000s] per spec §2 (ln(0.01)=-4.6, ln(1000)=6.9).
let tau: Vec<f32> = (0..cfg.n_hid)
.map(|_| {
let u: f32 = r.gen_range(0.0..1.0);
(-4.6 + u * (4.6 + 6.9)).exp()
})
.collect();
let head_scale = scale_rec;
let heads_w: Vec<f32> = (0..N_HORIZONS * cfg.n_hid).map(|_| r.gen_range(-head_scale..head_scale)).collect();
let heads_b: Vec<f32> = vec![0.0; N_HORIZONS];
let proj_w: Vec<f32> = (0..PROJ_DIM * cfg.n_hid).map(|_| r.gen_range(-head_scale..head_scale)).collect();
let proj_b: Vec<f32> = vec![0.0; PROJ_DIM];
let proj_g: Vec<f32> = vec![1.0; PROJ_DIM];
let proj_n: Vec<f32> = vec![0.0; PROJ_DIM];
// PRNG consumed by CfC weights only; PerceptionTrainer's ChaCha
// chain (also seeded from cfg.seed) drives the higher-up layers
// (VSN, LN, attn, GRN heads) — see trainer/perception.rs.
let stg_bid_px = unsafe { MappedF32Buffer::new(10) }
.map_err(|e| anyhow::anyhow!("stg bid_px alloc: {e}"))?;
let stg_bid_sz = unsafe { MappedF32Buffer::new(10) }
.map_err(|e| anyhow::anyhow!("stg bid_sz alloc: {e}"))?;
let stg_ask_px = unsafe { MappedF32Buffer::new(10) }
.map_err(|e| anyhow::anyhow!("stg ask_px alloc: {e}"))?;
let stg_ask_sz = unsafe { MappedF32Buffer::new(10) }
.map_err(|e| anyhow::anyhow!("stg ask_sz alloc: {e}"))?;
let stg_regime = unsafe { MappedF32Buffer::new(REGIME_DIM) }
.map_err(|e| anyhow::anyhow!("stg regime alloc: {e}"))?;
let upload = |host: &[f32]| -> Result<CudaSlice<f32>> {
let mut buf = stream.alloc_zeros::<f32>(host.len()).context("upload alloc")?;
stream.memcpy_htod(host, &mut buf).context("upload htod")?;
Ok(buf)
};
Ok(Self {
cfg: cfg.clone(),
w_in_d: upload(&stream, &w_in)?,
w_rec_d: upload(&stream, &w_rec)?,
b_d: upload(&stream, &b)?,
tau_d: upload(&stream, &tau)?,
heads_w_d: upload(&stream, &heads_w)?,
heads_b_d: upload(&stream, &heads_b)?,
proj_w_d: upload(&stream, &proj_w)?,
proj_b_d: upload(&stream, &proj_b)?,
proj_g_d: upload(&stream, &proj_g)?,
proj_n_d: upload(&stream, &proj_n)?,
h_ping: stream.alloc_zeros::<f32>(cfg.n_hid).context("h_ping alloc")?,
h_pong: stream.alloc_zeros::<f32>(cfg.n_hid).context("h_pong alloc")?,
bid_px_d: stream.alloc_zeros::<f32>(10).context("bid_px alloc")?,
bid_sz_d: stream.alloc_zeros::<f32>(10).context("bid_sz alloc")?,
ask_px_d: stream.alloc_zeros::<f32>(10).context("ask_px alloc")?,
ask_sz_d: stream.alloc_zeros::<f32>(10).context("ask_sz alloc")?,
prev_bid_sz_d: stream.alloc_zeros::<f32>(10).context("prev_bid_sz alloc")?,
prev_ask_sz_d: stream.alloc_zeros::<f32>(10).context("prev_ask_sz alloc")?,
regime_d: stream.alloc_zeros::<f32>(REGIME_DIM).context("regime alloc")?,
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")?,
// === Inference 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("vsn_w alloc")?,
vsn_b_d: stream.alloc_zeros::<f32>(cfg.n_in)
.context("vsn_b alloc")?,
w_in_d: upload(&w_in)?,
w_rec_d: upload(&w_rec)?,
b_d: upload(&b)?,
tau_d: upload(&tau)?,
// Inference layers — zero-initialised here; PerceptionTrainer's
// construction overwrites them with its own seeded init via
// memcpy_htod. Standalone trunk callers (tests, fxt-backtest's
// random-init fallback) get zero-initialised inference layers,
// which is fine — the layers are still well-formed device
// tensors of the correct shape, just not trained.
vsn_w_d: stream.alloc_zeros::<f32>(cfg.n_in * cfg.n_in).context("vsn_w alloc")?,
vsn_b_d: stream.alloc_zeros::<f32>(cfg.n_in).context("vsn_b alloc")?,
mamba2_stack_1: None,
mamba2_stack_2: None,
ln_a_gain_d: stream.alloc_zeros::<f32>(cfg.n_hid)
.context("ln_a_gain alloc")?,
ln_a_bias_d: stream.alloc_zeros::<f32>(cfg.n_hid)
.context("ln_a_bias alloc")?,
ln_b_gain_d: stream.alloc_zeros::<f32>(cfg.n_hid)
.context("ln_b_gain alloc")?,
ln_b_bias_d: stream.alloc_zeros::<f32>(cfg.n_hid)
.context("ln_b_bias alloc")?,
attn_q_d: stream.alloc_zeros::<f32>(cfg.n_hid)
.context("attn_q alloc")?,
heads_w1_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM * cfg.n_hid)
.context("heads_w1 alloc")?,
heads_b1_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM)
.context("heads_b1 alloc")?,
heads_w2_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM * HEAD_MID_DIM)
.context("heads_w2 alloc")?,
heads_b2_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM)
.context("heads_b2 alloc")?,
heads_w_gate_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM)
.context("heads_w_gate alloc")?,
heads_b_gate_d: stream.alloc_zeros::<f32>(N_HORIZONS)
.context("heads_b_gate alloc")?,
heads_w_main_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM)
.context("heads_w_main alloc")?,
heads_b_main_d: stream.alloc_zeros::<f32>(N_HORIZONS)
.context("heads_b_main alloc")?,
heads_w_skip_d: stream.alloc_zeros::<f32>(N_HORIZONS * cfg.n_hid)
.context("heads_w_skip alloc")?,
heads_b_skip_d: stream.alloc_zeros::<f32>(N_HORIZONS)
.context("heads_b_skip alloc")?,
ln_a_gain_d: stream.alloc_zeros::<f32>(cfg.n_hid).context("ln_a_gain alloc")?,
ln_a_bias_d: stream.alloc_zeros::<f32>(cfg.n_hid).context("ln_a_bias alloc")?,
ln_b_gain_d: stream.alloc_zeros::<f32>(cfg.n_hid).context("ln_b_gain alloc")?,
ln_b_bias_d: stream.alloc_zeros::<f32>(cfg.n_hid).context("ln_b_bias alloc")?,
attn_q_d: stream.alloc_zeros::<f32>(cfg.n_hid).context("attn_q alloc")?,
heads_w1_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM * cfg.n_hid).context("heads_w1 alloc")?,
heads_b1_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM).context("heads_b1 alloc")?,
heads_w2_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM * HEAD_MID_DIM).context("heads_w2 alloc")?,
heads_b2_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM).context("heads_b2 alloc")?,
heads_w_gate_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM).context("heads_w_gate alloc")?,
heads_b_gate_d: stream.alloc_zeros::<f32>(N_HORIZONS).context("heads_b_gate alloc")?,
heads_w_main_d: stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM).context("heads_w_main alloc")?,
heads_b_main_d: stream.alloc_zeros::<f32>(N_HORIZONS).context("heads_b_main alloc")?,
heads_w_skip_d: stream.alloc_zeros::<f32>(N_HORIZONS * cfg.n_hid).context("heads_w_skip alloc")?,
heads_b_skip_d: stream.alloc_zeros::<f32>(N_HORIZONS).context("heads_b_skip alloc")?,
stream,
_snap_module: snap_module,
_step_module: step_module,
_heads_module: heads_module,
_proj_module: proj_module,
snap_fn,
step_fn,
heads_fn,
proj_fn,
_ln_module: ln_module,
_vsn_module: vsn_module,
_attn_module: attn_module,
@@ -364,265 +261,9 @@ impl CfcTrunk {
step_batched_fn,
heads_grn_fwd_fn,
transpose_3d_fn,
stg_bid_px,
stg_bid_sz,
stg_ask_px,
stg_ask_sz,
stg_regime,
graph_a: None,
})
}
fn upload_pre_allocated(
stream: &Arc<CudaStream>,
host: &[f32],
staging: &MappedF32Buffer,
dst: &mut CudaSlice<f32>,
) -> Result<()> {
let n = host.len();
assert_eq!(n, dst.len());
assert_eq!(n, staging.len);
staging.write_from_slice(host);
if n > 0 {
let nbytes = n * std::mem::size_of::<f32>();
unsafe {
let (dst_ptr, _g) = dst.device_ptr_mut(stream);
cudarc::driver::result::memcpy_dtod_async(
dst_ptr,
staging.dev_ptr,
nbytes,
stream.cu_stream(),
)
.context("upload_pre_allocated DtoD")?;
}
}
Ok(())
}
/// Capture the four perception kernels into CUDA Graph A. The
/// scalars (dt_s, ts_ns, prev_mid, trade flow, trade count) at the
/// time of capture are baked in; production usage must re-capture
/// when those would change. For Phase A training the scalars come
/// from the trainer's per-batch context and capture happens fresh
/// each call.
pub fn capture_graph_a(&mut self, input: &Mbp10RawInput) -> Result<()> {
// Stage raw input into device buffers BEFORE capture using
// pre-allocated mapped-pinned staging (host-malloc during
// capture trips CUDA_ERROR_STREAM_CAPTURE_ISOLATION).
Self::upload_pre_allocated(&self.stream, &input.bid_px, &self.stg_bid_px, &mut self.bid_px_d)?;
Self::upload_pre_allocated(&self.stream, &input.bid_sz, &self.stg_bid_sz, &mut self.bid_sz_d)?;
Self::upload_pre_allocated(&self.stream, &input.ask_px, &self.stg_ask_px, &mut self.ask_px_d)?;
Self::upload_pre_allocated(&self.stream, &input.ask_sz, &self.stg_ask_sz, &mut self.ask_sz_d)?;
Self::upload_pre_allocated(&self.stream, &input.regime, &self.stg_regime, &mut self.regime_d)?;
self.stream.synchronize().context("graph A: pre-capture sync")?;
// cudarc's default event tracking creates cross-stream deps that
// the capture engine rejects with STREAM_CAPTURE_ISOLATION.
// Mirror the working pattern in crates/ml/.../fused_training.rs.
let ctx = self.stream.context().clone();
let _ = ctx.check_err();
unsafe { ctx.disable_event_tracking(); }
let begin_result = self
.stream
.begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED);
if let Err(e) = begin_result {
unsafe { ctx.enable_event_tracking(); }
let _ = ctx.check_err();
return Err(anyhow::anyhow!("graph A begin_capture: {e}"));
}
let dispatch_result = self.dispatch_perception(input);
let graph_result = self
.stream
.end_capture(CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
unsafe { ctx.enable_event_tracking(); }
let _ = ctx.check_err();
dispatch_result?;
let graph = graph_result
.context("graph A end_capture")?
.ok_or_else(|| anyhow::anyhow!("end_capture returned None — no work was captured"))?;
self.graph_a = Some(graph);
Ok(())
}
/// Replay the captured graph. Returns probs[5] + projection[8].
/// Caller is responsible for updating the trunk's pre-allocated
/// input buffers (via `update_input_buffers`) before calling this.
pub fn perception_forward_captured(&mut self) -> Result<([f32; N_HORIZONS], [f32; PROJ_DIM])> {
let graph = self
.graph_a
.as_ref()
.context("Graph A not captured — call capture_graph_a first")?;
graph.launch().context("graph A launch")?;
self.stream.synchronize().context("graph A sync")?;
let probs_vec = download(&self.stream, &self.probs_d)?;
let proj_vec = download(&self.stream, &self.proj_out_d)?;
let mut probs = [0f32; N_HORIZONS];
probs.copy_from_slice(&probs_vec);
let mut proj = [0f32; PROJ_DIM];
proj.copy_from_slice(&proj_vec);
Ok((probs, proj))
}
fn dispatch_perception(&mut self, input: &Mbp10RawInput) -> Result<()> {
let prev_mid = input.prev_mid;
let trade_signed_vol = input.trade_signed_vol;
let trade_count = input.trade_count as i32;
let ts_ns = input.ts_ns as i64;
let prev_ts_ns = input.prev_ts_ns as i64;
let tick_size = ES_TICK_SIZE;
let cfg1 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 };
{
let mut launch = self.stream.launch_builder(&self.snap_fn);
launch
.arg(&self.bid_px_d).arg(&self.bid_sz_d).arg(&self.ask_px_d).arg(&self.ask_sz_d)
.arg(&self.prev_bid_sz_d).arg(&self.prev_ask_sz_d)
.arg(&self.regime_d)
.arg(&prev_mid).arg(&trade_signed_vol).arg(&trade_count)
.arg(&ts_ns).arg(&prev_ts_ns).arg(&tick_size)
.arg(&mut self.snap_feat_d);
unsafe { launch.launch(cfg1).context("snap launch")?; }
}
let dt_s = input.ts_ns.saturating_sub(input.prev_ts_ns) as f32 * 1e-9;
let n_in_i = self.cfg.n_in as i32;
let n_hid_i = self.cfg.n_hid as i32;
let block_dim = 128.min(self.cfg.n_hid as u32);
let grid_dim = ((self.cfg.n_hid as u32) + block_dim - 1) / block_dim;
let cfg2 = LaunchConfig { grid_dim: (grid_dim, 1, 1), block_dim: (block_dim, 1, 1), shared_mem_bytes: 0 };
{
let mut launch = self.stream.launch_builder(&self.step_fn);
launch
.arg(&self.w_in_d).arg(&self.w_rec_d).arg(&self.b_d).arg(&self.tau_d)
.arg(&self.snap_feat_d).arg(&self.h_ping)
.arg(&dt_s).arg(&n_in_i).arg(&n_hid_i)
.arg(&mut self.h_pong);
unsafe { launch.launch(cfg2).context("step launch")?; }
}
// NOTE: ping-pong swap intentionally OMITTED inside the captured
// path. Inside Graph A we always write to h_pong from h_ping; the
// swap happens via a DtoD copy outside the captured region after
// each replay. Doing the swap as `std::mem::swap` would change
// pointer identity, which breaks captured kernel args.
let cfg3 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (N_HORIZONS as u32, 1, 1), shared_mem_bytes: 0 };
{
let mut launch = self.stream.launch_builder(&self.heads_fn);
launch.arg(&self.heads_w_d).arg(&self.heads_b_d).arg(&self.h_pong).arg(&mut self.probs_d);
unsafe { launch.launch(cfg3).context("heads launch")?; }
}
let cfg4 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (PROJ_DIM as u32, 1, 1), shared_mem_bytes: 0 };
{
let mut launch = self.stream.launch_builder(&self.proj_fn);
launch
.arg(&self.proj_w_d).arg(&self.proj_b_d)
.arg(&self.proj_g_d).arg(&self.proj_n_d)
.arg(&self.h_pong).arg(&mut self.proj_out_d);
unsafe { launch.launch(cfg4).context("proj launch")?; }
}
Ok(())
}
/// Mirror snapshot input into the pre-allocated device buffers
/// without dispatching the kernels (used between Graph A replays).
pub fn update_input_buffers(&mut self, input: &Mbp10RawInput) -> Result<()> {
Self::upload_pre_allocated(&self.stream, &input.bid_px, &self.stg_bid_px, &mut self.bid_px_d)?;
Self::upload_pre_allocated(&self.stream, &input.bid_sz, &self.stg_bid_sz, &mut self.bid_sz_d)?;
Self::upload_pre_allocated(&self.stream, &input.ask_px, &self.stg_ask_px, &mut self.ask_px_d)?;
Self::upload_pre_allocated(&self.stream, &input.ask_sz, &self.stg_ask_sz, &mut self.ask_sz_d)?;
Ok(())
}
/// Sequential dispatch: snap_features -> cfc_step -> heads -> projection.
/// Returns (probs[5], projection[8]). Updates h_ping/h_pong in place.
pub fn forward_snapshot(
&mut self,
input: &Mbp10RawInput,
) -> Result<([f32; N_HORIZONS], [f32; PROJ_DIM])> {
// Stage raw input into pre-allocated device buffers via mapped-pinned.
upload_into(&self.stream, &input.bid_px, &mut self.bid_px_d)?;
upload_into(&self.stream, &input.bid_sz, &mut self.bid_sz_d)?;
upload_into(&self.stream, &input.ask_px, &mut self.ask_px_d)?;
upload_into(&self.stream, &input.ask_sz, &mut self.ask_sz_d)?;
upload_into(&self.stream, &input.regime, &mut self.regime_d)?;
// Step 1: snap_feature_assemble.
let prev_mid = input.prev_mid;
let trade_signed_vol = input.trade_signed_vol;
let trade_count = input.trade_count as i32;
let ts_ns = input.ts_ns as i64;
let prev_ts_ns = input.prev_ts_ns as i64;
let tick_size = ES_TICK_SIZE;
let cfg1 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 };
{
let mut launch = self.stream.launch_builder(&self.snap_fn);
launch
.arg(&self.bid_px_d).arg(&self.bid_sz_d).arg(&self.ask_px_d).arg(&self.ask_sz_d)
.arg(&self.prev_bid_sz_d).arg(&self.prev_ask_sz_d)
.arg(&self.regime_d)
.arg(&prev_mid).arg(&trade_signed_vol).arg(&trade_count)
.arg(&ts_ns).arg(&prev_ts_ns).arg(&tick_size)
.arg(&mut self.snap_feat_d);
unsafe { launch.launch(cfg1).context("snap launch")?; }
}
// Step 2: cfc_step (h_ping -> h_pong).
let dt_s = input.ts_ns.saturating_sub(input.prev_ts_ns) as f32 * 1e-9;
let n_in_i = self.cfg.n_in as i32;
let n_hid_i = self.cfg.n_hid as i32;
let block_dim = 128.min(self.cfg.n_hid as u32);
let grid_dim = ((self.cfg.n_hid as u32) + block_dim - 1) / block_dim;
let cfg2 = LaunchConfig { grid_dim: (grid_dim, 1, 1), block_dim: (block_dim, 1, 1), shared_mem_bytes: 0 };
{
let mut launch = self.stream.launch_builder(&self.step_fn);
launch
.arg(&self.w_in_d).arg(&self.w_rec_d).arg(&self.b_d).arg(&self.tau_d)
.arg(&self.snap_feat_d).arg(&self.h_ping)
.arg(&dt_s).arg(&n_in_i).arg(&n_hid_i)
.arg(&mut self.h_pong);
unsafe { launch.launch(cfg2).context("step launch")?; }
}
std::mem::swap(&mut self.h_ping, &mut self.h_pong);
// Step 3: heads.
let cfg3 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (N_HORIZONS as u32, 1, 1), shared_mem_bytes: 0 };
{
let mut launch = self.stream.launch_builder(&self.heads_fn);
launch.arg(&self.heads_w_d).arg(&self.heads_b_d).arg(&self.h_ping).arg(&mut self.probs_d);
unsafe { launch.launch(cfg3).context("heads launch")?; }
}
// Step 4: projection.
let cfg4 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (PROJ_DIM as u32, 1, 1), shared_mem_bytes: 0 };
{
let mut launch = self.stream.launch_builder(&self.proj_fn);
launch
.arg(&self.proj_w_d).arg(&self.proj_b_d)
.arg(&self.proj_g_d).arg(&self.proj_n_d)
.arg(&self.h_ping).arg(&mut self.proj_out_d);
unsafe { launch.launch(cfg4).context("proj launch")?; }
}
// Copy current bid_sz/ask_sz into prev_*_sz for next-step OFI baseline (DtoD).
copy_dtod(&self.stream, &self.bid_sz_d, &mut self.prev_bid_sz_d)?;
copy_dtod(&self.stream, &self.ask_sz_d, &mut self.prev_ask_sz_d)?;
self.stream.synchronize().context("trunk forward sync")?;
// Readback (slow-path for now; Task 11 captures into Graph A and
// keeps everything on-device).
let probs_vec = download(&self.stream, &self.probs_d)?;
let proj_vec = download(&self.stream, &self.proj_out_d)?;
let mut probs = [0f32; N_HORIZONS];
probs.copy_from_slice(&probs_vec);
let mut proj = [0f32; PROJ_DIM];
proj.copy_from_slice(&proj_vec);
Ok((probs, proj))
}
/// Serialise the trunk weights + config to a bincode file. Used by
/// downstream backtest tooling (fxt-backtest --checkpoint) to run
@@ -798,10 +439,6 @@ impl CfcTrunk {
Ok(trunk)
}
pub fn snapshot_hidden(&self) -> Result<Vec<f32>> {
download(&self.stream, &self.h_ping)
}
pub fn config(&self) -> &CfcConfig { &self.cfg }
}
@@ -812,7 +449,7 @@ mod checkpoint_tests {
#[test]
#[ignore = "requires CUDA"]
fn save_load_roundtrip_preserves_weights() {
fn save_load_roundtrip_preserves_cfc_weights() {
let dev = match MlDevice::cuda(0) {
Ok(d) => d,
Err(e) => {
@@ -826,89 +463,20 @@ mod checkpoint_tests {
let path = dir.path().join("trunk.bin");
trunk_a.save_checkpoint(&path).expect("save_checkpoint");
let trunk_b = CfcTrunk::load_checkpoint(&dev, &cfg, &path).expect("load_checkpoint");
// Compare each weight tensor bit-exact between the two trunks.
let read = |slice: &CudaSlice<f32>, n: usize| -> Vec<f32> {
let mut v = vec![0.0; n];
trunk_a.stream.memcpy_dtoh(slice, v.as_mut_slice()).unwrap();
let read = |trunk: &CfcTrunk, slice: &CudaSlice<f32>| -> Vec<f32> {
let mut v = vec![0.0; slice.len()];
trunk.stream.memcpy_dtoh(slice, v.as_mut_slice()).unwrap();
v
};
let read_b = |slice: &CudaSlice<f32>, n: usize| -> Vec<f32> {
let mut v = vec![0.0; n];
trunk_b.stream.memcpy_dtoh(slice, v.as_mut_slice()).unwrap();
v
};
assert_eq!(read(&trunk_a.w_in_d, cfg.n_hid * cfg.n_in), read_b(&trunk_b.w_in_d, cfg.n_hid * cfg.n_in));
assert_eq!(read(&trunk_a.w_rec_d, cfg.n_hid * cfg.n_hid), read_b(&trunk_b.w_rec_d, cfg.n_hid * cfg.n_hid));
assert_eq!(read(&trunk_a.b_d, cfg.n_hid), read_b(&trunk_b.b_d, cfg.n_hid));
assert_eq!(read(&trunk_a.tau_d, cfg.n_hid), read_b(&trunk_b.tau_d, cfg.n_hid));
assert_eq!(read(&trunk_a.heads_w_d, N_HORIZONS * cfg.n_hid), read_b(&trunk_b.heads_w_d, N_HORIZONS * cfg.n_hid));
assert_eq!(read(&trunk_a.heads_b_d, N_HORIZONS), read_b(&trunk_b.heads_b_d, N_HORIZONS));
assert_eq!(read(&trunk_a.proj_w_d, PROJ_DIM * cfg.n_hid), read_b(&trunk_b.proj_w_d, PROJ_DIM * cfg.n_hid));
assert_eq!(read(&trunk_a.proj_b_d, PROJ_DIM), read_b(&trunk_b.proj_b_d, PROJ_DIM));
assert_eq!(read(&trunk_a.proj_g_d, PROJ_DIM), read_b(&trunk_b.proj_g_d, PROJ_DIM));
assert_eq!(read(&trunk_a.proj_n_d, PROJ_DIM), read_b(&trunk_b.proj_n_d, PROJ_DIM));
assert_eq!(read(&trunk_a, &trunk_a.w_in_d), read(&trunk_b, &trunk_b.w_in_d));
assert_eq!(read(&trunk_a, &trunk_a.w_rec_d), read(&trunk_b, &trunk_b.w_rec_d));
assert_eq!(read(&trunk_a, &trunk_a.b_d), read(&trunk_b, &trunk_b.b_d));
assert_eq!(read(&trunk_a, &trunk_a.tau_d), read(&trunk_b, &trunk_b.tau_d));
assert_eq!(read(&trunk_a, &trunk_a.vsn_w_d), read(&trunk_b, &trunk_b.vsn_w_d));
assert_eq!(read(&trunk_a, &trunk_a.ln_a_gain_d), read(&trunk_b, &trunk_b.ln_a_gain_d));
assert_eq!(read(&trunk_a, &trunk_a.attn_q_d), read(&trunk_b, &trunk_b.attn_q_d));
assert_eq!(read(&trunk_a, &trunk_a.heads_w_gate_d), read(&trunk_b, &trunk_b.heads_w_gate_d));
}
}
// Helpers ---------------------------------------------------------------
fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
let n = host.len();
let staging = unsafe { MappedF32Buffer::new(n) }
.map_err(|e| anyhow::anyhow!("trunk upload staging: {e}"))?;
staging.write_from_slice(host);
let mut dst = stream.alloc_zeros::<f32>(n).context("trunk upload alloc")?;
if n > 0 {
let nbytes = n * std::mem::size_of::<f32>();
unsafe {
let (dst_ptr, _g) = dst.device_ptr_mut(stream);
cudarc::driver::result::memcpy_dtod_async(dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream())
.context("trunk upload DtoD")?;
}
}
Ok(dst)
}
fn upload_into(stream: &Arc<CudaStream>, host: &[f32], dst: &mut CudaSlice<f32>) -> Result<()> {
let n = host.len();
assert_eq!(n, dst.len());
let staging = unsafe { MappedF32Buffer::new(n) }
.map_err(|e| anyhow::anyhow!("upload_into staging: {e}"))?;
staging.write_from_slice(host);
let nbytes = n * std::mem::size_of::<f32>();
unsafe {
let (dst_ptr, _g) = dst.device_ptr_mut(stream);
cudarc::driver::result::memcpy_dtod_async(dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream())
.context("upload_into DtoD")?;
}
Ok(())
}
fn copy_dtod(stream: &Arc<CudaStream>, src: &CudaSlice<f32>, dst: &mut CudaSlice<f32>) -> Result<()> {
let n = src.len();
assert_eq!(n, dst.len());
let nbytes = n * std::mem::size_of::<f32>();
unsafe {
let (src_ptr, _g1) = src.device_ptr(stream);
let (dst_ptr, _g2) = dst.device_ptr_mut(stream);
cudarc::driver::result::memcpy_dtod_async(dst_ptr, src_ptr, nbytes, stream.cu_stream())
.context("copy_dtod")?;
}
Ok(())
}
fn download(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>> {
let n = src.len();
let staging = unsafe { MappedF32Buffer::new(n) }
.map_err(|e| anyhow::anyhow!("trunk download staging: {e}"))?;
let nbytes = n * std::mem::size_of::<f32>();
unsafe {
let (src_ptr, _g) = src.device_ptr(stream);
cudarc::driver::result::memcpy_dtod_async(staging.dev_ptr, src_ptr, nbytes, stream.cu_stream())
.context("trunk download DtoD")?;
}
stream.synchronize().context("trunk download sync")?;
Ok(staging.read_all())
}

38
crates/ml-alpha/src/trainer/perception.rs
View File

@@ -139,19 +139,12 @@ pub struct PerceptionTrainer {
_step_module: Arc<CudaModule>,
_heads_module: Arc<CudaModule>,
_bce_module: Arc<CudaModule>,
/// Batched snap_feature kernel — processes all B*K snapshots per
/// training step in one launch.
snap_batched_fn: CudaFunction,
// Forward kernels are owned by `self.trunk` (snap_batched_fn,
// step_batched_fn, heads_grn_fwd_fn, transpose_3d_fn). The trainer
// keeps only the training-only handles below.
bce_fn: CudaFunction,
/// Batched cfc/heads forward + backward kernels — used by
/// `step_batched()` for all training updates regardless of B.
/// At cfg.n_batch == 1 they produce identical results to the
/// unbatched kernels (the inner per-thread loop runs once).
step_batched_fn: CudaFunction,
step_bwd_batched_fn: CudaFunction,
heads_grn_fwd_fn: CudaFunction,
heads_grn_bwd_fn: CudaFunction,
transpose_3d_fn: CudaFunction,
// X3: Mamba2 stack 1 weights moved to `self.trunk.mamba2_stack_1`.
// Access via `self.trunk.mamba2_l1()` / `mamba2_l1_mut()`. Gradient
@@ -314,7 +307,6 @@ pub struct PerceptionTrainer {
/// feeding Mamba2 backward.
grad_ln_in_d: GpuTensor,
/// Cached LN kernel handles + module lifetime owner.
ln_fwd_fn: CudaFunction,
ln_bwd_fn: CudaFunction,
ln_reduce_fn: CudaFunction,
_ln_module: Arc<CudaModule>,
@@ -347,7 +339,6 @@ pub struct PerceptionTrainer {
grad_vsn_b_d: CudaSlice<f32>,
pub opt_vsn_w: AdamW,
pub opt_vsn_b: AdamW,
vsn_fwd_fn: CudaFunction,
vsn_bwd_fn: CudaFunction,
_vsn_module: Arc<CudaModule>,
@@ -362,7 +353,6 @@ pub struct PerceptionTrainer {
attn_weights_d: CudaSlice<f32>,
grad_attn_q_d: CudaSlice<f32>,
pub opt_attn_q: AdamW,
attn_fwd_fn: CudaFunction,
attn_bwd_fn: CudaFunction,
_attn_module: Arc<CudaModule>,
@@ -551,9 +541,6 @@ impl PerceptionTrainer {
let ln_module = ctx
.load_cubin(LAYER_NORM_CUBIN.to_vec())
.context("layer_norm cubin")?;
let ln_fwd_fn = ln_module
.load_function("layer_norm_fwd")
.context("layer_norm_fwd symbol")?;
let ln_bwd_fn = ln_module
.load_function("layer_norm_bwd")
.context("layer_norm_bwd symbol")?;
@@ -563,18 +550,12 @@ impl PerceptionTrainer {
let vsn_module = ctx
.load_cubin(VARIABLE_SELECTION_CUBIN.to_vec())
.context("variable_selection cubin")?;
let vsn_fwd_fn = vsn_module
.load_function("variable_selection_fwd")
.context("variable_selection_fwd symbol")?;
let vsn_bwd_fn = vsn_module
.load_function("variable_selection_bwd")
.context("variable_selection_bwd symbol")?;
let attn_module = ctx
.load_cubin(ATTENTION_POOL_CUBIN.to_vec())
.context("attention_pool cubin")?;
let attn_fwd_fn = attn_module
.load_function("attention_pool_fwd")
.context("attention_pool_fwd symbol")?;
let attn_bwd_fn = attn_module
.load_function("attention_pool_bwd")
.context("attention_pool_bwd symbol")?;
@@ -584,17 +565,11 @@ impl PerceptionTrainer {
let reduce_axis0_fn = reduce_module
.load_function("reduce_axis0")
.context("reduce_axis0 symbol")?;
let snap_batched_fn = snap_module.load_function("snap_feature_assemble_batched")?;
let bce_fn = bce_module.load_function("bce_multi_horizon_forward_backward")?;
let step_batched_fn = step_module.load_function("cfc_step_batched")?;
let step_bwd_batched_fn = step_module.load_function("cfc_step_backward_batched")?;
let heads_grn_fwd_fn = heads_module
.load_function("multi_horizon_heads_grn_fwd_batched")
.context("heads GRN fwd symbol")?;
let heads_grn_bwd_fn = heads_module
.load_function("multi_horizon_heads_grn_bwd_batched")
.context("heads GRN bwd symbol")?;
let transpose_3d_fn = step_module.load_function("transpose_3d_swap_01")?;
// Mamba2 encoder: in_dim=FEATURE_DIM (32), hidden=HIDDEN_DIM (128).
let mamba2 = Mamba2Block::new(
@@ -922,7 +897,6 @@ impl PerceptionTrainer {
grad_ln_bias_d: stream.alloc_zeros::<f32>(HIDDEN_DIM)?,
grad_ln_in_d: GpuTensor::zeros(&[cfg.n_batch, k, HIDDEN_DIM], &stream)
.map_err(|e| anyhow::anyhow!("grad_ln_in_d alloc: {e}"))?,
ln_fwd_fn,
ln_bwd_fn,
ln_reduce_fn,
_ln_module: ln_module,
@@ -938,7 +912,6 @@ impl PerceptionTrainer {
grad_vsn_b_d: stream.alloc_zeros::<f32>(FEATURE_DIM)?,
opt_vsn_w,
opt_vsn_b,
vsn_fwd_fn,
vsn_bwd_fn,
_vsn_module: vsn_module,
// Attention pool (Phase 3).
@@ -946,7 +919,6 @@ impl PerceptionTrainer {
attn_weights_d,
grad_attn_q_d,
opt_attn_q,
attn_fwd_fn,
attn_bwd_fn,
_attn_module: attn_module,
log_sigma_h_d,
@@ -1046,13 +1018,9 @@ impl PerceptionTrainer {
_step_module: step_module,
_heads_module: heads_module,
_bce_module: bce_module,
snap_batched_fn,
bce_fn,
step_batched_fn,
step_bwd_batched_fn,
heads_grn_fwd_fn,
heads_grn_bwd_fn,
transpose_3d_fn,
mamba2_adamw,
mamba2_fwd_scratch,
mamba2_bwd_scratch,