Mamba M1: real selective-scan forward (replaces the passthrough stub)

The MambaBlock forward was a stub — SelectiveScan::forward passed input
through, discretize returned zeros, conv1d was a no-op, and B/C were
randn per call, leaving conv1d_weight/dt_proj/A_log as dead weights with
zero temporal mixing. This implements a genuine S6 selective-scan:

- New params: x_proj [d_inner, dt_rank+2*d_state] (data-dependent dt,B,C),
  dt_bias [d_inner], D [d_inner] (skip). Added to new()/new_seeded() and
  the persistence contract (persistence_tensors/from_persistence_tensors).
- Real forward (CPU f32, looped — backbone is small): in_proj -> causal
  depthwise conv1d -> SiLU -> x_proj->(dt,B,C) -> delta=softplus(dt.dt_proj
  +dt_bias) -> A=-exp(A_log) -> sequential scan h=dA.h+dBu, y=sum C.h + D.u
  -> gate by SiLU(z) -> out_proj. Residual moved OUT (canonical).
  Numerically-stable silu_f32/softplus_f32 helpers.
- The scan runs inline (not via the immature rtx-tensor autograd tape);
  the analytic backward lands in M2 per docs/phase3_real_ssm_spec.md.

New tests/real_selective_scan.rs (4 cases, all green): liveness (each
formerly-dead weight now moves the output), causality (no future
leakage), seeded determinism, and finite/non-constant output.

Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>

crates/training/rtx-transformers/src/layers/mamba.rs

@@ -264,11 +264,23 @@ pub struct MambaBlock {
     conv1d_bias: Option<Tensor>,
     /// State space parameter A
     A_log: Tensor,
+    /// Data-dependent projection of the activated branch into
+    /// `(dt, B, C)` — `[d_inner, dt_rank + 2*d_state]`. This is what
+    /// makes the scan *selective*: `B` and `C` are functions of the
+    /// input rather than the fabricated `randn` of the old stub.
+    x_proj: Tensor,
     /// Time step projection
     dt_proj: Tensor,
+    /// Time-step bias added before softplus — `[d_inner]`.
+    dt_bias: Tensor,
+    /// Per-channel skip connection `D` — `[d_inner]`.
+    d_skip: Tensor,
     /// Output projection
     out_proj: Tensor,
-    /// Selective scan operation
+    /// Selective scan operation. Retained for API/cache compatibility
+    /// (`SelectiveScan` is public and exercised by tests); the
+    /// `MambaBlock` forward now runs the scan inline in f32.
+    #[allow(dead_code)]
     selective_scan: SelectiveScan,
 }
 
@@ -287,7 +299,10 @@ impl MambaBlock {
             None
         };
         let A_log = Tensor::randn(&[d_inner, config.d_state], device)?;
+        let x_proj = Tensor::randn(&[d_inner, dt_rank + 2 * config.d_state], device)?;
         let dt_proj = Tensor::randn(&[dt_rank, d_inner], device)?;
+        let dt_bias = Tensor::randn(&[d_inner], device)?;
+        let d_skip = Tensor::randn(&[d_inner], device)?;
         let out_proj = Tensor::randn(&[d_inner, config.d_model], device)?;
 
         let selective_scan = SelectiveScan::new(d_inner, config.d_state);
@@ -299,7 +314,10 @@ impl MambaBlock {
             conv1d_weight,
             conv1d_bias,
             A_log,
+            x_proj,
             dt_proj,
+            dt_bias,
+            d_skip,
             out_proj,
             selective_scan,
         })
@@ -341,7 +359,11 @@ impl MambaBlock {
             None
         };
         let A_log = Tensor::randn_seeded(&[d_inner, config.d_state], device, next_seed())?;
+        let x_proj =
+            Tensor::randn_seeded(&[d_inner, dt_rank + 2 * config.d_state], device, next_seed())?;
         let dt_proj = Tensor::randn_seeded(&[dt_rank, d_inner], device, next_seed())?;
+        let dt_bias = Tensor::randn_seeded(&[d_inner], device, next_seed())?;
+        let d_skip = Tensor::randn_seeded(&[d_inner], device, next_seed())?;
         let out_proj = Tensor::randn_seeded(&[d_inner, config.d_model], device, next_seed())?;
 
         let selective_scan = SelectiveScan::new(d_inner, config.d_state);
