From cecc08a12283f8fc0bf2ab342b895aebddc39079 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 19 May 2026 09:19:26 +0200 Subject: [PATCH] =?UTF-8?q?chore(ml-alpha):=20deep=20cleanup=20=E2=80=94?= =?UTF-8?q?=20delete=20all=20V1=20dead=20code?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- crates/ml-alpha/src/cfc/trunk.rs | 580 +++------------------- crates/ml-alpha/src/trainer/perception.rs | 38 +- 2 files changed, 77 insertions(+), 541 deletions(-) diff --git a/crates/ml-alpha/src/cfc/trunk.rs b/crates/ml-alpha/src/cfc/trunk.rs index 589027c53..0d5ba002b 100644 --- a/crates/ml-alpha/src/cfc/trunk.rs +++ b/crates/ml-alpha/src/cfc/trunk.rs @@ -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, heads_b_skip: Vec, } -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.` in a subsequent commit (X10b). @@ -117,21 +113,14 @@ pub struct CfcTrunk { cfg: CfcConfig, stream: Arc, - // Cubin modules (kept alive for the function refs below) + // Cubin modules — kept alive for the function refs below. _snap_module: Arc, _step_module: Arc, _heads_module: Arc, - _proj_module: Arc, - 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.) and X11 (full capture_graph_a). _ln_module: Arc, _vsn_module: Arc, _attn_module: Arc, + // 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, pub w_rec_d: CudaSlice, pub b_d: CudaSlice, pub tau_d: CudaSlice, - // 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, - heads_b_d: CudaSlice, - proj_w_d: CudaSlice, - proj_b_d: CudaSlice, - proj_g_d: CudaSlice, - proj_n_d: CudaSlice, - // === 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, pub vsn_b_d: CudaSlice, pub mamba2_stack_1: Option, @@ -172,9 +149,9 @@ pub struct CfcTrunk { pub ln_b_gain_d: CudaSlice, pub ln_b_bias_d: CudaSlice, pub attn_q_d: CudaSlice, - // 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, pub heads_b1_d: CudaSlice, pub heads_w2_d: CudaSlice, @@ -185,40 +162,6 @@ pub struct CfcTrunk { pub heads_b_main_d: CudaSlice, pub heads_w_skip_d: CudaSlice, pub heads_b_skip_d: CudaSlice, - - // Ping-pong hidden state - h_ping: CudaSlice, - h_pong: CudaSlice, - - // Pre-allocated per-step scratch (raw input + features + outputs) - bid_px_d: CudaSlice, - bid_sz_d: CudaSlice, - ask_px_d: CudaSlice, - ask_sz_d: CudaSlice, - prev_bid_sz_d: CudaSlice, - prev_ask_sz_d: CudaSlice, - regime_d: CudaSlice, - snap_feat_d: CudaSlice, - probs_d: CudaSlice, - proj_out_d: CudaSlice, - - // 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, } impl CfcTrunk { @@ -226,16 +169,13 @@ impl CfcTrunk { let stream: Arc = 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 = (0..cfg.n_hid * cfg.cfc_n_in).map(|_| r.gen_range(-scale_in..scale_in)).collect(); - let w_rec: Vec = (0..cfg.n_hid * cfg.n_hid).map(|_| r.gen_range(-scale_rec..scale_rec)).collect(); + let w_in: Vec = (0..cfg.n_hid * cfg.cfc_n_in) + .map(|_| r.gen_range(-scale_in..scale_in)).collect(); + let w_rec: Vec = (0..cfg.n_hid * cfg.n_hid) + .map(|_| r.gen_range(-scale_rec..scale_rec)).collect(); let b: Vec = 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 = (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 = (0..N_HORIZONS * cfg.n_hid).map(|_| r.gen_range(-head_scale..head_scale)).collect(); - let heads_b: Vec = vec![0.0; N_HORIZONS]; - let proj_w: Vec = (0..PROJ_DIM * cfg.n_hid).map(|_| r.gen_range(-head_scale..head_scale)).collect(); - let proj_b: Vec = vec![0.0; PROJ_DIM]; - let proj_g: Vec = vec![1.0; PROJ_DIM]; - let proj_n: Vec = 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> { + let mut buf = stream.alloc_zeros::(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::(cfg.n_hid).context("h_ping alloc")?, - h_pong: stream.alloc_zeros::(cfg.n_hid).context("h_pong alloc")?, - bid_px_d: stream.alloc_zeros::(10).context("bid_px alloc")?, - bid_sz_d: stream.alloc_zeros::(10).context("bid_sz alloc")?, - ask_px_d: stream.alloc_zeros::(10).context("ask_px alloc")?, - ask_sz_d: stream.alloc_zeros::(10).context("ask_sz alloc")?, - prev_bid_sz_d: stream.alloc_zeros::(10).context("prev_bid_sz alloc")?, - prev_ask_sz_d: stream.alloc_zeros::(10).context("prev_ask_sz alloc")?, - regime_d: stream.alloc_zeros::(REGIME_DIM).context("regime alloc")?, - snap_feat_d: stream.alloc_zeros::(FEATURE_DIM).context("snap_feat alloc")?, - probs_d: stream.alloc_zeros::(N_HORIZONS).context("probs alloc")?, - proj_out_d: stream.alloc_zeros::(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::(cfg.n_in * cfg.n_in) - .context("vsn_w alloc")?, - vsn_b_d: stream.alloc_zeros::(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::(cfg.n_in * cfg.n_in).context("vsn_w alloc")?, + vsn_b_d: stream.alloc_zeros::(cfg.n_in).context("vsn_b alloc")?, mamba2_stack_1: None, mamba2_stack_2: None, - ln_a_gain_d: stream.alloc_zeros::(cfg.n_hid) - .context("ln_a_gain alloc")?, - ln_a_bias_d: stream.alloc_zeros::(cfg.n_hid) - .context("ln_a_bias alloc")?, - ln_b_gain_d: stream.alloc_zeros::(cfg.n_hid) - .context("ln_b_gain alloc")?, - ln_b_bias_d: stream.alloc_zeros::(cfg.n_hid) - .context("ln_b_bias alloc")?, - attn_q_d: stream.alloc_zeros::(cfg.n_hid) - .context("attn_q alloc")?, - heads_w1_d: stream.alloc_zeros::(N_HORIZONS * HEAD_MID_DIM * cfg.n_hid) - .context("heads_w1 alloc")?, - heads_b1_d: stream.alloc_zeros::(N_HORIZONS * HEAD_MID_DIM) - .context("heads_b1 alloc")?, - heads_w2_d: stream.alloc_zeros::(N_HORIZONS * HEAD_MID_DIM * HEAD_MID_DIM) - .context("heads_w2 alloc")?, - heads_b2_d: stream.alloc_zeros::(N_HORIZONS * HEAD_MID_DIM) - .context("heads_b2 alloc")?, - heads_w_gate_d: stream.alloc_zeros::(N_HORIZONS * HEAD_MID_DIM) - .context("heads_w_gate alloc")?, - heads_b_gate_d: stream.alloc_zeros::(N_HORIZONS) - .context("heads_b_gate alloc")?, - heads_w_main_d: stream.alloc_zeros::(N_HORIZONS * HEAD_MID_DIM) - .context("heads_w_main alloc")?, - heads_b_main_d: stream.alloc_zeros::(N_HORIZONS) - .context("heads_b_main alloc")?, - heads_w_skip_d: stream.alloc_zeros::(N_HORIZONS * cfg.n_hid) - .context("heads_w_skip alloc")?, - heads_b_skip_d: stream.alloc_zeros::(N_HORIZONS) - .context("heads_b_skip alloc")?, + ln_a_gain_d: stream.alloc_zeros::(cfg.n_hid).context("ln_a_gain alloc")?, + ln_a_bias_d: stream.alloc_zeros::(cfg.n_hid).context("ln_a_bias alloc")?, + ln_b_gain_d: stream.alloc_zeros::(cfg.n_hid).context("ln_b_gain alloc")?, + ln_b_bias_d: stream.alloc_zeros::(cfg.n_hid).context("ln_b_bias alloc")?, + attn_q_d: stream.alloc_zeros::(cfg.n_hid).context("attn_q alloc")?, + heads_w1_d: stream.alloc_zeros::(N_HORIZONS * HEAD_MID_DIM * cfg.n_hid).context("heads_w1 alloc")?, + heads_b1_d: stream.alloc_zeros::(N_HORIZONS * HEAD_MID_DIM).context("heads_b1 alloc")?, + heads_w2_d: stream.alloc_zeros::(N_HORIZONS * HEAD_MID_DIM * HEAD_MID_DIM).context("heads_w2 alloc")?, + heads_b2_d: stream.alloc_zeros::(N_HORIZONS * HEAD_MID_DIM).context("heads_b2 alloc")?, + heads_w_gate_d: stream.alloc_zeros::(N_HORIZONS * HEAD_MID_DIM).context("heads_w_gate alloc")?, + heads_b_gate_d: stream.alloc_zeros::(N_HORIZONS).context("heads_b_gate alloc")?, + heads_w_main_d: stream.alloc_zeros::(N_HORIZONS * HEAD_MID_DIM).context("heads_w_main alloc")?, + heads_b_main_d: stream.alloc_zeros::(N_HORIZONS).context("heads_b_main alloc")?, + heads_w_skip_d: stream.alloc_zeros::(N_HORIZONS * cfg.n_hid).context("heads_w_skip alloc")?, + heads_b_skip_d: stream.alloc_zeros::(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, - host: &[f32], - staging: &MappedF32Buffer, - dst: &mut CudaSlice, - ) -> 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::(); - 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> { - 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, n: usize| -> Vec { - let mut v = vec![0.0; n]; - trunk_a.stream.memcpy_dtoh(slice, v.as_mut_slice()).unwrap(); + let read = |trunk: &CfcTrunk, slice: &CudaSlice| -> Vec { + let mut v = vec![0.0; slice.len()]; + trunk.stream.memcpy_dtoh(slice, v.as_mut_slice()).unwrap(); v }; - let read_b = |slice: &CudaSlice, n: usize| -> Vec { - 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, host: &[f32]) -> Result> { - 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::(n).context("trunk upload alloc")?; - if n > 0 { - let nbytes = n * std::mem::size_of::(); - 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, host: &[f32], dst: &mut CudaSlice) -> 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::(); - 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, src: &CudaSlice, dst: &mut CudaSlice) -> Result<()> { - let n = src.len(); - assert_eq!(n, dst.len()); - let nbytes = n * std::mem::size_of::(); - 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, src: &CudaSlice) -> Result> { - 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::(); - 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()) -} diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 3566adbcf..2a598e4ad 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -139,19 +139,12 @@ pub struct PerceptionTrainer { _step_module: Arc, _heads_module: Arc, _bce_module: Arc, - /// 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, @@ -347,7 +339,6 @@ pub struct PerceptionTrainer { grad_vsn_b_d: CudaSlice, pub opt_vsn_w: AdamW, pub opt_vsn_b: AdamW, - vsn_fwd_fn: CudaFunction, vsn_bwd_fn: CudaFunction, _vsn_module: Arc, @@ -362,7 +353,6 @@ pub struct PerceptionTrainer { attn_weights_d: CudaSlice, grad_attn_q_d: CudaSlice, pub opt_attn_q: AdamW, - attn_fwd_fn: CudaFunction, attn_bwd_fn: CudaFunction, _attn_module: Arc, @@ -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::(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::(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,