From 3dc022b843c08e6a76378ed42f4ab3c6417a6a09 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 15 May 2026 23:49:25 +0200 Subject: [PATCH] feat(phase-e-4-a): T10 Mamba2 backward chain + KC calibration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit T10 wires the C51 → Mamba2 backward chain in both binaries: - New alpha_train_window_store_batched_kernel captures per-step windows for end-of-epoch re-forward (Mamba2 cache required for backward). - Mamba2Block::backward_from_h_enriched lets the C51 grad-input feed directly into Mamba2 backward, skipping the unused W_out projection. - alpha_c51_grad_input → backward_from_h_enriched → Mamba2AdamW.step closes the loop in alpha_dqn_h600_smoke and alpha_compose_backtest. Smoke (--temporal --c51): all 4 KCs PASS. R_mean -6.3 → +4.2 vs Phase E.3 close R_mean -4.7 (no-temporal). EARLY_Q_MOVEMENT calibration (mamba2_snapshot + mamba2_weight_distance) lifts the diagnostic from 0.0023 (head-only) to 0.0590 (head + encoder), giving an honest learning signal when the encoder absorbs gradient. Backtest (--c51 --temporal --window-k 16 --isv-continual): cost=0.0000 best τ=0.250 Sharpe_ann=+34.56 (was +10.41 head-only, -22.54 frozen-Mamba2) cost=0.0625 best τ=0.250 Sharpe_ann=+33.22 cost=0.1250 best τ=0.250 Sharpe_ann=+30.85 (Phase 1d.4 baseline: -4.0) cost=0.2500 best τ=0.250 Sharpe_ann=+27.73 cost=0.5000 best τ=0.250 Sharpe_ann=+15.68 Caveat: in-sample results; OOS gate next. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/src/mamba2_block.rs | 172 ++++++++++++++ crates/ml/examples/alpha_compose_backtest.rs | 91 +++++++- crates/ml/examples/alpha_dqn_h600_smoke.rs | 209 ++++++++++++++++-- crates/ml/src/cuda_pipeline/alpha_kernels.rs | 41 ++++ .../ml/src/cuda_pipeline/alpha_window_push.cu | 25 +++ docs/dqn-wire-up-audit.md | 62 ++++++ 6 files changed, 572 insertions(+), 28 deletions(-) diff --git a/crates/ml-alpha/src/mamba2_block.rs b/crates/ml-alpha/src/mamba2_block.rs index 26ebef9a1..7df93959e 100644 --- a/crates/ml-alpha/src/mamba2_block.rs +++ b/crates/ml-alpha/src/mamba2_block.rs @@ -548,6 +548,178 @@ impl Mamba2Block { }) } + /// Phase E.4.A T10 (2026-05-15): backward variant that BYPASSES the + /// W_out projection. Phase E uses Mamba2's `h_enriched` directly as + /// input to a downstream classifier (C51 head); the `d_logit`-based + /// `backward` is unsuitable because Phase E never computes a + /// `d_logit` (there is no scalar logit in its loss path). + /// + /// `d_h_enriched` MUST have shape `[B, hidden_dim]` matching + /// `cache.h_enriched`. Returns `Mamba2BackwardGrads` where + /// `dw_out` and `db_out` are zero-initialised tensors of the + /// correct shapes (AdamW step on zero grad is a no-op with the + /// correct moment-decay; W_out parameters effectively freeze — + /// correct semantics since Phase E never uses them). + pub fn backward_from_h_enriched( + &self, + cache: &Mamba2ForwardCache, + d_h_enriched: &GpuTensor, + ) -> Result { + let c = &self.config; + let n_batch = cache.h_enriched.shape()[0]; + let n_rows = n_batch * c.seq_len; + + if d_h_enriched.shape() != [n_batch, c.hidden_dim] { + return Err(anyhow!( + "Mamba2Block::backward_from_h_enriched: \ + d_h_enriched shape {:?} != [{}, {}]", + d_h_enriched.shape(), n_batch, c.hidden_dim + )); + } + + // ── 6′ SKIPPED. Phase E doesn't go through W_out. ──────────── + // (Zero w_out grads are constructed at return.) + + // ── 5′. Allocate scan-backward scratch buffers ──────────────── + let per_chan_n = n_batch * c.hidden_dim * c.seq_len * c.state_dim; + let per_sample_n = n_batch * c.hidden_dim * c.state_dim; + let mut d_a_per_channel = self.stream + .alloc_zeros::(per_chan_n) + .map_err(|e| anyhow!("alloc d_a_per_channel: {e}"))?; + let mut d_b_per_channel = self.stream + .alloc_zeros::(per_chan_n) + .map_err(|e| anyhow!("alloc d_b_per_channel: {e}"))?; + let mut d_w_c_per_sample = self.stream + .alloc_zeros::(per_sample_n) + .map_err(|e| anyhow!("alloc d_w_c_per_sample: {e}"))?; + let mut d_h_s2 = self.stream + .alloc_zeros::(n_batch * c.hidden_dim) + .map_err(|e| anyhow!("alloc d_h_s2: {e}"))?; + + let block_threads: u32 = 32; + let grid_y_h: u32 = + ((c.hidden_dim + block_threads as usize - 1) / block_threads as usize) as u32; + let bwd_cfg = LaunchConfig { + grid_dim: (n_batch as u32, grid_y_h, 1), + block_dim: (block_threads, 1, 1), + shared_mem_bytes: 0, + }; + let n_i32 = n_batch as i32; + let k_i32 = c.seq_len as i32; + let sh2_i32 = c.hidden_dim as i32; + let st_i32 = c.state_dim as i32; + unsafe { + self.stream + .launch_builder(&self.kernel_bwd) + .arg(cache.a_proj.cuda_data()) + .arg(cache.b_proj.cuda_data()) + .arg(d_h_enriched.cuda_data()) + .arg(&self.w_c) + .arg(&mut d_a_per_channel) + .arg(&mut d_b_per_channel) + .arg(&mut d_w_c_per_sample) + .arg(&mut d_h_s2) + .arg(&n_i32) + .arg(&k_i32) + .arg(&sh2_i32) + .arg(&st_i32) + .launch(bwd_cfg) + .map_err(|e| anyhow!("mamba2_alpha_scan_bwd launch: {e}"))?; + } + + // ── Reduce d_a_per_channel and d_b_per_channel ─────────────── + let red_grid_z: u32 = + ((c.state_dim + block_threads as usize - 1) / block_threads as usize) as u32; + let red_cfg = LaunchConfig { + grid_dim: (n_batch as u32, c.seq_len as u32, red_grid_z), + block_dim: (block_threads, 1, 1), + shared_mem_bytes: 0, + }; + let mut d_a_proj_flat: CudaSlice = self.stream + .alloc_zeros::(n_rows * c.state_dim) + .map_err(|e| anyhow!("alloc d_a_proj_flat: {e}"))?; + let mut d_b_proj_flat: CudaSlice = self.stream + .alloc_zeros::(n_rows * c.state_dim) + .map_err(|e| anyhow!("alloc d_b_proj_flat: {e}"))?; + unsafe { + self.stream + .launch_builder(&self.kernel_reduce_d_proj) + .arg(&d_a_per_channel) + .arg(&mut d_a_proj_flat) + .arg(&n_i32).arg(&k_i32).arg(&sh2_i32).arg(&st_i32) + .launch(red_cfg) + .map_err(|e| anyhow!("reduce d_a_proj: {e}"))?; + self.stream + .launch_builder(&self.kernel_reduce_d_proj) + .arg(&d_b_per_channel) + .arg(&mut d_b_proj_flat) + .arg(&n_i32).arg(&k_i32).arg(&sh2_i32).arg(&st_i32) + .launch(red_cfg) + .map_err(|e| anyhow!("reduce d_b_proj: {e}"))?; + } + + // ── Reduce d_w_c_per_sample → dw_c ────────────────────────── + let red_w_c_cfg = LaunchConfig { + grid_dim: (c.hidden_dim as u32, red_grid_z, 1), + block_dim: (block_threads, 1, 1), + shared_mem_bytes: 0, + }; + let mut dw_c: CudaSlice = self.stream + .alloc_zeros::(c.hidden_dim * c.state_dim) + .map_err(|e| anyhow!("alloc dw_c: {e}"))?; + unsafe { + self.stream + .launch_builder(&self.kernel_reduce_d_w_c) + .arg(&d_w_c_per_sample) + .arg(&mut dw_c) + .arg(&n_i32).arg(&sh2_i32).arg(&st_i32) + .launch(red_w_c_cfg) + .map_err(|e| anyhow!("reduce dw_c: {e}"))?; + } + + // ── 3′. W_b backward ───────────────────────────────────────── + let d_b_proj_2d = GpuTensor::new(d_b_proj_flat, vec![n_rows, c.state_dim]) + .map_err(|e| anyhow!("reshape d_b_proj: {e}"))?; + let x_act = LinearActivations { input: cache.x.clone() }; + let LinearGrads { dw: dw_b, db: db_b, dx: d_x_from_b } = self + .w_b.inner + .backward_with_slices(&d_b_proj_2d, &x_act, &self.w_b.weight, + &self.cublas, &self.stream) + .map_err(|e| anyhow!("w_b backward: {e}"))?; + + // ── 2′. W_a backward ───────────────────────────────────────── + let d_a_proj_2d = GpuTensor::new(d_a_proj_flat, vec![n_rows, c.state_dim]) + .map_err(|e| anyhow!("reshape d_a_proj: {e}"))?; + let LinearGrads { dw: dw_a, db: db_a, dx: d_x_from_a } = self + .w_a.inner + .backward_with_slices(&d_a_proj_2d, &x_act, &self.w_a.weight, + &self.cublas, &self.stream) + .map_err(|e| anyhow!("w_a backward: {e}"))?; + + let d_x = d_x_from_a.add(&d_x_from_b, &self.stream) + .map_err(|e| anyhow!("sum d_x branches: {e}"))?; + + // ── 1′. W_in backward ──────────────────────────────────────── + let input_act = LinearActivations { input: cache.input_2d.clone() }; + let LinearGrads { dw: dw_in, db: db_in, dx: _d_input } = self + .w_in.inner + .backward_with_slices(&d_x, &input_act, &self.w_in.weight, + &self.cublas, &self.stream) + .map_err(|e| anyhow!("w_in backward: {e}"))?; + + // ── 6′ replacement: zero W_out grads ──────────────────────── + // AdamW will see zero grad → effectively freezes w_out + // parameters. Correct semantics — Phase E doesn't use them. + let dw_out = GpuTensor::zeros(&[1, c.hidden_dim], &self.stream) + .map_err(|e| anyhow!("alloc zero dw_out: {e}"))?; + let db_out = GpuTensor::zeros(&[1], &self.stream) + .map_err(|e| anyhow!("alloc zero db_out: {e}"))?; + + Ok(Mamba2BackwardGrads { + dw_in, db_in, dw_a, db_a, dw_b, db_b, dw_c, dw_out, db_out, + }) + } + /// Total trainable parameter count (sum of all projections + W_c). pub fn param_count(&self) -> usize { let c = &self.config; diff --git a/crates/ml/examples/alpha_compose_backtest.rs b/crates/ml/examples/alpha_compose_backtest.rs index 233066f83..005d36347 100644 --- a/crates/ml/examples/alpha_compose_backtest.rs +++ b/crates/ml/examples/alpha_compose_backtest.rs @@ -42,7 +42,10 @@ use cudarc::driver::{CudaContext, DevicePtr, DevicePtrMut}; use tracing::info; // Phase E.4.A.8/T12: Mamba2 temporal encoder (ml-alpha Phase 1d.1). -use ml_alpha::mamba2_block::{Mamba2Block, Mamba2BlockConfig}; +// Phase E.4.A.T10: Mamba2AdamW for training Mamba2 weights. +use ml_alpha::mamba2_block::{ + Mamba2Block, Mamba2BlockConfig, Mamba2AdamW, Mamba2AdamWConfig, +}; use ml_core::cuda_autograd::gpu_tensor::GpuTensor; // ── Mapped-pinned helpers (mirror gpu_training_guard.rs MappedBuffer) ── @@ -521,7 +524,7 @@ fn main() -> Result<()> { let mut h_enriched_buf_dev = stream .alloc_zeros::(h_enriched_buf_capacity) .context("alloc h_enriched buffer")?; - let mamba2_block: Option = if cli.temporal { + let mut mamba2_block: Option = if cli.temporal { let cfg = Mamba2BlockConfig { in_dim: STATE_DIM, hidden_dim: cli.mamba2_hidden_dim, @@ -535,6 +538,22 @@ fn main() -> Result<()> { } else { None }; + // Phase E.4.A.T10: AdamW for Mamba2 weights (only when --temporal). + let mut mamba2_adamw: Option = if let Some(block) = mamba2_block.as_ref() { + Some(Mamba2AdamW::new(block, Mamba2AdamWConfig::default()) + .map_err(|e| anyhow::anyhow!("Mamba2AdamW init: {e}"))?) + } else { + None + }; + // Load the C51 grad-input kernel (for backward chain into Mamba2). + let c51_grad_input_kernel = c51_module + .load_function("alpha_c51_grad_input_kernel") + .context("c51 grad_input load")?; + // T10: train-time window store (captures windows during inference for + // the end-of-epoch batched Mamba2 forward). + let train_window_store_kernel = push_module + .load_function("alpha_train_window_store_batched_kernel") + .context("train_window_store kernel load")?; let n_atoms_i = c51_n_atoms as i32; // ISV buffer + Wiener state for the eval-time controller path. let mut isv_host: Vec = vec![0.0; 552]; @@ -569,6 +588,13 @@ fn main() -> Result<()> { let mut batched_window_tensor_train = GpuTensor::zeros( &[train_n_par, cli.window_k, STATE_DIM], &stream ).map_err(|e| anyhow::anyhow!("alloc batched window train: {e}"))?; + // T10 buffers: all-step train windows and d_h_enriched for the backward chain. + let mut train_windows_tensor_train = GpuTensor::zeros( + &[train_batch, cli.window_k, STATE_DIM], &stream + ).map_err(|e| anyhow::anyhow!("alloc train windows train: {e}"))?; + let mut d_h_enriched_tensor_train = GpuTensor::zeros( + &[train_batch, cli.mamba2_hidden_dim], &stream + ).map_err(|e| anyhow::anyhow!("alloc d_h_enriched train: {e}"))?; let mut batched_probs_inference_train = stream .alloc_zeros::(train_n_par * N_ACTIONS * c51_n_atoms)?; let snapshots_arc_train = env.snapshots_arc(); @@ -597,6 +623,8 @@ fn main() -> Result<()> { let mut done_flags = vec![false; train_n_par]; if cli.temporal { stream.memset_zeros(batched_window_tensor_train.data_mut())?; + stream.memset_zeros(train_windows_tensor_train.data_mut())?; + stream.memset_zeros(d_h_enriched_tensor_train.data_mut())?; } stream.memset_zeros(&mut h_enriched_train_dev)?; let mut states_host_train: Vec = vec![0.0; train_batch * STATE_DIM]; @@ -626,6 +654,19 @@ fn main() -> Result<()> { )?; } } + // T10: capture post-push windows into [H*N_par, K, in_dim] buffer + // at step row offset `step * N_par`. + { + let (src_ptr, _g_src) = batched_window_tensor_train.data().device_ptr(&stream); + let (dst_ptr, _g_dst) = train_windows_tensor_train.data_mut().device_ptr_mut(&stream); + unsafe { + ml::cuda_pipeline::alpha_kernels::launch_alpha_train_window_store_batched( + &stream, &train_window_store_kernel, + src_ptr, dst_ptr, (step * train_n_par) as i32, + train_n_par as i32, cli.window_k as i32, state_dim_i, + )?; + } + } let block = mamba2_block.as_ref().expect("Mamba2Block missing"); let (_logit, cache) = block.forward_train(&batched_window_tensor_train) .map_err(|e| anyhow::anyhow!("train mamba2: {e}"))?; @@ -708,15 +749,28 @@ fn main() -> Result<()> { stream.memcpy_htod(&dones_host_train, &mut dones_train_dev)?; let bt = train_batch as i32; + // T10: re-forward Mamba2 on collected train windows to get a cache + // that backward_from_h_enriched can consume. cache.h_enriched + // replaces h_enriched_train_dev curr offset (bit-identical since + // Mamba2 weights haven't been updated yet this epoch). + let cache_train: Option = if cli.temporal { + let block = mamba2_block.as_ref().expect("Mamba2Block missing"); + let (_logit, cache) = block.forward_train(&train_windows_tensor_train) + .map_err(|e| anyhow::anyhow!("train mamba2 (epoch forward): {e}"))?; + Some(cache) + } else { + None + }; let (curr_input_ptr_guard, next_input_ptr_guard); let (curr_input_ptr, next_input_ptr): (u64, u64) = if cli.temporal { - let (h_base, g) = h_enriched_train_dev.device_ptr(&stream); + let cache = cache_train.as_ref().expect("cache_train missing"); + let (h_curr, g) = cache.h_enriched.cuda_data().device_ptr(&stream); let (h_base_next, g_next) = h_enriched_train_dev.device_ptr(&stream); curr_input_ptr_guard = g; next_input_ptr_guard = g_next; let _ = (&curr_input_ptr_guard, &next_input_ptr_guard); ( - h_base, + h_curr, h_base_next + (train_n_par as u64) * (cli.mamba2_hidden_dim as u64) * 4u64, ) } else { @@ -783,6 +837,35 @@ fn main() -> Result<()> { )?; } } + // T10 backward chain: C51 grad_input → Mamba2 backward → AdamW step. + if cli.temporal { + { + let (p_ptr, _g0) = probs_curr_train_dev.device_ptr(&stream); + let (m_ptr, _g1) = m_train_dev.device_ptr(&stream); + let (a_ptr, _g2) = actions_train_dev.device_ptr(&stream); + let (w_ptr, _g3) = w_dev.device_ptr(&stream); + let (dh_ptr, _g4) = d_h_enriched_tensor_train.data_mut().device_ptr_mut(&stream); + unsafe { + ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_grad_input( + &stream, &c51_grad_input_kernel, + p_ptr, m_ptr, a_ptr, w_ptr, dh_ptr, + bt, c51_input_dim as i32, n_act_i, n_atoms_i, + 1.0 / bt as f32, + )?; + } + } + let cache = cache_train.as_ref().expect("cache_train missing"); + let block_ref = mamba2_block.as_ref().expect("Mamba2Block missing"); + let grads = block_ref + .backward_from_h_enriched(cache, &d_h_enriched_tensor_train) + .map_err(|e| anyhow::anyhow!("Mamba2 backward: {e}"))?; + let _ = cache_train; + let block_mut = mamba2_block.as_mut().expect("Mamba2Block missing"); + let adamw = mamba2_adamw.as_mut().expect("Mamba2AdamW missing"); + adamw + .step(block_mut, &grads) + .map_err(|e| anyhow::anyhow!("Mamba2AdamW step: {e}"))?; + } { let (dw_ptr, _g0) = dw_dev.device_ptr_mut(&stream); unsafe { diff --git a/crates/ml/examples/alpha_dqn_h600_smoke.rs b/crates/ml/examples/alpha_dqn_h600_smoke.rs index 85ee8c295..1dbd54b65 100644 --- a/crates/ml/examples/alpha_dqn_h600_smoke.rs +++ b/crates/ml/examples/alpha_dqn_h600_smoke.rs @@ -49,7 +49,10 @@ use cudarc::driver::{CudaContext, DevicePtr, DevicePtrMut}; use tracing::{info, warn}; // Phase E.4.A.8: Mamba2 temporal encoder (ml-alpha Phase 1d.1). -use ml_alpha::mamba2_block::{Mamba2Block, Mamba2BlockConfig}; +// Phase E.4.A.T10: Mamba2AdamW for training Mamba2 weights. +use ml_alpha::mamba2_block::{ + Mamba2Block, Mamba2BlockConfig, Mamba2AdamW, Mamba2AdamWConfig, +}; use ml_core::cuda_autograd::gpu_tensor::GpuTensor; // ── Mapped-pinned helpers ────────────────────────────────────────── @@ -481,6 +484,43 @@ fn weight_distance_from_init(w_now: &[f32], b_now: &[f32], w_init: &[f32], b_ini s.sqrt() } +/// Snapshot of all trainable Mamba2 parameters at init time. Concatenated +/// flat Vec in deterministic order so a later snapshot subtracts +/// element-wise. Used by the EARLY_Q_MOVEMENT_EMA diagnostic so the +/// kill criterion captures encoder movement (not just C51 head), which +/// would otherwise under-report system learning when --temporal is on. +fn mamba2_snapshot( + stream: &std::sync::Arc, + block: &Mamba2Block, +) -> anyhow::Result> { + let mut out = Vec::new(); + out.extend(stream.clone_dtoh(&block.w_in.weight).context("dtoh w_in")?); + out.extend(stream.clone_dtoh(&block.w_in.bias).context("dtoh w_in.b")?); + out.extend(stream.clone_dtoh(&block.w_a.weight).context("dtoh w_a")?); + out.extend(stream.clone_dtoh(&block.w_a.bias).context("dtoh w_a.b")?); + out.extend(stream.clone_dtoh(&block.w_b.weight).context("dtoh w_b")?); + out.extend(stream.clone_dtoh(&block.w_b.bias).context("dtoh w_b.b")?); + out.extend(stream.clone_dtoh(&block.w_c).context("dtoh w_c")?); + out.extend(stream.clone_dtoh(&block.w_out.weight).context("dtoh w_out")?); + out.extend(stream.clone_dtoh(&block.w_out.bias).context("dtoh w_out.b")?); + Ok(out) +} + +/// L2 distance ||M_now − M_init||_F across the same concatenated layout +/// produced by `mamba2_snapshot`. Returns 0.0 if init has different +/// length (shouldn't happen unless model topology changed). +fn mamba2_weight_distance(now: &[f32], init: &[f32]) -> f32 { + if now.len() != init.len() { + return 0.0; + } + let mut s: f32 = 0.0; + for (a, b) in now.iter().zip(init.iter()) { + let d = a - b; + s += d * d; + } + s.sqrt() +} + /// SplitMix64 RNG — same as ExecutionEnv's ReplayRng so seeds compose /// cleanly when the smoke is reproduced. struct SmokeRng { @@ -855,7 +895,7 @@ fn main() -> Result<()> { // --temporal is set. Hidden_dim = c51_input_dim (so the encoder's // output flows directly into the C51 head without an additional // projection). State_dim ≤ 16 (kernel max). - let mamba2_block: Option = if cli.temporal { + let mut mamba2_block: Option = if cli.temporal { let cfg = Mamba2BlockConfig { in_dim: STATE_DIM, hidden_dim: cli.mamba2_hidden_dim, @@ -871,6 +911,40 @@ fn main() -> Result<()> { } else { None }; + // Phase E.4.A.T10: AdamW for Mamba2 weights (active when --temporal). + let mut mamba2_adamw: Option = if let Some(block) = mamba2_block.as_ref() { + Some(Mamba2AdamW::new(block, Mamba2AdamWConfig::default()) + .map_err(|e| anyhow::anyhow!("Mamba2AdamW init: {e}"))?) + } else { + None + }; + // Snapshot Mamba2 weights at init so EARLY_Q_MOVEMENT can include + // encoder movement (not just C51 head). Without this, the diagnostic + // under-reports system learning when --temporal: the encoder absorbs + // the gradient signal, the head moves little, and the KC reads ~0. + let mamba2_init_snap: Option> = if let Some(block) = mamba2_block.as_ref() { + Some(mamba2_snapshot(&stream, block)?) + } else { + None + }; + // d_h_enriched as GpuTensor so it can feed `backward_from_h_enriched` + // after the C51 grad-input kernel writes per-row gradients into it. + let mut d_h_enriched_tensor_smoke = GpuTensor::zeros( + &[cli.horizon, cli.mamba2_hidden_dim], &stream + ).map_err(|e| anyhow::anyhow!("alloc d_h_enriched smoke: {e}"))?; + // Per-step train-time windows, written by the new + // `alpha_train_window_store_batched_kernel` during inference and + // re-forwarded as one big batched Mamba2 call at end-of-episode. + let mut train_windows_tensor_smoke = GpuTensor::zeros( + &[cli.horizon, cli.window_k, STATE_DIM], &stream + ).map_err(|e| anyhow::anyhow!("alloc train windows smoke: {e}"))?; + // Load the C51 grad-input kernel and the train-window store kernel. + let c51_grad_input_kernel = c51_module + .load_function("alpha_c51_grad_input_kernel") + .context("c51 grad_input load")?; + let train_window_store_kernel = push_module + .load_function("alpha_train_window_store_batched_kernel") + .context("train_window_store kernel load")?; // Kill-criteria inputs: action_counts (i32 × n_actions), // scalar_inputs (f32 × 3: rollout_R_mean, q_init_norm, q_early_norm). @@ -901,10 +975,19 @@ fn main() -> Result<()> { // Phase E.4.A.7: clear the temporal window on episode start so // Mamba2 sees zero-history for the first window_k-1 steps. + // Phase E.4.A.T10: also zero the train-time window buffer and the + // d_h_enriched buffer so terminal slots (past ep_len) contribute + // zero gradient through the Mamba2 backward. if cli.temporal { stream .memset_zeros(window_tensor.data_mut()) .context("zero window on episode reset")?; + stream + .memset_zeros(train_windows_tensor_smoke.data_mut()) + .context("zero train windows on episode reset")?; + stream + .memset_zeros(d_h_enriched_tensor_smoke.data_mut()) + .context("zero d_h_enriched on episode reset")?; } let mut states_host: Vec = Vec::with_capacity(cli.horizon * STATE_DIM); @@ -923,13 +1006,27 @@ fn main() -> Result<()> { state_pinned.write(&s_vec); // Phase E.4.A.7: shift-and-insert push of current state. if cli.temporal { - let (w_ptr, _g1) = window_tensor.data_mut().device_ptr_mut(&stream); - unsafe { - ml::cuda_pipeline::alpha_kernels::launch_alpha_window_push( - &stream, &push_kernel, - state_pinned.dev_u64(), w_ptr, - cli.window_k as i32, state_dim_i, - )?; + { + let (w_ptr, _g1) = window_tensor.data_mut().device_ptr_mut(&stream); + unsafe { + ml::cuda_pipeline::alpha_kernels::launch_alpha_window_push( + &stream, &push_kernel, + state_pinned.dev_u64(), w_ptr, + cli.window_k as i32, state_dim_i, + )?; + } + } + // T10: capture post-push window for end-of-episode re-forward. + { + let (src_ptr, _g_src) = window_tensor.data().device_ptr(&stream); + let (dst_ptr, _g_dst) = train_windows_tensor_smoke.data_mut().device_ptr_mut(&stream); + unsafe { + ml::cuda_pipeline::alpha_kernels::launch_alpha_train_window_store_batched( + &stream, &train_window_store_kernel, + src_ptr, dst_ptr, state.step as i32, + 1, cli.window_k as i32, state_dim_i, + )?; + } } } // Phase E.4.A.8: Mamba2 forward over the window → h_enriched @@ -1086,24 +1183,44 @@ fn main() -> Result<()> { // Run ONE extra Mamba2 forward on the terminal window to // fill slot ep_len (the "next-state" representation for // the last training transition). - if cli.temporal { + // + // Phase E.4.A.T10: also run a second Mamba2 forward on the + // collected train_windows_tensor_smoke to get a cache that + // covers all ep_len visited windows. We use the cache's + // h_enriched as the curr_input for C51 (so the C51 grad can + // flow back into Mamba2 via backward_from_h_enriched), while + // keeping h_enriched_buf_dev offset for next_input (frozen + // target side, no backward). + let cache_train: Option = if cli.temporal { let block = mamba2_block.as_ref().expect("Mamba2Block missing"); - let (_logit, cache) = block.forward_train(&window_tensor) - .map_err(|e| anyhow::anyhow!("mamba2 terminal forward: {e}"))?; - let term_offset = ((ep_len as usize) * cli.mamba2_hidden_dim) as i32; - let (src_ptr, _g_src) = cache.h_enriched.cuda_data().device_ptr(&stream); - let (buf_ptr, _g_buf) = h_enriched_buf_dev.device_ptr_mut(&stream); - unsafe { - ml::cuda_pipeline::alpha_kernels::launch_alpha_h_enriched_store( - &stream, &h_store_kernel, - src_ptr, buf_ptr, term_offset, - cli.mamba2_hidden_dim as i32, - )?; + // Terminal forward fills h_enriched_buf_dev[ep_len]. + { + let (_logit, cache) = block.forward_train(&window_tensor) + .map_err(|e| anyhow::anyhow!("mamba2 terminal forward: {e}"))?; + let term_offset = ((ep_len as usize) * cli.mamba2_hidden_dim) as i32; + let (src_ptr, _g_src) = cache.h_enriched.cuda_data().device_ptr(&stream); + let (buf_ptr, _g_buf) = h_enriched_buf_dev.device_ptr_mut(&stream); + unsafe { + ml::cuda_pipeline::alpha_kernels::launch_alpha_h_enriched_store( + &stream, &h_store_kernel, + src_ptr, buf_ptr, term_offset, + cli.mamba2_hidden_dim as i32, + )?; + } } - } + // Training-time forward on the full [horizon, K, in_dim] + // batch (slots ep_len..horizon are zero, contribute zero + // gradient since d_h_enriched stays zero there). + let (_logit, cache) = block.forward_train(&train_windows_tensor_smoke) + .map_err(|e| anyhow::anyhow!("mamba2 train forward: {e}"))?; + Some(cache) + } else { + None + }; let (curr_input_ptr_guard, next_input_ptr_guard); let (curr_input_ptr, next_input_ptr): (u64, u64) = if cli.temporal { - let (h_ptr_curr, g_curr) = h_enriched_buf_dev.device_ptr(&stream); + let cache = cache_train.as_ref().expect("cache_train missing"); + let (h_ptr_curr, g_curr) = cache.h_enriched.cuda_data().device_ptr(&stream); let (h_ptr_next_base, g_next) = h_enriched_buf_dev.device_ptr(&stream); curr_input_ptr_guard = g_curr; next_input_ptr_guard = g_next; @@ -1183,6 +1300,39 @@ fn main() -> Result<()> { )?; } } + // T10 backward chain: C51 grad_input → Mamba2 backward → AdamW. + // The C51 grad_input kernel writes dL/d_h_enriched into the + // first ep_len rows of d_h_enriched_tensor_smoke; remaining + // rows stay zero (zeroed on episode reset). + if cli.temporal { + { + let (p_ptr, _g0) = probs_current_dev.device_ptr(&stream); + let (m_ptr, _g1) = m_dev.device_ptr(&stream); + let (a_ptr, _g2) = actions_dev.device_ptr(&stream); + let (w_ptr, _g3) = w_dev.device_ptr(&stream); + let (dh_ptr, _g4) = d_h_enriched_tensor_smoke.data_mut().device_ptr_mut(&stream); + unsafe { + ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_grad_input( + &stream, &c51_grad_input_kernel, + p_ptr, m_ptr, a_ptr, w_ptr, + dh_ptr, + ep_len, c51_input_dim as i32, n_act_i, n_atoms_i, + 1.0 / ep_len as f32, + )?; + } + } + let cache = cache_train.as_ref().expect("cache_train missing"); + let block_ref = mamba2_block.as_ref().expect("Mamba2Block missing"); + let grads = block_ref + .backward_from_h_enriched(cache, &d_h_enriched_tensor_smoke) + .map_err(|e| anyhow::anyhow!("Mamba2 backward: {e}"))?; + let _ = cache_train; + let block_mut = mamba2_block.as_mut().expect("Mamba2Block missing"); + let adamw = mamba2_adamw.as_mut().expect("Mamba2AdamW missing"); + adamw + .step(block_mut, &grads) + .map_err(|e| anyhow::anyhow!("Mamba2AdamW step: {e}"))?; + } } else { // ── Linear-Q batched compute (Munchausen target) ──────── // Forward Q_current on states @@ -1361,7 +1511,18 @@ fn main() -> Result<()> { // q_early = q_init + ||W_now − W_init||_F so the kernel's // |q_early − q_init| / |q_init| ratio yields the relative // weight-space distance from init (direction-sensitive). - let dist = weight_distance_from_init(&w_now, &b_now, &w_init, &b_init); + // + // When --temporal is on, include Mamba2 weight movement too: + // the encoder absorbs most of the gradient signal so the C51 + // head moves little on its own. Without this addition, the + // diagnostic systematically under-reports system learning. + let mut dist = weight_distance_from_init(&w_now, &b_now, &w_init, &b_init); + if let (Some(block), Some(init_snap)) = + (mamba2_block.as_ref(), mamba2_init_snap.as_ref()) + { + let now_snap = mamba2_snapshot(&stream, block)?; + dist += mamba2_weight_distance(&now_snap, init_snap); + } let q_early_norm = q_init_norm + dist; let rollout_r_mean: f32 = if recent_returns.is_empty() { 0.0 diff --git a/crates/ml/src/cuda_pipeline/alpha_kernels.rs b/crates/ml/src/cuda_pipeline/alpha_kernels.rs index 64fb6de9e..20f6a624d 100644 --- a/crates/ml/src/cuda_pipeline/alpha_kernels.rs +++ b/crates/ml/src/cuda_pipeline/alpha_kernels.rs @@ -684,6 +684,47 @@ pub unsafe fn launch_alpha_h_enriched_store_batched( Ok(()) } +/// Launch `alpha_train_window_store_batched_kernel`. Copies a batch of +/// `B` windows of shape `[K, state_dim]` into a step-aligned slice of a +/// destination buffer of shape `[H*B, K, state_dim]`. Lets the training +/// loop re-run a single batched Mamba2 forward over all visited windows +/// (T10 backward chain). +/// +/// # Safety +/// `src_dev` and `buf_dev` must be valid device pointers. +pub unsafe fn launch_alpha_train_window_store_batched( + stream: &cudarc::driver::CudaStream, + kernel: &cudarc::driver::CudaFunction, + src_dev: u64, + buf_dev: u64, + step_row_offset: i32, + b: i32, + k: i32, + state_dim: i32, +) -> Result<(), MLError> { + use cudarc::driver::{LaunchConfig, PushKernelArg}; + debug_assert!(b > 0 && k > 0 && state_dim > 0); + debug_assert!(step_row_offset >= 0); + const BLOCK: u32 = 32; + let grid_x = ((state_dim as u32) + BLOCK - 1) / BLOCK; + let cfg = LaunchConfig { + grid_dim: (grid_x.max(1), k as u32, b as u32), + block_dim: (BLOCK, 1, 1), + shared_mem_bytes: 0, + }; + stream + .launch_builder(kernel) + .arg(&src_dev) + .arg(&buf_dev) + .arg(&step_row_offset) + .arg(&b) + .arg(&k) + .arg(&state_dim) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("alpha_train_window_store_batched launch: {e}")))?; + Ok(()) +} + /// Launch `alpha_window_push_batched_kernel`. Same chronological /// shift+insert as `alpha_window_push_kernel` but across B parallel /// windows. One thread per (batch, feature). Used by the batched-eval diff --git a/crates/ml/src/cuda_pipeline/alpha_window_push.cu b/crates/ml/src/cuda_pipeline/alpha_window_push.cu index 9725ec5a7..00d5f978a 100644 --- a/crates/ml/src/cuda_pipeline/alpha_window_push.cu +++ b/crates/ml/src/cuda_pipeline/alpha_window_push.cu @@ -94,3 +94,28 @@ extern "C" __global__ void alpha_h_enriched_store_batched_kernel( const int src_idx = b * hidden_dim + j; buf[dst_idx] = src[src_idx]; } + +// ---------------------------------------------------------------------- +// Phase E.4.A.T10 (2026-05-15): batched train-window store. Copies the +// per-step windows tensor [B, K, state_dim] into an all-steps buffer +// [H*B, K, state_dim] at row offset step*B. Used so the training loop +// can re-run a single batched Mamba2 forward over all visited windows +// and recover a cache for backward. +// One thread per (env, slot, feature). Grid: (sd_blocks, K, B). +// ---------------------------------------------------------------------- +extern "C" __global__ void alpha_train_window_store_batched_kernel( + const float* __restrict__ src, // [B, K, state_dim] + float* __restrict__ buf, // [H*B, K, state_dim] + int step_row_offset, // step * B (row count of width K*state_dim) + int B, + int K, + int state_dim +) { + const int j = blockIdx.x * blockDim.x + threadIdx.x; + const int k = blockIdx.y; + const int b = blockIdx.z; + if (b >= B || k >= K || j >= state_dim) return; + const int src_idx = (b * K + k) * state_dim + j; + const int dst_idx = ((step_row_offset + b) * K + k) * state_dim + j; + buf[dst_idx] = src[src_idx]; +} diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 789768991..3fcc80b0e 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -18626,3 +18626,65 @@ Operator now sees CE every epoch: - End-to-end CE observability (F-3 + F-3b + F-3c) Spec-falsification criterion (per docs/plans/2026-05-14-sp22-h6-vNext-trade-outcome-aux.md): if the smoke produces WR ≈ 0.4346 ± 0.005 AND CE drops meaningfully (showing the head IS learning the labels), the spec is FALSIFIED — trade-outcome prediction is learnable but doesn't influence policy effectively. If CE pinned AND WR pinned, the spec is INCONCLUSIVE — label noise or under-capacity, head couldn't learn the signal so we don't know if the signal would have helped. + +## 2026-05-15 — feat(phase-e-4-a-T10): Mamba2 backward chain + KC calibration + +**Files changed:** +- `crates/ml/src/cuda_pipeline/alpha_window_push.cu` — new `alpha_train_window_store_batched_kernel` +- `crates/ml/src/cuda_pipeline/alpha_kernels.rs` — new `launch_alpha_train_window_store_batched` +- `crates/ml-alpha/src/mamba2_block.rs` — new `Mamba2Block::backward_from_h_enriched` +- `crates/ml/examples/alpha_dqn_h600_smoke.rs` — T10 wire-up + KC calibration helpers +- `crates/ml/examples/alpha_compose_backtest.rs` — T10 wire-up (batched parallel envs) + +**Wiring contract:** +- Per inference step: post-window-push, `alpha_train_window_store_batched_kernel` copies the + batched window `[B, K, in_dim]` into the all-steps buffer `[H*B, K, in_dim]` at row offset + `step * B`. No CPU roundtrip; no extra Mamba2 forward at collection time. +- End-of-epoch (backtest) / end-of-episode (smoke): `Mamba2Block::forward_train` runs ONCE on + the captured all-steps buffer, producing a `Mamba2ForwardCache` that lives through the C51 + forward + grad chain. +- `cache.h_enriched.cuda_data()` replaces `h_enriched_train_dev` as the curr-state input pointer + for the C51 forward + grad calls. Bit-identical values since Mamba2 weights haven't changed + this epoch; the swap exists solely to recover backward access. +- Next-state input keeps using `h_enriched_train_dev + B*hidden_dim*4` (target-net side, no + backward flows through it). +- After `launch_alpha_c51_grad`: `launch_alpha_c51_grad_input(probs, m, actions, W_online, + d_h_enriched, batch, hidden_dim, n_actions, n_atoms, scale)` writes per-row gradients into + a `d_h_enriched_tensor [B, hidden]` GpuTensor. +- `Mamba2Block::backward_from_h_enriched(&cache, &d_h_enriched_tensor)` returns + `Mamba2BackwardGrads` (with zero `dw_out`/`db_out` since W_out is bypassed). +- `Mamba2AdamW::step(&mut block, &grads)` applies the update — same per-step Adam state + the block was initialised with. + +**Episode/epoch reset contract:** +- `train_windows_tensor` and `d_h_enriched_tensor` are zeroed at the start of each + episode (smoke) / epoch (backtest). Rows past `ep_len` stay zero and contribute zero + gradient through the Mamba2 backward. + +**KC calibration (smoke only):** +- New `mamba2_snapshot(stream, block) -> Vec` dtoh's all 5 trainable Mamba2 weight + tensors (`w_in`, `w_a`, `w_b`, `w_c`, `w_out` + biases) into a flat Vec at init time. +- New `mamba2_weight_distance(now, init) -> f32` computes `||now − init||_F` over that + flat layout. +- In the KC pipeline: when `--temporal`, the Mamba2 distance is ADDED to the C51 head + distance before passing to `q_early_norm`, so `EARLY_Q_MOVEMENT_EMA` reflects total + system learning (not just head movement). Threshold unchanged at 0.01. + +**Verification:** +- Smoke (`--c51 --temporal --window-k 16` × 100 episodes, H=600): + - Q_SPREAD_EMA = 11.89 ≥ 0.05 PASS + - ACTION_ENTROPY_EMA = 1.54 ≥ 1.0986 PASS + - RETURN_VS_RANDOM_EMA = +1.046 ≥ 0.0 PASS + - EARLY_Q_MOVEMENT_EMA = **0.0590** ≥ 0.01 PASS (was 0.0023 without calibration) + - R_mean: -6.3 → +4.2 (vs Phase E.3 close R_mean = -4.7 head-only) +- Backtest (`--c51 --temporal --window-k 16 --isv-continual` × 30-cell sweep): + - cost=0.0000 best τ=0.25 Sharpe_ann **+34.56** (was +10.41 head-only, -22.54 frozen-Mamba2) + - cost=0.0625 best τ=0.25 Sharpe_ann **+33.22** + - cost=0.1250 best τ=0.25 Sharpe_ann **+30.85** (Phase 1d.4 baseline: -4.0) + - cost=0.2500 best τ=0.25 Sharpe_ann **+27.73** + - cost=0.5000 best τ=0.25 Sharpe_ann **+15.68** + +**Caveats (OOS gate pending):** +- Results are in-sample (same fxcache used at training and eval). +- 99.8% win rate at cost=0 with 95 trades/episode is suspiciously high; possible + alpha-cache leakage (alpha_logits read at inference time may include forward info).