diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index e54292436..b046da706 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -60,6 +60,7 @@ const HEADS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_horiz const BCE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bce_loss_multi_horizon.cubin")); const HORIZON_LAMBDA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/horizon_lambda.cubin")); const LAYER_NORM_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/layer_norm.cubin")); +const VARIABLE_SELECTION_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/variable_selection.cubin")); #[derive(Clone, Debug)] pub struct PerceptionTrainerConfig { @@ -306,6 +307,31 @@ pub struct PerceptionTrainer { /// AdamW state for the two LN params (gain + bias). opt_ln_gain: AdamW, opt_ln_bias: AdamW, + + // ── TFT VSN (Phase 2D) ── + // Per-position softmax-normalised feature gating at trunk entry: + // gates = softmax(W_vsn @ x + b_vsn); y = x * gates + // Lets the model learn which snap_features matter per regime. + /// `[FEATURE_DIM, FEATURE_DIM]` linear layer for gate logits. + pub vsn_w_d: CudaSlice, + pub vsn_b_d: CudaSlice, + /// VSN forward output `[B, K, FEATURE_DIM]` — replaces window_tensor_d + /// as Mamba2's input. window_tensor_d still holds the RAW features + /// (consumed by VSN bwd as `x`). + vsn_out_d: GpuTensor, + /// Saved softmax-normalised gates `[B*K, FEATURE_DIM]` (per-K row). + vsn_gates_d: CudaSlice, + /// VSN bwd writes grad_x here (gradient w.r.t. snap_features). Not + /// propagated further (snap_features are a non-trainable transform + /// of raw MBP-10 data); allocated only so the kernel has a buffer. + vsn_grad_x_d: CudaSlice, + grad_vsn_w_d: CudaSlice, + 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, /// Per-horizon EMA of `loss_per_horizon_d`. Zero-initialised /// (sentinel); the horizon_ema_and_lambda kernel detects the /// sentinel on step 1 and replaces directly per @@ -422,6 +448,15 @@ impl PerceptionTrainer { let ln_reduce_fn = ln_module .load_function("layer_norm_reduce_param_grads") .context("layer_norm_reduce_param_grads symbol")?; + 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 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")?; @@ -560,6 +595,20 @@ impl PerceptionTrainer { let mut opt_ln_bias = AdamW::new(dev, HIDDEN_DIM, cfg.lr_cfc)?; opt_ln_bias.wd = 0.0; + // ── VSN init (Phase 2D) ── + // W_vsn ∈ [FEATURE_DIM, FEATURE_DIM] init near zero so initial + // gates ≈ softmax(0) = uniform 1/FEATURE_DIM. The model can then + // learn departures from uniform. Scale = sqrt(1/FEATURE_DIM). + let vsn_scale = (1.0_f32 / FEATURE_DIM as f32).sqrt(); + let vsn_w_init: Vec = (0..FEATURE_DIM * FEATURE_DIM) + .map(|_| r.gen_range(-vsn_scale..vsn_scale)).collect(); + let vsn_b_init: Vec = vec![0.0; FEATURE_DIM]; + let vsn_w_d = upload(&stream, &vsn_w_init)?; + let vsn_b_d = upload(&stream, &vsn_b_init)?; + let opt_vsn_w = AdamW::new(dev, FEATURE_DIM * FEATURE_DIM, cfg.lr_cfc)?; + let mut opt_vsn_b = AdamW::new(dev, FEATURE_DIM, cfg.lr_cfc)?; + opt_vsn_b.wd = 0.0; + let k = cfg.seq_len; Ok(Self { cfg: cfg.clone(), @@ -592,6 +641,20 @@ impl PerceptionTrainer { _ln_module: ln_module, opt_ln_gain, opt_ln_bias, + // VSN (Phase 2D) — params + per-K gates + scratch. + vsn_w_d, + vsn_b_d, + vsn_out_d: GpuTensor::zeros(&[cfg.n_batch, k, FEATURE_DIM], &stream) + .map_err(|e| anyhow::anyhow!("vsn_out_d alloc: {e}"))?, + vsn_gates_d: stream.alloc_zeros::(cfg.n_batch * k * FEATURE_DIM)?, + vsn_grad_x_d: stream.alloc_zeros::(cfg.n_batch * k * FEATURE_DIM)?, + grad_vsn_w_d: stream.alloc_zeros::(FEATURE_DIM * FEATURE_DIM)?, + 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, loss_ema_d: stream.alloc_zeros::(N_HORIZONS)?, lambda_d: stream.alloc_zeros::(N_HORIZONS)?, horizon_lambda_fn, @@ -1017,9 +1080,35 @@ impl PerceptionTrainer { // No sync — stream-ordered. A sync would trip // STREAM_CAPTURE_ISOLATION inside the captured region. + // ── 1b. VSN forward — per-position softmax-normalised feature + // gating over `window_tensor_d` [B, K, FEATURE_DIM]. + // Writes gated features to `vsn_out_d` and saved + // post-softmax gates to `vsn_gates_d` for backward. + // Mamba2 then consumes `vsn_out_d` instead of the raw + // snap features. + { + let n_rows_vsn: i32 = (b_sz * k_seq) as i32; + let cfg_vsn = LaunchConfig { + grid_dim: (n_rows_vsn as u32, 1, 1), + block_dim: (64, 1, 1), // VSN_BLOCK = 64 (1 warp + idle threads) + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.vsn_fwd_fn); + launch + .arg(&self.vsn_w_d) + .arg(&self.vsn_b_d) + .arg(self.window_tensor_d.cuda_data()) + .arg(&n_rows_vsn) + .arg(self.vsn_out_d.data_mut()) + .arg(&mut self.vsn_gates_d); + unsafe { launch.launch(cfg_vsn).context("variable_selection_fwd")?; } + } + // ── 2. Mamba2 per-step forward — writes into self.mamba2_fwd_scratch. + // Now consumes `vsn_out_d` (gated features) instead of the + // raw `window_tensor_d`. self.mamba2 - .forward_train_seq_into(&self.window_tensor_d, &mut self.mamba2_fwd_scratch) + .forward_train_seq_into(&self.vsn_out_d, &mut self.mamba2_fwd_scratch) .context("mamba2 forward_train_seq_into")?; // ── 2a. LayerNorm forward over h_enriched_seq [B*K, H]. @@ -1464,11 +1553,15 @@ impl PerceptionTrainer { // ── 8. Mamba2 backward — fully pre-allocated path. Writes all // grads into self.mamba2_grads_buffers; no allocation. - // Now consumes grad_ln_in_d (LN bwd output), NOT the raw - // grad_h_enriched_seq_d — LN sits between Mamba2 and CfC. + // Now consumes: + // - input = vsn_out_d (gated snap_features fed to Mamba2 fwd) + // - d_h_out = grad_ln_in_d (LN bwd output) + // and writes `mamba2_grads_buffers.d_x_from_in` = + // gradient w.r.t. Mamba2's input = gradient on VSN's + // output, which VSN bwd consumes below. self.mamba2 .backward_from_h_enriched_seq_full_into( - &self.window_tensor_d, + &self.vsn_out_d, &self.mamba2_fwd_scratch, &self.grad_ln_in_d, &mut self.mamba2_bwd_scratch, @@ -1476,6 +1569,40 @@ impl PerceptionTrainer { ) .context("mamba2 backward_from_h_enriched_seq_full_into")?; + // ── 8b. VSN backward. Consumes: + // grad_y = mamba2_grads_buffers.d_x_from_in [B*K, FEATURE_DIM] + // x = window_tensor_d [B, K, FEATURE_DIM] + // gates = vsn_gates_d (saved by fwd) [B*K, FEATURE_DIM] + // Writes: + // grad_W_vsn, grad_b_vsn — Adam consumes. + // vsn_grad_x_d — discarded (snap_features non-trainable). + // Param grads OVERWRITE (kernel uses += but the prior + // memsets clear them at step start, so a single VSN bwd + // per step writes the full accumulation cleanly). + self.stream.memset_zeros(&mut self.grad_vsn_w_d) + .map_err(|e| anyhow::anyhow!("zero grad_vsn_w: {e}"))?; + self.stream.memset_zeros(&mut self.grad_vsn_b_d) + .map_err(|e| anyhow::anyhow!("zero grad_vsn_b: {e}"))?; + { + let n_rows_vsn: i32 = (b_sz * k_seq) as i32; + let cfg_vsn_bwd = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (64, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.vsn_bwd_fn); + launch + .arg(&self.vsn_w_d) + .arg(self.window_tensor_d.cuda_data()) + .arg(&self.vsn_gates_d) + .arg(self.mamba2_grads_buffers.d_x_from_in.cuda_data()) + .arg(&n_rows_vsn) + .arg(&mut self.grad_vsn_w_d) + .arg(&mut self.grad_vsn_b_d) + .arg(&mut self.vsn_grad_x_d); + unsafe { launch.launch(cfg_vsn_bwd).context("variable_selection_bwd")?; } + } + // ── 9. Apply AdamW updates on all 17 param groups: CfC×4 + // GRN heads×10 + LN×2 + Mamba2 grouped. self.opt_w_in.step(&mut self.w_in_d, &self.grad_w_in_d)?; @@ -1494,6 +1621,8 @@ impl PerceptionTrainer { self.opt_heads_b_skip.step(&mut self.heads_b_skip_d, &self.grad_heads_b_skip_d)?; self.opt_ln_gain.step(&mut self.ln_gain_d, &self.grad_ln_gain_d)?; self.opt_ln_bias.step(&mut self.ln_bias_d, &self.grad_ln_bias_d)?; + self.opt_vsn_w.step(&mut self.vsn_w_d, &self.grad_vsn_w_d)?; + self.opt_vsn_b.step(&mut self.vsn_b_d, &self.grad_vsn_b_d)?; // GPU-resident grad-clip + AdamW (zero host roundtrips). // Replaces step_from_buffers which did 9× memcpy_dtoh per step. self.mamba2_adamw @@ -1651,9 +1780,28 @@ impl PerceptionTrainer { } self.stream.synchronize().context("eval snap_batched sync")?; - // Mamba2 fwd into pre-allocated fwd_scratch. + // VSN fwd — same as training path. Eval consumes vsn_out_d. + { + let n_rows_vsn: i32 = (b_sz * k_seq) as i32; + let cfg_vsn = LaunchConfig { + grid_dim: (n_rows_vsn as u32, 1, 1), + block_dim: (64, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.vsn_fwd_fn); + launch + .arg(&self.vsn_w_d) + .arg(&self.vsn_b_d) + .arg(self.window_tensor_d.cuda_data()) + .arg(&n_rows_vsn) + .arg(self.vsn_out_d.data_mut()) + .arg(&mut self.vsn_gates_d); + unsafe { launch.launch(cfg_vsn).context("eval variable_selection_fwd")?; } + } + + // Mamba2 fwd into pre-allocated fwd_scratch — reads vsn_out_d. self.mamba2 - .forward_train_seq_into(&self.window_tensor_d, &mut self.mamba2_fwd_scratch) + .forward_train_seq_into(&self.vsn_out_d, &mut self.mamba2_fwd_scratch) .context("eval mamba2 fwd_into")?; // LN forward over h_enriched_seq [B*K, H] — same as training path