SMT E3a: expose Mamba's recurrent state as a single-step updater (equivalence-pinned)

MambaRecurrence reconstructs MambaBlock's selective-scan as an explicit
step(state, x_t) -> (state', y_t) recurrence over host f32 weights, with the
recurrent state (h, conv-ring) carried between steps. This is the substrate the
SMT memory updater is behaviorally cloned on (E3b) — no BPTT, one step at a time.

Dimensions are derived from the persisted tensor shapes, so mamba.rs is left
untouched (it is near the 1250-LOC cap). silu/softplus are byte-identical copies
of the forward's.

Equivalence pin (the plan's highest-risk item): stepping a window one token at a
time from a zero state reproduces the full-sequence forward EXACTLY —
max_abs_diff = 0.0 (bit-identical), on an active_block with wide Delta so the
scan genuinely drives the output. Plus a fresh-state determinism test.

clippy(-D warnings) + fmt clean on the new module and test.

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

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

@@ -0,0 +1,274 @@
+//! `MambaRecurrence` — the single-timestep recurrent view of a [`MambaBlock`]
+//! (Supervised Memory Training, Phase E3).
+//!
+//! The full-sequence [`MambaBlock::forward`] runs the selective scan inline over
+//! a whole window. To behaviorally-clone Mamba as a *memory updater*
+//! `f(state, x_t) → (state', y_t)` (the SMT core — no backprop-through-time), we
+//! need the recurrence exposed one step at a time. This module reconstructs
+//! exactly that recurrence from a block's persisted weights, with **bit-level
+//! fidelity** to the full forward: stepping a sequence through `step` and
+//! concatenating the per-step outputs reproduces `forward` (pinned by
+//! `tests/mamba_step_equivalence.rs`).
+//!
+//! The recurrent state is `(h, conv)`:
+//! - `h ∈ ℝ^{d_inner · d_state}` — the selective-scan SSM state `h[j·n + k]`.
+//! - `conv ∈ ℝ^{(d_conv − 1) · d_inner}` — the causal depthwise-conv ring of the
+//!   last `d_conv − 1` pre-activation inputs (zero-initialized = the forward's
+//!   left zero-padding).
+//!
+//! Dimensions are derived from the persisted tensor shapes, so no access to the
+//! block's private config is needed and `mamba.rs` is left untouched.
+
+use crate::Result;
+use crate::layers::mamba::MambaBlock;
+
+/// Numerically-stable SiLU `x·σ(x)` — identical to the one in `mamba.rs` so the
+/// stepwise path matches the forward bit-for-bit.
+#[inline]
+fn silu(x: f32) -> f32 {
+    let s = if x >= 0.0 {
+        1.0 / (1.0 + (-x).exp())
+    } else {
+        let e = x.exp();
+        e / (1.0 + e)
+    };
+    x * s
+}
+
+/// Numerically-stable softplus `ln(1 + eˣ)` — identical to the one in `mamba.rs`.
+#[inline]
+fn softplus(x: f32) -> f32 {
+    x.max(0.0) + (1.0 + (-x.abs()).exp()).ln()
+}
+
+/// The recurrent state carried between [`MambaRecurrence::step`] calls.
+#[derive(Clone, Debug)]
+pub struct MambaState {
+    /// Selective-scan SSM state, row-major `[d_inner, d_state]` as `h[j*n + k]`.
+    h: Vec<f32>,
+    /// Causal-conv ring of the last `d_conv − 1` pre-activation inputs,
+    /// row-major `[d_conv - 1, d_inner]` (index 0 = oldest).
+    conv: Vec<f32>,
+}
+
+impl MambaState {
+    /// Read-only view of the SSM state `h` (`[d_inner · d_state]`). This is the
+    /// "memory" the SMT cloning supervises.
+    #[must_use]
+    pub fn hidden(&self) -> &[f32] {
+        &self.h
+    }
+}
+
+/// A `MambaBlock` reconstituted as an explicit step-by-step recurrence over host
+/// `f32` weights.
+pub struct MambaRecurrence {
+    d_model: usize,
+    d_inner: usize,
+    d_state: usize,
+    dt_rank: usize,
+    d_conv: usize,
+    in_proj: Vec<f32>,  // [d_model, 2*d_inner]
+    conv_w: Vec<f32>,   // [d_inner, d_conv]
+    conv_b: Vec<f32>,   // [d_inner] (zeros if the block had no conv bias)
+    a_log: Vec<f32>,    // [d_inner, d_state]
+    x_proj: Vec<f32>,   // [d_inner, dt_rank + 2*d_state]
+    dt_proj: Vec<f32>,  // [dt_rank, d_inner]
+    dt_bias: Vec<f32>,  // [d_inner]
+    d_skip: Vec<f32>,   // [d_inner]
+    out_proj: Vec<f32>, // [d_inner, d_model]
+}
+
+impl MambaRecurrence {
+    /// Reconstruct the recurrence from a block's persisted weights.
+    ///
+    /// # Errors
+    /// Propagates any tensor `to_vec` failure.
+    pub fn from_block(block: &MambaBlock) -> Result<Self> {
+        let mut in_proj = Vec::new();
+        let mut conv_w = Vec::new();
+        let mut conv_b = None;
+        let mut a_log = Vec::new();
+        let mut x_proj = Vec::new();
+        let mut dt_proj = Vec::new();
+        let mut dt_bias = Vec::new();
+        let mut d_skip = Vec::new();
+        let mut out_proj = Vec::new();
+        // Shapes we recover dimensions from.
+        let mut in_proj_shape = Vec::new();
+        let mut conv_shape = Vec::new();
+        let mut a_shape = Vec::new();
+        let mut x_proj_cols = 0usize;
+
+        for (name, tensor) in block.persistence_tensors() {
+            let dims = tensor.shape().dims().to_vec();
+            let data = tensor.to_vec()?;
+            match name {
+                "in_proj" => {
+                    in_proj_shape = dims;
+                    in_proj = data;
+                }
+                "conv1d_weight" => {
+                    conv_shape = dims;
+                    conv_w = data;
+                }
+                "conv1d_bias" => conv_b = Some(data),
+                "A_log" => {
+                    a_shape = dims;
+                    a_log = data;
+                }
+                "x_proj" => {
+                    x_proj_cols = dims[1];
+                    x_proj = data;
+                }
+                "dt_proj" => dt_proj = data,
+                "dt_bias" => dt_bias = data,
+                "D" => d_skip = data,
+                "out_proj" => out_proj = data,
+                _ => {}
+            }
+        }
+
+        let d_model = in_proj_shape[0];
+        let d_inner = in_proj_shape[1] / 2;
+        let d_conv = conv_shape[2]; // [d_inner, 1, d_conv]
+        let d_state = a_shape[1];
+        let dt_rank = x_proj_cols - 2 * d_state;
+        let conv_b = conv_b.unwrap_or_else(|| vec![0.0f32; d_inner]);
+
+        Ok(Self {
+            d_model,
+            d_inner,
+            d_state,
+            dt_rank,
+            d_conv,
+            in_proj,
+            conv_w,
+            conv_b,
+            a_log,
+            x_proj,
+            dt_proj,
+            dt_bias,
+            d_skip,
+            out_proj,
+        })
+    }
+
+    /// Input/output token width (`d_model`).
+    #[must_use]
+    pub const fn d_model(&self) -> usize {
+        self.d_model
+    }
+
+    /// Size of the recurrent SSM state `h` (`d_inner · d_state`).
+    #[must_use]
+    pub const fn state_size(&self) -> usize {
+        self.d_inner * self.d_state
+    }
+
+    /// A fresh zero state (matches the forward's zero conv-padding + zero `h`).
+    #[must_use]
+    pub fn init_state(&self) -> MambaState {
+        MambaState {
+            h: vec![0.0f32; self.d_inner * self.d_state],
+            conv: vec![0.0f32; self.d_conv.saturating_sub(1) * self.d_inner],
+        }
+    }
+
+    /// Advance one timestep: consume `x_t` (`[d_model]`), mutate `state`, and
+    /// return the output `y_t` (`[d_model]`). Bit-identical to the corresponding
+    /// timestep of [`MambaBlock::forward`].
+    ///
+    /// # Panics
+    /// Panics if `x_t.len() != d_model`.
+    #[must_use]
+    pub fn step(&self, x_t: &[f32], state: &mut MambaState) -> Vec<f32> {
+        let d = self.d_inner;
+        let n = self.d_state;
+        let kc = self.d_conv;
+        let dbc = self.dt_rank + 2 * n;
+        assert_eq!(x_t.len(), self.d_model, "x_t must be [d_model]");
+
+        // (1) in_proj → x_in (scan branch) + z (gate branch).
+        let mut x_in = vec![0.0f32; d];
+        let mut z = vec![0.0f32; d];
+        for (j, (xij, zj)) in x_in.iter_mut().zip(z.iter_mut()).enumerate() {
+            let mut sx = 0.0f32;
+            let mut sz = 0.0f32;
+            for (m, &xm) in x_t.iter().enumerate() {
+                sx += xm * self.in_proj[m * (2 * d) + j];
+                sz += xm * self.in_proj[m * (2 * d) + d + j];
+            }
+            *xij = sx;
+            *zj = sz;
+        }
+
+        // (2-3) causal depthwise conv over the last kc pre-activation inputs,
+        // then SiLU. window[kk] = conv-ring for kk < kc-1, else the current x_in.
+        let mut u = vec![0.0f32; d];
+        for j in 0..d {
+            let mut acc = self.conv_b[j];
+            for kk in 0..kc {
+                let v = if kk + 1 < kc {
+                    state.conv[kk * d + j]
+                } else {
+                    x_in[j]
+                };
+                acc += self.conv_w[j * kc + kk] * v;
+            }
+            u[j] = silu(acc);
+        }
+        // Slide the conv ring: drop the oldest, append the current x_in.
+        if kc > 1 {
+            for k in 0..(kc - 2) {
+                for j in 0..d {
+                    state.conv[k * d + j] = state.conv[(k + 1) * d + j];
+                }
+            }
+            for j in 0..d {
+                state.conv[(kc - 2) * d + j] = x_in[j];
+            }
+        }
+
+        // (4) x_proj → (dt_logits, B, C).
+        let mut xdbl = vec![0.0f32; dbc];
+        for (q, xq) in xdbl.iter_mut().enumerate() {
+            let mut s = 0.0f32;
+            for (j, &uj) in u.iter().enumerate() {
+                s += uj * self.x_proj[j * dbc + q];
+            }
+            *xq = s;
+        }
+
+        // (5) delta = softplus(dt_bias + dt_logits @ dt_proj).
+        let mut delta = vec![0.0f32; d];
+        for (j, dj) in delta.iter_mut().enumerate() {
+            let mut s = self.dt_bias[j];
+            for r in 0..self.dt_rank {
+                s += xdbl[r] * self.dt_proj[r * d + j];
+            }
+            *dj = softplus(s);
+        }
+
+        // (6-9) selective scan + skip + gate + out_proj.
+        let mut out = vec![0.0f32; self.d_model];
+        for j in 0..d {
+            let dj = delta[j];
+            let uj = u[j];
+            let mut y = self.d_skip[j] * uj;
+            for k in 0..n {
+                let a = -self.a_log[j * n + k].exp();
+                let da = (dj * a).exp();
+                let dbu = dj * xdbl[self.dt_rank + k] * uj; // dj * B[k] * uj
+                let hv = da * state.h[j * n + k] + dbu;
+                state.h[j * n + k] = hv;
+                y += xdbl[self.dt_rank + n + k] * hv; // C[k] * h'
+            }
+            let yg = y * silu(z[j]);
+            for m in 0..self.d_model {
+                out[m] += yg * self.out_proj[j * self.d_model + m];
+            }
+        }
+        out
+    }
+}

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

@@ -55,6 +55,7 @@ pub mod mixture_of_experts;
 
 // State-space models (Mamba)
 pub mod mamba;
+pub mod mamba_step;
 pub mod metal_mamba;
 
 // SMT predictive-state teacher (the time-parallel oracle for Supervised

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

@@ -0,0 +1,126 @@
+//! D309 equivalence pin — the single-step `MambaRecurrence` is numerically
+//! identical to the full-sequence `MambaBlock::forward`.
+//!
+//! This is the precondition the SMT plan flags as the highest remaining risk:
+//! before a recurrent updater can be behaviorally cloned, its exposed
+//! `(state, x_t) → (state', y_t)` step MUST reproduce the forward exactly.
+//! We run a window through `forward`, then through `step` one timestep at a
+//! time from a zero state, and require the outputs to match to the f32 floor.
+//!
+//! The block uses the `active_block` trick (wide `Δ`) so the scan genuinely
+//! contributes — otherwise the near-identity default init would make the test
+//! pass trivially without exercising the recurrence.
+
+use std::collections::HashMap;
+
+use rtx_tensor::{Device, Tensor};
+use rtx_transformers::layers::mamba::{MambaBlock, MambaConfig};
+use rtx_transformers::layers::mamba_step::MambaRecurrence;
+
+const D_MODEL: usize = 8;
+const D_STATE: usize = 16;
+const D_CONV: usize = 4;
+const L: usize = 7;
+
+fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 {
+    a.iter()
+        .zip(b)
+        .map(|(x, y)| (x - y).abs())
+        .fold(0.0, f32::max)
+}
+
+/// A seeded block with `dt_bias` zeroed (`Δ = softplus(0) ≈ 0.69`) so the scan
+/// genuinely drives the output — mirrors `real_selective_scan.rs::active_block`.
+fn active_block(dev: &Device, seed: u64) -> MambaBlock {
+    let cfg = MambaConfig::new(D_MODEL, D_STATE, D_CONV);
+    let base = MambaBlock::new_seeded(cfg.clone(), dev, seed).expect("base");
+    let d = cfg.get_d_inner();
+    let mut map = HashMap::new();
+    for (name, t) in base.persistence_tensors() {
+        if name == "dt_bias" {
+            map.insert(
+                name.to_string(),
+                Tensor::from_vec(vec![0.0f32; d], &[d], dev).expect("dtb"),
+            );
+        } else {
+            map.insert(name.to_string(), t.clone());
+        }
+    }
+    MambaBlock::from_persistence_tensors(cfg, dev, map).expect("rebuild")
+}
+
+/// Deterministic input window `[1, L, D_MODEL]`.
+fn input(dev: &Device) -> (Tensor, Vec<f32>) {
+    let mut s = 0x57E9u64;
+    let data: Vec<f32> = (0..L * D_MODEL)
+        .map(|_| {
+            s = s
+                .wrapping_mul(6_364_136_223_846_793_005)
+                .wrapping_add(1_442_695_040_888_963_407);
+            ((s >> 40) as f32 / (1u64 << 24) as f32) * 2.0 - 1.0
+        })
+        .collect();
+    let t = Tensor::from_vec(data.clone(), &[1, L, D_MODEL], dev).expect("input");
+    (t, data)
+}
+
+#[test]
+fn step_matches_full_forward() {
+    let dev = Device::cpu();
+    let block = active_block(&dev, 0xD309);
+    let (x, flat) = input(&dev);
+
+    // Full-sequence reference.
+    let full = block
+        .forward(&x)
+        .expect("forward")
+        .output
+        .to_vec()
+        .expect("vec");
+    assert_eq!(full.len(), L * D_MODEL);
+
+    // Stepwise: feed one token at a time from a zero state.
+    let rec = MambaRecurrence::from_block(&block).expect("recurrence");
+    assert_eq!(rec.d_model(), D_MODEL);
+    let mut state = rec.init_state();
+    let mut stepwise = Vec::with_capacity(L * D_MODEL);
+    for li in 0..L {
+        let x_t = &flat[li * D_MODEL..(li + 1) * D_MODEL];
+        stepwise.extend(rec.step(x_t, &mut state));
+    }
+
+    let diff = max_abs_diff(&full, &stepwise);
+    eprintln!("D309 step-vs-forward max_abs_diff = {diff:.2e}");
+    assert!(
+        diff < 1e-4,
+        "stepwise recurrence must match the full forward (max_abs_diff {diff:.2e})"
+    );
+
+    // The state must be non-trivial (the scan actually ran).
+    assert_eq!(rec.state_size(), state.hidden().len());
+    assert!(
+        state.hidden().iter().any(|&v| v.abs() > 1e-6),
+        "SSM state should be populated after stepping"
+    );
+}
+
+#[test]
+fn fresh_state_is_zero_and_deterministic() {
+    let dev = Device::cpu();
+    let block = active_block(&dev, 7);
+    let rec = MambaRecurrence::from_block(&block).expect("recurrence");
+    let s = rec.init_state();
+    assert!(s.hidden().iter().all(|&v| v == 0.0), "init state h is zero");
+
+    // Same input from a fresh state → identical output (stateless determinism).
+    let (_, flat) = input(&dev);
+    let run = |rec: &MambaRecurrence| {
+        let mut st = rec.init_state();
+        let mut out = Vec::new();
+        for li in 0..L {
+            out.extend(rec.step(&flat[li * D_MODEL..(li + 1) * D_MODEL], &mut st));
+        }
+        out
+    };
+    assert_eq!(run(&rec), run(&rec), "stepping must be deterministic");
+}