From 88a6db6eaea304c2c1a399fa124f2c25b6065f5b Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 15 May 2026 01:51:20 +0200 Subject: [PATCH] =?UTF-8?q?feat(ml-alpha):=20Mamba2=20backward=20pass=20?= =?UTF-8?q?=E2=80=94=20analytical,=20GPU-pure=20(Phase=201d.1,=20session?= =?UTF-8?q?=203)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit End-to-end analytical backward through all six stages of the forward chain. No atomicAdd, no SPSA, no host roundtrip during gradient flow — each scratch buffer is per-channel-unique so concurrent writes don't collide; cross-channel reduction is a separate kernel. New API: - GpuLinear::backward_with_slices(dy, act, weight, cublas, stream) → LinearGrads{dw, db, dx} Mirror of forward_with_slices added to ml-core; no GpuVarStore lookup. - Mamba2BackwardGrads holds dw_in/db_in/dw_a/db_a/dw_b/db_b/dw_c/dw_out/db_out - Mamba2Block::backward(&cache, &d_logit) → Mamba2BackwardGrads Chain (reverse of forward): 6′. W_out backward (cuBLAS sgemm) → d_h_enriched, dw_out, db_out 5′. mamba2_alpha_scan_bwd kernel → d_a/d_b/d_w_c per-channel/sample + d_h_s2 (identity passthrough) ↳ mamba2_alpha_reduce_d_proj × 2 → d_a_proj, d_b_proj [N, K, state] ↳ mamba2_alpha_reduce_d_w_c → dw_c [hidden, state] 3′. W_b backward → d_x_from_b, dw_b, db_b 2′. W_a backward → d_x_from_a, dw_a, db_a ↳ d_x = d_x_from_a + d_x_from_b 1′. W_in backward → dw_in, db_in (d_input discarded) Memory cost per backward call (for Phase 1d.1 sizes N=64, sh2=32, K=16, state=8): ~1.1 MiB scratch — fits comfortably on RTX 3050. Tests (9 passing on real GPU): - Backward produces all 9 grads with correct shapes - All grads finite through the 6-stage chain - dw_in non-zero (gradient flows to the input projection, proving the full chain is wired — not silently zero-ing somewhere) - Backward rejects wrong d_logit shape - + all previously-passing forward / config / shape tests Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/src/mamba2_block.rs | 309 ++++++++++++++++++++- crates/ml-core/src/cuda_autograd/linear.rs | 75 +++++ 2 files changed, 380 insertions(+), 4 deletions(-) diff --git a/crates/ml-alpha/src/mamba2_block.rs b/crates/ml-alpha/src/mamba2_block.rs index b88210c8b..c966bbbc1 100644 --- a/crates/ml-alpha/src/mamba2_block.rs +++ b/crates/ml-alpha/src/mamba2_block.rs @@ -40,7 +40,7 @@ use anyhow::{anyhow, Result}; use cudarc::cublas::CudaBlas; use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; use ml_core::cuda_autograd::gpu_tensor::GpuTensor; -use ml_core::cuda_autograd::linear::OwnedGpuLinear; +use ml_core::cuda_autograd::linear::{LinearActivations, LinearGrads, OwnedGpuLinear}; /// Hard kernel limits. The forward and backward kernels use register/local /// arrays sized at compile-time: `float x[16]` (state register) and @@ -89,10 +89,31 @@ impl Mamba2BlockConfig { } } +/// All parameter gradients produced by [`Mamba2Block::backward`]. +/// +/// Tensor shapes match the corresponding forward parameters: +/// - `dw_in / db_in` : `[hidden_dim, in_dim] / [hidden_dim]` +/// - `dw_a / db_a` : `[state_dim, hidden_dim] / [state_dim]` +/// - `dw_b / db_b` : `[state_dim, hidden_dim] / [state_dim]` +/// - `dw_c` : `[hidden_dim, state_dim]` (stored as flat slice; +/// layout matches the `W_c` weight) +/// - `dw_out / db_out` : `[1, hidden_dim] / [1]` +pub struct Mamba2BackwardGrads { + pub dw_in: GpuTensor, + pub db_in: GpuTensor, + pub dw_a: GpuTensor, + pub db_a: GpuTensor, + pub dw_b: GpuTensor, + pub db_b: GpuTensor, + pub dw_c: CudaSlice, + pub dw_out: GpuTensor, + pub db_out: GpuTensor, +} + /// Saved activations + per-stage tensors that the backward pass needs. Returned -/// by [`Mamba2Block::forward_train`]; consumed by the backward path (lands in -/// session 3). For inference-only ([`Mamba2Block::forward`]) the cache is -/// computed but immediately dropped. +/// by [`Mamba2Block::forward_train`]; consumed by [`Mamba2Block::backward`]. +/// For inference-only ([`Mamba2Block::forward`]) the cache is computed but +/// immediately dropped. pub struct Mamba2ForwardCache { /// Input as reshaped 2-D `[N*K, in_dim]` — needed to backprop into W_in. pub input_2d: GpuTensor, @@ -324,6 +345,202 @@ impl Mamba2Block { self.forward_train(input).map(|(logit, _cache)| logit) } + // ── Backward pass ───────────────────────────────────────────────── + + /// GPU-native backward pass. `d_logit` is the upstream gradient + /// `[n_batch, 1]` flowing back from the loss; `cache` is the + /// `Mamba2ForwardCache` returned by [`forward_train`]. Returns all + /// nine parameter gradients in a [`Mamba2BackwardGrads`]. + /// + /// Chain (mirrors forward in reverse): + /// 6′. W_out backward : d_logit + h_enriched_act → d_h_enriched, dw_out, db_out + /// 5′. scan backward : (a_proj, b_proj, d_h_enriched, W_c) + /// → d_a_per_channel, d_b_per_channel, + /// d_w_c_per_sample, d_h_s2 (=d_h_enriched) + /// ↳ reduce d_a_per_channel across j → d_a_proj [N, K, state] + /// ↳ reduce d_b_per_channel across j → d_b_proj [N, K, state] + /// ↳ reduce d_w_c_per_sample across i → dw_c [sh2, state] + /// 3′. W_b backward : d_b_proj[N*K, state] + x_act → d_x_from_b, dw_b, db_b + /// 2′. W_a backward : d_a_proj[N*K, state] + x_act → d_x_from_a, dw_a, db_a + /// ↳ d_x = d_x_from_a + d_x_from_b + /// 1′. W_in backward : d_x[N*K, hidden] + input_act → (d_input dropped), dw_in, db_in + pub fn backward( + &self, + cache: &Mamba2ForwardCache, + d_logit: &GpuTensor, + ) -> Result { + let c = &self.config; + let n_batch = cache.h_enriched.shape()[0]; + let n_rows = n_batch * c.seq_len; + + if d_logit.shape() != [n_batch, 1] { + return Err(anyhow!( + "Mamba2Block::backward: d_logit shape {:?} != [{}, 1]", + d_logit.shape(), n_batch + )); + } + + // ── 6′. Output projection backward ──────────────────────────── + let h_enriched_act = LinearActivations { input: cache.h_enriched.clone() }; + let LinearGrads { dw: dw_out, db: db_out, dx: d_h_enriched } = self + .w_out + .inner + .backward_with_slices(d_logit, &h_enriched_act, &self.w_out.weight, + &self.cublas, &self.stream) + .map_err(|e| anyhow!("w_out backward: {e}"))?; + + // ── 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 ({} floats): {e}", per_chan_n))?; + 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}"))?; + + // ── Launch scan backward kernel ─────────────────────────────── + 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 across j ─────── + 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_proj (a) launch: {e}"))?; + } + unsafe { + 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_proj (b) launch: {e}"))?; + } + + // ── Reduce d_w_c_per_sample across i → dw_c[sh2, state] ─────── + let red_w_c_grid_y: u32 = red_grid_z; + let red_w_c_cfg = LaunchConfig { + grid_dim: (c.hidden_dim as u32, red_w_c_grid_y, 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_d_w_c launch: {e}"))?; + } + + // ── 3′. W_b backward (uses x_act = cache.x) ─────────────────── + 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}"))?; + + // d_x = d_x_from_a + d_x_from_b + 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}"))?; + + 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; @@ -367,6 +584,90 @@ mod tests { assert!(cfg.validate().is_err(), "seq_len > 32 must be rejected (backward x_hist limit)"); } + #[test] + fn test_mamba2_block_backward_returns_finite_grads_with_correct_shapes() { + let stream = match cuda_stream_or_skip() { + Some(s) => s, + None => return, + }; + let cfg = Mamba2BlockConfig { + in_dim: 8, hidden_dim: 8, state_dim: 4, seq_len: 8, + }; + let n_batch = 3; + let block = Mamba2Block::new(cfg.clone(), Arc::clone(&stream)).expect("init"); + + // Forward with cache. + let n_in = n_batch * cfg.seq_len * cfg.in_dim; + let h: Vec = (0..n_in).map(|i| ((i as f32) * 0.137).sin()).collect(); + let d = stream.clone_htod(&h).expect("htod"); + let input = GpuTensor::new(d, vec![n_batch, cfg.seq_len, cfg.in_dim]).expect("input"); + let (logit, cache) = block.forward_train(&input).expect("forward_train"); + + // Pretend d_logit is just sigmoid(logit) - 0.5 (a BCE-shaped scalar). + let logit_h = logit.to_host(&stream).expect("logit dtoh"); + let d_logit_h: Vec = logit_h.iter() + .map(|z| 1.0 / (1.0 + (-z).exp()) - 0.5) + .collect(); + let d_logit_dev = stream.clone_htod(&d_logit_h).expect("htod"); + let d_logit = GpuTensor::new(d_logit_dev, vec![n_batch, 1]).expect("d_logit tensor"); + + let grads = block.backward(&cache, &d_logit).expect("backward"); + + // Shape assertions match each weight's forward shape. + assert_eq!(grads.dw_in.shape(), &[cfg.hidden_dim, cfg.in_dim], "dw_in shape"); + assert_eq!(grads.db_in.shape(), &[cfg.hidden_dim], "db_in shape"); + assert_eq!(grads.dw_a.shape(), &[cfg.state_dim, cfg.hidden_dim], "dw_a shape"); + assert_eq!(grads.db_a.shape(), &[cfg.state_dim], "db_a shape"); + assert_eq!(grads.dw_b.shape(), &[cfg.state_dim, cfg.hidden_dim], "dw_b shape"); + assert_eq!(grads.db_b.shape(), &[cfg.state_dim], "db_b shape"); + assert_eq!(grads.dw_c.len(), cfg.hidden_dim * cfg.state_dim, "dw_c numel"); + assert_eq!(grads.dw_out.shape(), &[1, cfg.hidden_dim], "dw_out shape"); + assert_eq!(grads.db_out.shape(), &[1], "db_out shape"); + + // Finiteness for every gradient tensor. + for (name, g) in [ + ("dw_in", &grads.dw_in), ("db_in", &grads.db_in), + ("dw_a", &grads.dw_a), ("db_a", &grads.db_a), + ("dw_b", &grads.dw_b), ("db_b", &grads.db_b), + ("dw_out",&grads.dw_out),("db_out",&grads.db_out), + ] { + let h = g.to_host(&stream).expect("dtoh"); + for (i, &v) in h.iter().enumerate() { + assert!(v.is_finite(), "{} grad index {} non-finite: {}", name, i, v); + } + } + // dw_c is a raw CudaSlice (no GpuTensor wrapper); read directly. + let mut dw_c_host = vec![0.0_f32; grads.dw_c.len()]; + stream.memcpy_dtoh(&grads.dw_c, &mut dw_c_host).expect("dw_c dtoh"); + for (i, &v) in dw_c_host.iter().enumerate() { + assert!(v.is_finite(), "dw_c grad index {} non-finite: {}", i, v); + } + + // At least one element of dw_in should be non-zero (gradient + // actually flowed through the chain — not silently zero). + let dw_in_h = grads.dw_in.to_host(&stream).expect("dw_in dtoh"); + let any_nonzero = dw_in_h.iter().any(|&v| v.abs() > 1e-12); + assert!(any_nonzero, "dw_in entirely zero — gradient didn't propagate to W_in"); + } + + #[test] + fn test_mamba2_block_backward_rejects_wrong_d_logit_shape() { + let stream = match cuda_stream_or_skip() { + Some(s) => s, + None => return, + }; + let cfg = Mamba2BlockConfig { + in_dim: 8, hidden_dim: 8, state_dim: 4, seq_len: 8, + }; + let n_batch = 3; + let block = Mamba2Block::new(cfg.clone(), Arc::clone(&stream)).expect("init"); + let input = GpuTensor::zeros(&[n_batch, cfg.seq_len, cfg.in_dim], &stream).expect("input"); + let (_logit, cache) = block.forward_train(&input).expect("forward_train"); + // Wrong d_logit shape: [n_batch, 2] instead of [n_batch, 1]. + let bad = GpuTensor::zeros(&[n_batch, 2], &stream).expect("bad d_logit"); + assert!(block.backward(&cache, &bad).is_err()); + } + #[test] fn test_mamba2_block_forward_train_returns_cache() { let stream = match cuda_stream_or_skip() { diff --git a/crates/ml-core/src/cuda_autograd/linear.rs b/crates/ml-core/src/cuda_autograd/linear.rs index 1e5371964..492c7a5b2 100644 --- a/crates/ml-core/src/cuda_autograd/linear.rs +++ b/crates/ml-core/src/cuda_autograd/linear.rs @@ -310,6 +310,81 @@ impl GpuLinear { Ok((y, activations)) } + /// Backward pass with raw `CudaSlice` weights (no `GpuVarStore` lookup). + /// Mirror of [`Self::forward_with_slices`]; produces the same `LinearGrads` + /// as [`Self::backward`]. + /// + /// - `dL/dW = dY^T @ X` (cuBLAS sgemm) + /// - `dL/db = sum(dY, axis=0)` (reduce kernel) + /// - `dL/dX = dY @ W` (cuBLAS sgemm) + pub fn backward_with_slices( + &self, + dy: &GpuTensor, + activations: &LinearActivations, + weight: &CudaSlice, + cublas: &CudaBlas, + stream: &Arc, + ) -> Result { + let batch = dy.shape()[0]; + let in_dim = self.in_dim; + let out_dim = self.out_dim; + + // ── dW[out, in] = dY^T[out, B] @ X[B, in] ────────────────── + // col-major: dW_col[in, out] = X_col[in, B] @ dY_col[out, B]^T + // transA=N, transB=T, m=in, n=out, k=B + let mut dw = GpuTensor::zeros(&[out_dim, in_dim], stream)?; + let x_ptr = raw_ptr(&activations.input.data, stream); + let dy_ptr = raw_ptr(&dy.data, stream); + let dw_ptr = raw_ptr_mut(&mut dw.data, stream); + unsafe { + gemm_ex_f32( + cublas, + cublasOperation_t::CUBLAS_OP_N, + cublasOperation_t::CUBLAS_OP_T, + in_dim as i32, + out_dim as i32, + batch as i32, + x_ptr, + in_dim as i32, + dy_ptr, + out_dim as i32, + dw_ptr, + in_dim as i32, + "dW_with_slices", + )?; + } + + // ── db[out] = sum(dY[B, out], axis=0) ─────────────────────── + let db = reduce_sum_axis0(dy, batch, out_dim, stream)?; + + // ── dX[B, in] = dY[B, out] @ W[out, in] ──────────────────── + // col-major: dX_col[in, B] = W_col[in, out] @ dY_col[out, B] + // transA=N, transB=N, m=in, n=B, k=out + let mut dx = GpuTensor::zeros(&[batch, in_dim], stream)?; + let w_ptr = raw_ptr(weight, stream); + let dy_ptr2 = raw_ptr(&dy.data, stream); + let dx_ptr = raw_ptr_mut(&mut dx.data, stream); + unsafe { + gemm_ex_f32( + cublas, + cublasOperation_t::CUBLAS_OP_N, + cublasOperation_t::CUBLAS_OP_N, + in_dim as i32, + batch as i32, + out_dim as i32, + w_ptr, + in_dim as i32, + dy_ptr2, + out_dim as i32, + dx_ptr, + in_dim as i32, + "dX_with_slices", + )?; + } + + Ok(LinearGrads { dw, db, dx }) + } + /// Backward pass: given `dL/dY` (upstream gradient), compute weight, bias, /// and input gradients. ///