D309: fix MambaRecurrence step — conv1d history buffer + new ergonomic API

The single-step recurrence was missing a causal conv1d history buffer, so
stepwise outputs diverged from the full-sequence forward (max_abs_diff ≈ 6e-2).
Add `conv_buf: Vec<f32>` to `MambaState` (oldest-first per channel), thread it
through `MambaRecurrence::step` so the kernel sees the correct `kc-1` prior
x_in values, and shift the buffer after each step.

Ergonomic additions:
- `MambaRecurrence::init_state()` — zero-initialised state with correct dims
- `MambaRecurrence::state_size()` / `d_model()` — accessor methods
- `MambaState::hidden()` — slice accessor for the SSM h vector
- `MambaState: PartialEq` — enables determinism assertions in tests

Both D309 tests now pass: `step_matches_full_forward` and
`fresh_state_is_zero_and_deterministic`.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

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

@@ -16,22 +16,32 @@ pub type ThinkError = TransformerError;
 
 /// Recurrent hidden state for a single-channel Mamba SSM.
 ///
-/// The vector `h` has length `d_inner * d_state`; it is laid out as
+/// `h` carries the SSM hidden state `[d_inner * d_state]`; `conv_buf` carries
-/// `h[j * d_state + n]` for channel `j` and state dimension `n`.
+/// the last `d_conv − 1` x_in values per channel `[d_inner * (d_conv − 1)]`
-#[derive(Debug, Clone)]
+/// so the causal conv1d can be computed correctly in a single-step setting.
+#[derive(Debug, Clone, PartialEq)]
 pub struct MambaState {
-    /// Flattened hidden state `[d_inner * d_state]`.
+    /// Flattened SSM hidden state `[d_inner * d_state]`.
     pub h: Vec<f32>,
+    /// Causal-conv history buffer `[d_inner * (d_conv - 1)]`, oldest-first per channel.
+    pub conv_buf: Vec<f32>,
 }
 
 impl MambaState {
     /// Create a zero-initialised state for the given dimensions.
     #[must_use]
-    pub fn zeros(d_inner: usize, d_state: usize) -> Self {
+    pub fn zeros(d_inner: usize, d_state: usize, d_conv: usize) -> Self {
         Self {
             h: vec![0.0f32; d_inner * d_state],
+            conv_buf: vec![0.0f32; d_inner * d_conv.saturating_sub(1)],
         }
     }
+
+    /// The SSM hidden-state slice `[d_inner * d_state]`.
+    #[must_use]
+    pub fn hidden(&self) -> &[f32] {
+        &self.h
+    }
 }
 
 // ─── recurrence cell ─────────────────────────────────────────────────────────
@@ -76,6 +86,24 @@ pub struct MambaRecurrence {
 }
 
 impl MambaRecurrence {
+    /// Return the model input dimension.
+    #[must_use]
+    pub fn d_model(&self) -> usize {
+        self.d_model
+    }
+
+    /// Return `d_inner * d_state` — the length of the hidden-state vector.
+    #[must_use]
+    pub fn state_size(&self) -> usize {
+        self.d_inner * self.d_state
+    }
+
+    /// Create a zero-initialised [`MambaState`] compatible with this recurrence cell.
+    #[must_use]
+    pub fn init_state(&self) -> MambaState {
+        MambaState::zeros(self.d_inner, self.d_state, self.d_conv)
+    }
+
     /// Extract all weights from `block` into CPU `Vec<f32>` buffers.
     ///
     /// # Errors
@@ -152,41 +180,35 @@ impl MambaRecurrence {
 
     /// Perform one recurrent step of the S6 selective scan.
     ///
-    /// `x` is the current input token `[d_model]`.  The conv1d context is
+    /// `x` is the current input token `[d_model]`.  `state` is updated in
-    /// carried inside `state` as the last `d_conv − 1` values of the inner
+    /// place; initialise with [`MambaRecurrence::init_state`] for a fresh
-    /// projection; for a fully stateless call initialize with
+    /// zero-state run.
-    /// [`MambaState::zeros`].
     ///
-    /// Returns `(output [d_model], new_state)`.
+    /// Returns the output token `[d_model]`.
     ///
-    /// # Errors
+    /// # Panics
-    /// Returns `TransformerError` if `x.len() != d_model` or `state.h.len()`
+    /// Panics if `x.len() != d_model` or `state.h.len() != d_inner * d_state`
-    /// doesn't match `d_inner * d_state`.
+    /// (programmer error — callers should always pass matching dimensions).
-    pub fn step(
+    pub fn step(&self, x: &[f32], state: &mut MambaState) -> Vec<f32> {
-        &self,
-        x: &[f32],
-        state: &MambaState,
-    ) -> Result<(Vec<f32>, MambaState), TransformerError> {
         let d = self.d_inner;
         let n = self.d_state;
         let dm = self.d_model;
         let kc = self.d_conv;
         let dt_rank = self.dt_rank;
 
-        if x.len() != dm {
+        assert_eq!(
-            return Err(TransformerError::Generic(format!(
+            x.len(),
-                "MambaRecurrence::step: x.len()={} != d_model={}",
+            dm,
-                x.len(),
+            "MambaRecurrence::step: x.len()={} != d_model={dm}",
-                dm
+            x.len()
-            )));
+        );
-        }
+        assert_eq!(
-        if state.h.len() != d * n {
+            state.h.len(),
-            return Err(TransformerError::Generic(format!(
+            d * n,
-                "MambaRecurrence::step: state.h.len()={} != d_inner*d_state={}",
+            "MambaRecurrence::step: state.h.len()={} != d_inner*d_state={}",
-                state.h.len(),
+            state.h.len(),
-                d * n
+            d * n
-            )));
+        );
-        }
 
         // ── 1. in_proj → x_in (SSM branch) + z (gate branch) ──────────────
         // in_proj: [d_model, 2*d_inner]
@@ -204,15 +226,29 @@ impl MambaRecurrence {
             z[j] = sz;
         }
 
-        // ── 2. causal depthwise conv1d (single-step: context = zeros) → SiLU → u ──
+        // ── 2. causal depthwise conv1d (with history buffer) → SiLU → u ──
-        // For a single step we don't have a history buffer, so the conv reduces to
+        // conv_buf holds the last (kc-1) x_in values per channel, oldest-first.
-        // a single tap (the last kernel coefficient) applied to the current x_in.
         // conv1d_weight layout: [d_inner, 1, d_conv] = [d_inner * d_conv] flat.
+        let cb = kc.saturating_sub(1); // history buffer length per channel
         let mut u = vec![0.0f32; d];
         for j in 0..d {
-            let acc = self.conv1d_bias[j] + x_in[j] * self.conv1d_weight[j * kc + (kc - 1)];
+            let mut acc = self.conv1d_bias[j];
+            for kk in 0..cb {
+                acc += self.conv1d_weight[j * kc + kk] * state.conv_buf[j * cb + kk];
+            }
+            acc += self.conv1d_weight[j * kc + (kc - 1)] * x_in[j];
             u[j] = silu_f32(acc);
         }
+        // Shift buffer left and append current x_in as newest entry.
+        for j in 0..d {
+            let base = j * cb;
+            for kk in 0..cb.saturating_sub(1) {
+                state.conv_buf[base + kk] = state.conv_buf[base + kk + 1];
+            }
+            if cb > 0 {
+                state.conv_buf[base + cb - 1] = x_in[j];
+            }
+        }
 
         // ── 3. x_proj → (dt_raw [dt_rank], B [d_state], C [d_state]) ───────
         let dbc = dt_rank + 2 * n;
@@ -255,6 +291,7 @@ impl MambaRecurrence {
             }
             y_raw[j] = yj;
         }
+        state.h = h_new;
 
         // ── 6. gate: y_gated = y_raw * silu(z) ──────────────────────────────
         let mut y_gated = vec![0.0f32; d];
@@ -272,7 +309,7 @@ impl MambaRecurrence {
             }
         }
 
-        Ok((output, MambaState { h: h_new }))
+        output
     }
 }