diff --git a/crates/ml-alpha/src/mamba2_block.rs b/crates/ml-alpha/src/mamba2_block.rs index e10882c0a..c6d142215 100644 --- a/crates/ml-alpha/src/mamba2_block.rs +++ b/crates/ml-alpha/src/mamba2_block.rs @@ -145,6 +145,50 @@ pub struct Mamba2ForwardCacheSeq { pub h_enriched_seq: GpuTensor, } +/// Pre-allocated scratch for [`Mamba2Block::backward_from_h_enriched_seq_into`]. +/// Holds the 4 BIG per-step scratch buffers (the ones at ~6MB each that +/// dominated the per-call `alloc_zeros` churn). Smaller buffers +/// (`d_a_proj_flat`, `d_b_proj_flat`, `dw_c`) remain per-call — +/// they're tiny and feed directly into ownership-transferred +/// LinearGrads outputs, where pre-allocation would force a +/// refactor of ml-core's cuBLAS wrappers without proportional gain. +/// +/// Construct ONCE per (n_batch, seq_len, hidden_dim, state_dim) tuple. +pub struct Mamba2BackwardScratch { + pub d_a_per_channel: CudaSlice, // [N, sh2, K, state_d] + pub d_b_per_channel: CudaSlice, // [N, sh2, K, state_d] + pub d_w_c_per_sample: CudaSlice, // [N, sh2, state_d] + pub d_h_s2: CudaSlice, // [N, sh2] + pub n_batch: usize, + pub seq_len: usize, + pub hidden_dim: usize, + pub state_dim: usize, +} + +impl Mamba2BackwardScratch { + pub fn new( + stream: &Arc, + n_batch: usize, + seq_len: usize, + hidden_dim: usize, + state_dim: usize, + ) -> Result { + let per_chan_n = n_batch * hidden_dim * seq_len * state_dim; + let per_sample_n = n_batch * hidden_dim * state_dim; + Ok(Self { + d_a_per_channel: stream.alloc_zeros::(per_chan_n) + .map_err(|e| anyhow!("scratch d_a_per_channel: {e}"))?, + d_b_per_channel: stream.alloc_zeros::(per_chan_n) + .map_err(|e| anyhow!("scratch d_b_per_channel: {e}"))?, + d_w_c_per_sample: stream.alloc_zeros::(per_sample_n) + .map_err(|e| anyhow!("scratch d_w_c_per_sample: {e}"))?, + d_h_s2: stream.alloc_zeros::(n_batch * hidden_dim) + .map_err(|e| anyhow!("scratch d_h_s2: {e}"))?, + n_batch, seq_len, hidden_dim, state_dim, + }) + } +} + /// GPU-resident Mamba2 sequence block. /// /// Owns all parameter tensors (`W_in`, `W_a`, `W_b`, `W_c`, `W_out`) on the @@ -984,6 +1028,156 @@ impl Mamba2Block { }) } + /// Perf-tuned backward variant — same math as + /// [`backward_from_h_enriched_seq`] but uses caller-provided + /// pre-allocated scratch buffers instead of allocating 7 fresh + /// CudaSlices per call. Cuts ~10-20ms/step on the training hot + /// path by eliminating `cudaMalloc` churn. + /// + /// Callers MUST pre-allocate [`Mamba2BackwardScratch`] with shapes + /// matching the (n_batch, seq_len, hidden_dim, state_dim) used + /// here. The scratch buffer contents are overwritten by every call + /// (no carry between calls; each backward pass is self-contained). + pub fn backward_from_h_enriched_seq_into( + &self, + cache: &Mamba2ForwardCacheSeq, + d_h_enriched_seq: &GpuTensor, + scratch: &mut Mamba2BackwardScratch, + ) -> Result { + let c = &self.config; + let n_batch = cache.h_enriched_seq.shape()[0]; + let n_rows = n_batch * c.seq_len; + + // Validate scratch shapes match — wrong scratch silently + // produces wrong gradients (size mismatch → out-of-bounds + // writes). Assert eagerly. + anyhow::ensure!( + scratch.n_batch == n_batch + && scratch.seq_len == c.seq_len + && scratch.hidden_dim == c.hidden_dim + && scratch.state_dim == c.state_dim, + "Mamba2BackwardScratch shape mismatch: expected ({},{},{},{}) got ({},{},{},{})", + n_batch, c.seq_len, c.hidden_dim, c.state_dim, + scratch.n_batch, scratch.seq_len, scratch.hidden_dim, scratch.state_dim + ); + anyhow::ensure!( + d_h_enriched_seq.shape() == [n_batch, c.seq_len, c.hidden_dim], + "d_h_enriched_seq shape {:?} != [{}, {}, {}]", + d_h_enriched_seq.shape(), n_batch, c.seq_len, c.hidden_dim + ); + + 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_seq) + .arg(cache.a_proj.cuda_data()) + .arg(cache.b_proj.cuda_data()) + .arg(d_h_enriched_seq.cuda_data()) + .arg(&self.w_c) + .arg(&mut scratch.d_a_per_channel) + .arg(&mut scratch.d_b_per_channel) + .arg(&mut scratch.d_w_c_per_sample) + .arg(&mut scratch.d_h_s2) + .arg(&n_i32).arg(&k_i32).arg(&sh2_i32).arg(&st_i32) + .launch(bwd_cfg) + .map_err(|e| anyhow!("scan_bwd_seq launch: {e}"))?; + } + + // Small per-call buffers (~few KB each, no measurable cost). + 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 into: {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 into: {e}"))?; + unsafe { + self.stream + .launch_builder(&self.kernel_reduce_d_proj) + .arg(&scratch.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 seq into: {e}"))?; + self.stream + .launch_builder(&self.kernel_reduce_d_proj) + .arg(&scratch.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 seq into: {e}"))?; + } + + 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 into: {e}"))?; + unsafe { + self.stream + .launch_builder(&self.kernel_reduce_d_w_c) + .arg(&scratch.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 seq into: {e}"))?; + } + + 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 seq into: {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 into: {e}"))?; + + 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 seq into: {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 into: {e}"))?; + + let d_x = d_x_from_a.add(&d_x_from_b, &self.stream) + .map_err(|e| anyhow!("sum d_x branches into: {e}"))?; + + 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 into: {e}"))?; + + let dw_out = GpuTensor::zeros(&[1, c.hidden_dim], &self.stream) + .map_err(|e| anyhow!("alloc zero dw_out seq into: {e}"))?; + let db_out = GpuTensor::zeros(&[1], &self.stream) + .map_err(|e| anyhow!("alloc zero db_out seq into: {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-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index d2f144c0b..11596cc81 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -47,7 +47,7 @@ use rand_chacha::ChaCha8Rng; use crate::cfc::snap_features::{Mbp10RawInput, ES_TICK_SIZE, FEATURE_DIM, REGIME_DIM}; use crate::heads::{HIDDEN_DIM, N_HORIZONS}; use crate::mamba2_block::{ - Mamba2AdamW, Mamba2AdamWConfig, Mamba2Block, Mamba2BlockConfig, + Mamba2AdamW, Mamba2AdamWConfig, Mamba2BackwardScratch, Mamba2Block, Mamba2BlockConfig, }; use crate::pinned_mem::MappedF32Buffer; use crate::trainer::optim::AdamW; @@ -128,6 +128,9 @@ pub struct PerceptionTrainer { // Mamba2 encoder block + its optimizer pub mamba2: Mamba2Block, pub mamba2_adamw: Mamba2AdamW, + /// Pre-allocated scratch for the Mamba2 seq backward — eliminates + /// ~10-20ms/step of `alloc_zeros` churn in the hot path. + mamba2_bwd_scratch: Mamba2BackwardScratch, // CfC + heads weights + their AdamWs (6 groups — tau is trained now). pub w_in_d: CudaSlice, @@ -242,6 +245,9 @@ impl PerceptionTrainer { }, ) .context("Mamba2AdamW::new")?; + let mamba2_bwd_scratch = Mamba2BackwardScratch::new( + &stream, cfg.n_batch, cfg.seq_len, HIDDEN_DIM, cfg.mamba2_state_dim, + ).context("Mamba2BackwardScratch::new")?; // CfC weights (input = h_enriched [HIDDEN_DIM], output = [HIDDEN_DIM]) let mut r = ChaCha8Rng::seed_from_u64(cfg.seed); @@ -334,6 +340,7 @@ impl PerceptionTrainer { transpose_3d_fn, mamba2, mamba2_adamw, + mamba2_bwd_scratch, opt_w_in, opt_w_rec, opt_b, @@ -755,8 +762,8 @@ impl PerceptionTrainer { // ── 8. Mamba2 backward — single call consumes grad_h_enriched_seq. let mamba2_grads = self .mamba2 - .backward_from_h_enriched_seq(&cache, &grad_h_enriched_seq) - .context("mamba2 backward_from_h_enriched_seq")?; + .backward_from_h_enriched_seq_into(&cache, &grad_h_enriched_seq, &mut self.mamba2_bwd_scratch) + .context("mamba2 backward_from_h_enriched_seq_into")?; // ── 9. Apply AdamW updates on all 7 param groups (added tau). self.opt_w_in.step(&mut self.w_in_d, &self.grad_w_in_d)?;