@@ -353,7 +375,10 @@ impl MambaBlock {
             conv1d_weight,
             conv1d_bias,
             A_log,
+            x_proj,
             dt_proj,
+            dt_bias,
+            d_skip,
             out_proj,
             selective_scan,
         })
@@ -371,7 +396,10 @@ impl MambaBlock {
     /// - `conv1d_weight`
     /// - `conv1d_bias` (only present when `config.conv_bias`)
     /// - `A_log`
+    /// - `x_proj`
     /// - `dt_proj`
+    /// - `dt_bias`
+    /// - `D`
     /// - `out_proj`
     pub fn persistence_tensors(&self) -> Vec<(&'static str, &Tensor)> {
         let mut out: Vec<(&'static str, &Tensor)> = vec![
@@ -382,7 +410,10 @@ impl MambaBlock {
             out.push(("conv1d_bias", bias));
         }
         out.push(("A_log", &self.A_log));
+        out.push(("x_proj", &self.x_proj));
         out.push(("dt_proj", &self.dt_proj));
+        out.push(("dt_bias", &self.dt_bias));
+        out.push(("D", &self.d_skip));
         out.push(("out_proj", &self.out_proj));
         out
     }
@@ -439,9 +470,18 @@ impl MambaBlock {
         let A_log = take(&mut tensors, "A_log")?;
         assert_shape(&A_log, "A_log", &[d_inner, config.d_state])?;
 
+        let x_proj = take(&mut tensors, "x_proj")?;
+        assert_shape(&x_proj, "x_proj", &[d_inner, dt_rank + 2 * config.d_state])?;
+
         let dt_proj = take(&mut tensors, "dt_proj")?;
         assert_shape(&dt_proj, "dt_proj", &[dt_rank, d_inner])?;
 
+        let dt_bias = take(&mut tensors, "dt_bias")?;
+        assert_shape(&dt_bias, "dt_bias", &[d_inner])?;
+
+        let d_skip = take(&mut tensors, "D")?;
+        assert_shape(&d_skip, "D", &[d_inner])?;
+
         let out_proj = take(&mut tensors, "out_proj")?;
         assert_shape(&out_proj, "out_proj", &[d_inner, config.d_model])?;
 
@@ -454,7 +494,10 @@ impl MambaBlock {
             conv1d_weight,
             conv1d_bias,
             A_log,
+            x_proj,
             dt_proj,
+            dt_bias,
+            d_skip,
             out_proj,
             selective_scan,
         })
@@ -462,74 +505,143 @@ impl MambaBlock {
 
     /// Forward pass through `MambaBlock`
     pub fn forward(&self, x: &Tensor) -> Result<MambaOutput> {
-        let batch_size = x.shape().dims()[0];
+        let dims = x.shape().dims().to_vec();
-        let seq_len = x.shape().dims()[1];
+        let (b, l, d_model) = (dims[0], dims[1], dims[2]);
-        let d_model = x.shape().dims()[2];
+        let d = self.config.get_d_inner();
+        let n = self.config.d_state;
+        let dt_rank = self.config.get_dt_rank();
+        let kc = self.config.d_conv;
+        let dbc = dt_rank + 2 * n;
 
-        // Input projection
+        // Pull every parameter + the input to CPU f32 once. The
-        let projected = x.matmul(&self.in_proj)?; // [B, L, 2*D]
+        // backbone is small (d_inner = expand·d_model), so a direct
-        let (x_proj, res_proj) = self.split_projection(&projected)?;
+        // loop nest is both correct and fast, and keeps the forward
+        // analytically differentiable by hand (see the Phase-3 spec:
+        // we do not route the selective scan through the autograd
+        // tape).
+        let xv = x.to_vec()?; // [b, l, d_model]
+        let in_proj = self.in_proj.to_vec()?; // [d_model, 2d]
+        let conv_w = self.conv1d_weight.to_vec()?; // [d, 1, kc] = [d, kc]
+        let conv_b = match &self.conv1d_bias {
+            Some(t) => Some(t.to_vec()?),
+            None => None,
+        };
+        let a_log = self.A_log.to_vec()?; // [d, n]
+        let x_proj = self.x_proj.to_vec()?; // [d, dbc]
+        let dt_proj = self.dt_proj.to_vec()?; // [dt_rank, d]
+        let dt_bias = self.dt_bias.to_vec()?; // [d]
+        let d_skip = self.d_skip.to_vec()?; // [d]
+        let out_proj = self.out_proj.to_vec()?; // [d, d_model]
 
-        // Convolution (simplified 1D conv)
+        let mut out = vec![0.0f32; b * l * d_model];
-        let x_conv = self.apply_conv1d(&x_proj)?;
 
-        // Activation (SiLU)
+        for bi in 0..b {
-        let x_activated = self.silu(&x_conv)?;
+            // (1) in_proj → x_in (selective branch) + z (gate branch).
-
+            let mut x_in = vec![0.0f32; l * d];
-        // Prepare state space parameters
+            let mut z = vec![0.0f32; l * d];
-        let A = self.A_log.exp()?; // Convert log to actual values
+            for li in 0..l {
-        let delta = Tensor::randn(
+                let xr = &xv[(bi * l + li) * d_model..(bi * l + li) * d_model + d_model];
-            &[batch_size, seq_len, self.config.get_d_inner()],
+                for j in 0..d {
-            &self.device,
+                    let (mut sx, mut sz) = (0.0f32, 0.0f32);
-        )?;
+                    for m in 0..d_model {
-        let B = Tensor::randn(&[batch_size, seq_len, self.config.d_state], &self.device)?;
+                        let w = xr[m];
-        let C = Tensor::randn(&[batch_size, seq_len, self.config.d_state], &self.device)?;
+                        sx += w * in_proj[m * (2 * d) + j];
-
+                        sz += w * in_proj[m * (2 * d) + d + j];
-        // Selective scan operation
+                    }
-        let ssm_output = self
+                    x_in[li * d + j] = sx;
-            .selective_scan
+                    z[li * d + j] = sz;
-            .forward(&x_activated, &delta, &A, &B, &C, None)?;
+                }
-
+            }
-        // Gating with residual projection
+            // (2-3) causal depthwise conv1d (left-pad kc-1) → SiLU → u.
-        let gated = ssm_output.mul(&self.silu(&res_proj)?)?;
+            let mut u = vec![0.0f32; l * d];
-
+            for li in 0..l {
-        // Output projection
+                for j in 0..d {
-        let output = gated.matmul(&self.out_proj)?;
+                    let mut acc = conv_b.as_ref().map_or(0.0f32, |cb| cb[j]);
-
+                    for kk in 0..kc {
-        // Residual connection
+                        let src = li as isize - (kc as isize - 1) + kk as isize;
-        let final_output = output.add(x)?;
+                        if src >= 0 {
+                            acc += x_in[(src as usize) * d + j] * conv_w[j * kc + kk];
+                        }
+                    }
+                    u[li * d + j] = silu_f32(acc);
+                }
+            }
+            // (4-5) x_proj → (dt, B, C); delta = softplus(dt·dt_proj + dt_bias).
+            let mut bmat = vec![0.0f32; l * n];
+            let mut cmat = vec![0.0f32; l * n];
+            let mut delta = vec![0.0f32; l * d];
+            for li in 0..l {
+                let mut xdbl = vec![0.0f32; dbc];
+                for q in 0..dbc {
+                    let mut s = 0.0f32;
+                    for j in 0..d {
+                        s += u[li * d + j] * x_proj[j * dbc + q];
+                    }
+                    xdbl[q] = s;
+                }
+                for nn in 0..n {
+                    bmat[li * n + nn] = xdbl[dt_rank + nn];
+                    cmat[li * n + nn] = xdbl[dt_rank + n + nn];
+                }
+                for j in 0..d {
+                    let mut s = dt_bias[j];
+                    for r in 0..dt_rank {
+                        s += xdbl[r] * dt_proj[r * d + j];
+                    }
+                    delta[li * d + j] = softplus_f32(s);
+                }
+            }
+            // (6-9) selective scan with A = -exp(A_log); gate by SiLU(z);
+            // accumulate through out_proj. No residual here (the caller
+            // adds it — canonical Mamba placement).
+            let mut h = vec![0.0f32; d * n];
+            for li in 0..l {
+                for j in 0..d {
+                    let dj = delta[li * d + j];
+                    let uj = u[li * d + j];
+                    let mut y = d_skip[j] * uj;
+                    for nn in 0..n {
+                        let a = -a_log[j * n + nn].exp();
+                        let da = (dj * a).exp();
+                        let dbu = dj * bmat[li * n + nn] * uj;
+                        let hv = da * h[j * n + nn] + dbu;
+                        h[j * n + nn] = hv;
+                        y += cmat[li * n + nn] * hv;
+                    }
+                    let yg = y * silu_f32(z[li * d + j]);
+                    let orow = &out_proj[j * d_model..j * d_model + d_model];
+                    let obase = (bi * l + li) * d_model;
+                    for m in 0..d_model {
+                        out[obase + m] += yg * orow[m];
+                    }
+                }
+            }
+        }
 
+        let output = Tensor::from_vec(out, &[b, l, d_model], &self.device)?;
         Ok(MambaOutput {
-            output: final_output,
+            output,
             aux_info: None,
         })
     }
+}
 
-    /// Split projection into main and residual parts
+/// Numerically-stable SiLU (a.k.a. swish): `x · σ(x)`.
-    fn split_projection(&self, projected: &Tensor) -> Result<(Tensor, Tensor)> {
+#[inline]
-        let d_inner = self.config.get_d_inner();
+fn silu_f32(x: f32) -> f32 {
-        let x_proj = projected.narrow(2, 0, d_inner)?;
+    let s = if x >= 0.0 {
-        let res_proj = projected.narrow(2, d_inner, d_inner)?;
+        1.0 / (1.0 + (-x).exp())
-        Ok((x_proj, res_proj))
+    } else {
-    }
+        let e = x.exp();
+        e / (1.0 + e)
+    };
+    x * s
+}
 
-    /// Apply 1D convolution (simplified)
+/// Numerically-stable softplus: `ln(1 + eˣ) = max(x,0) + ln(1 + e^-|x|)`.
-    fn apply_conv1d(&self, x: &Tensor) -> Result<Tensor> {
+#[inline]
-        // Simplified conv1d - in practice would use proper convolution
+fn softplus_f32(x: f32) -> f32 {
-        let output = x.clone();
+    x.max(0.0) + (1.0 + (-x.abs()).exp()).ln()
-        if let Some(bias) = &self.conv1d_bias {
-            Ok(output.add(bias)?)
-        } else {
-            Ok(output)
-        }
-    }
-
-    /// `SiLU` activation function
-    fn silu(&self, x: &Tensor) -> Result<Tensor> {
-        // SiLU(x) = x * sigmoid(x)
-        let sigmoid_x = x.sigmoid()?;
-        Ok(x.mul(&sigmoid_x)?)
-    }
 }
 
 impl Layer for MambaBlock {

crates/training/rtx-transformers/tests/real_selective_scan.rs

@@ -0,0 +1,124 @@
+//! Milestone-1 gate for the real selective-scan Mamba (Phase-3 spec).
+//!
+//! These tests prove the `MambaBlock` forward is now a genuine S6
+//! selective-state-space model — not the prior passthrough stub:
+//!   - **Liveness**: every weight that was *dead* in the stub
+//!     (`conv1d_weight`, `A_log`, `dt_proj`, `x_proj`, `D`) now changes
+//!     the output when perturbed.
+//!   - **Causality**: output at step `l` depends only on inputs `≤ l`
+//!     (no future leakage through the causal conv + scan).
+//!   - **Determinism**: a seeded block is bit-stable across builds.
+//!   - **Scan dynamics**: a non-trivial input produces non-constant,
+//!     finite output with real temporal mixing.
+
+use std::collections::HashMap;
+
+use rtx_tensor::{Device, Tensor};
+use rtx_transformers::layers::mamba::{MambaBlock, MambaConfig};
+
+const D_MODEL: usize = 8;
+const D_STATE: usize = 16;
+const D_CONV: usize = 4;
+const L: usize = 6;
+
+fn cfg() -> MambaConfig {
+    MambaConfig::new(D_MODEL, D_STATE, D_CONV)
+}
+
+fn fwd(block: &MambaBlock, x: &Tensor) -> Vec<f32> {
+    block.forward(x).expect("forward").output.to_vec().expect("to_vec")
+}
+
+fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 {
+    a.iter().zip(b).map(|(x, y)| (x - y).abs()).fold(0.0, f32::max)
+}
+
+/// Clone all persistence tensors into a map, replacing `target` with a
+/// copy whose every element is shifted by `delta`.
+fn perturbed_map(block: &MambaBlock, target: &str, delta: f32, dev: &Device) -> HashMap<String, Tensor> {
+    let mut map = HashMap::new();
+    for (name, t) in block.persistence_tensors() {
+        if name == target {
+            let mut v = t.to_vec().expect("to_vec");
+            for x in v.iter_mut() {
+                *x += delta;
+            }
+            let pt = Tensor::from_vec(v, t.shape().dims(), dev).expect("from_vec");
+            map.insert(name.to_string(), pt);
+        } else {
+            map.insert(name.to_string(), t.clone());
+        }
+    }
+    map
+}
+
+#[test]
+fn every_formerly_dead_weight_is_live() {
+    let dev = Device::cpu();
+    let block = MambaBlock::new_seeded(cfg(), &dev, 42).expect("block");
+    let x = Tensor::randn_seeded(&[1, L, D_MODEL], &dev, 7).expect("x");
+    let y0 = fwd(&block, &x);
+
+    // In the stub these were all discarded; each must now move the output.
+    for name in ["conv1d_weight", "A_log", "dt_proj", "x_proj", "D"] {
+        let map = perturbed_map(&block, name, 0.5, &dev);
+        let b2 = MambaBlock::from_persistence_tensors(cfg(), &dev, map).expect("rebuild");
+        let y2 = fwd(&b2, &x);
+        let diff = max_abs_diff(&y0, &y2);
+        assert!(
+            diff > 1e-4,
+            "weight `{name}` is DEAD — perturbing it left the output unchanged (max_abs_diff={diff:.2e})"
+        );
+    }
+}
+
+#[test]
+fn forward_is_causal() {
+    let dev = Device::cpu();
+    let block = MambaBlock::new_seeded(cfg(), &dev, 11).expect("block");
+    let base = Tensor::randn_seeded(&[1, L, D_MODEL], &dev, 3).expect("x");
+    let y0 = fwd(&block, &base);
+
+    // Perturb ONLY the last timestep.
+    let mut xv = base.to_vec().expect("to_vec");
+    for m in 0..D_MODEL {
+        xv[(L - 1) * D_MODEL + m] += 1.0;
+    }
+    let x2 = Tensor::from_vec(xv, &[1, L, D_MODEL], &dev).expect("x2");
+    let y2 = fwd(&block, &x2);
+
+    // Outputs for steps 0..L-1 must be identical (no future leakage);
+    // the last step must change.
+    let prefix0 = &y0[..(L - 1) * D_MODEL];
+    let prefix2 = &y2[..(L - 1) * D_MODEL];
+    assert!(
+        max_abs_diff(prefix0, prefix2) < 1e-6,
+        "non-causal: changing input[L-1] altered an earlier output"
+    );
+    let last0 = &y0[(L - 1) * D_MODEL..];
+    let last2 = &y2[(L - 1) * D_MODEL..];
+    assert!(
+        max_abs_diff(last0, last2) > 1e-4,
+        "last output did not respond to the last input"
+    );
+}
+
+#[test]
+fn seeded_forward_is_deterministic() {
+    let dev = Device::cpu();
+    let a = MambaBlock::new_seeded(cfg(), &dev, 99).expect("a");
+    let b = MambaBlock::new_seeded(cfg(), &dev, 99).expect("b");
+    let x = Tensor::randn_seeded(&[1, L, D_MODEL], &dev, 5).expect("x");
+    assert_eq!(fwd(&a, &x), fwd(&b, &x), "same (config, seed) must be bit-exact");
+}
+
+#[test]
+fn output_is_finite_and_non_constant() {
+    let dev = Device::cpu();
+    let block = MambaBlock::new_seeded(cfg(), &dev, 1).expect("block");
+    let x = Tensor::randn_seeded(&[1, L, D_MODEL], &dev, 2).expect("x");
+    let y = fwd(&block, &x);
+    assert!(y.iter().all(|v| v.is_finite()), "non-finite output");
+    let spread = max_abs_diff(&y, &vec![y[0]; y.len()]);
+    assert!(spread > 1e-4, "output is suspiciously constant ({spread:.2e})");
+}