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]>
This commit is contained in:
omar sobh
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
96106cf888
commit
00cd527ed4
@@ -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,
|
||||
/// Output projection
|
||||
/// 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 seq_len = x.shape().dims()[1];
|
||||
let d_model = x.shape().dims()[2];
|
||||
let dims = x.shape().dims().to_vec();
|
||||
let (b, l, d_model) = (dims[0], dims[1], 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
|
||||
let projected = x.matmul(&self.in_proj)?; // [B, L, 2*D]
|
||||
let (x_proj, res_proj) = self.split_projection(&projected)?;
|
||||
// Pull every parameter + the input to CPU f32 once. The
|
||||
// backbone is small (d_inner = expand·d_model), so a direct
|
||||
// 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 x_conv = self.apply_conv1d(&x_proj)?;
|
||||
let mut out = vec![0.0f32; b * l * d_model];
|
||||
|
||||
// Activation (SiLU)
|
||||
let x_activated = self.silu(&x_conv)?;
|
||||
|
||||
// Prepare state space parameters
|
||||
let A = self.A_log.exp()?; // Convert log to actual values
|
||||
let delta = Tensor::randn(
|
||||
&[batch_size, seq_len, self.config.get_d_inner()],
|
||||
&self.device,
|
||||
)?;
|
||||
let B = Tensor::randn(&[batch_size, seq_len, self.config.d_state], &self.device)?;
|
||||
let C = Tensor::randn(&[batch_size, seq_len, self.config.d_state], &self.device)?;
|
||||
|
||||
// Selective scan operation
|
||||
let ssm_output = self
|
||||
.selective_scan
|
||||
.forward(&x_activated, &delta, &A, &B, &C, None)?;
|
||||
|
||||
// Gating with residual projection
|
||||
let gated = ssm_output.mul(&self.silu(&res_proj)?)?;
|
||||
|
||||
// Output projection
|
||||
let output = gated.matmul(&self.out_proj)?;
|
||||
|
||||
// Residual connection
|
||||
let final_output = output.add(x)?;
|
||||
for bi in 0..b {
|
||||
// (1) in_proj → x_in (selective branch) + z (gate branch).
|
||||
let mut x_in = vec![0.0f32; l * d];
|
||||
let mut z = vec![0.0f32; l * d];
|
||||
for li in 0..l {
|
||||
let xr = &xv[(bi * l + li) * d_model..(bi * l + li) * d_model + d_model];
|
||||
for j in 0..d {
|
||||
let (mut sx, mut sz) = (0.0f32, 0.0f32);
|
||||
for m in 0..d_model {
|
||||
let w = xr[m];
|
||||
sx += w * in_proj[m * (2 * d) + j];
|
||||
sz += w * in_proj[m * (2 * d) + d + j];
|
||||
}
|
||||
x_in[li * d + j] = sx;
|
||||
z[li * d + j] = sz;
|
||||
}
|
||||
}
|
||||
// (2-3) causal depthwise conv1d (left-pad kc-1) → SiLU → u.
|
||||
let mut u = vec![0.0f32; l * d];
|
||||
for li in 0..l {
|
||||
for j in 0..d {
|
||||
let mut acc = conv_b.as_ref().map_or(0.0f32, |cb| cb[j]);
|
||||
for kk in 0..kc {
|
||||
let src = li as isize - (kc as isize - 1) + kk as isize;
|
||||
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
|
||||
fn split_projection(&self, projected: &Tensor) -> Result<(Tensor, Tensor)> {
|
||||
let d_inner = self.config.get_d_inner();
|
||||
let x_proj = projected.narrow(2, 0, d_inner)?;
|
||||
let res_proj = projected.narrow(2, d_inner, d_inner)?;
|
||||
Ok((x_proj, res_proj))
|
||||
}
|
||||
/// Numerically-stable SiLU (a.k.a. swish): `x · σ(x)`.
|
||||
#[inline]
|
||||
fn silu_f32(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
|
||||
}
|
||||
|
||||
/// Apply 1D convolution (simplified)
|
||||
fn apply_conv1d(&self, x: &Tensor) -> Result<Tensor> {
|
||||
// Simplified conv1d - in practice would use proper convolution
|
||||
let output = x.clone();
|
||||
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)?)
|
||||
}
|
||||
/// Numerically-stable softplus: `ln(1 + eˣ) = max(x,0) + ln(1 + e^-|x|)`.
|
||||
#[inline]
|
||||
fn softplus_f32(x: f32) -> f32 {
|
||||
x.max(0.0) + (1.0 + (-x.abs()).exp()).ln()
|
||||
}
|
||||
|
||||
impl Layer for MambaBlock {
|
||||
|
||||
Reference in New Issue
Block a user