Restore Mage-VL mamba import compatibility

README.md

@@ -64,4 +64,6 @@ to reconstruct only the events supported by the recording, and links each claim
 - VibeASR.cpp pinned at `70b3ebb8ad75b5f37aee948df34f15cc84951d05`
 
 Mage-VL loading and codec preprocessing are derived from Microsoft's Apache-2.0 reference Space. ReplayForge's
-workflow, evidence model, privacy controls, report generator, and interface are original project work.
+pure-PyTorch `mamba_ssm` compatibility shim is also retained because Transformers validates Mage-VL's optional
+StreamMind import even though ReplayForge does not load that gate. ReplayForge's workflow, evidence model, privacy
+controls, report generator, and interface are original project work.

mamba_ssm/__init__.py

@@ -0,0 +1,26 @@
+"""Pure-PyTorch stand-in for the parts of `mamba-ssm` that Mage-VL needs.
+
+`microsoft/Mage-VL`'s remote code (`streammind_gate.py`) does
+
+    from mamba_ssm.models.mixer_seq_simple import create_block
+
+at module level, and `transformers.dynamic_module_utils.check_imports` imports
+every top-level dependency of every remote file before it will load the model —
+so `mamba_ssm` must be importable even though only the StreamMind cognition gate
+uses it.
+
+The real `mamba-ssm` ships CUDA extensions (`selective_scan_cuda`,
+`causal_conv1d_cuda`) with no wheel for the ZeroGPU Blackwell (sm_120) /
+torch-2.11 / cp312 runtime, and compiling them from source is not viable inside
+a Space build. This package therefore provides a faithful pure-PyTorch port of
+Mamba-1 (`mamba_ssm.modules.mamba_simple.Mamba` + `mamba_ssm.modules.block.Block`
+with the defaults `create_block` uses: rms_norm=False, fused_add_norm=False,
+residual_in_fp32=False, d_intermediate=0) built on the upstream
+`selective_scan_ref` reference recurrence, with identical parameter names and
+shapes so `streammind_gate.safetensors` loads with `strict=True`.
+
+The gate runs over a handful of EPFE tokens (one per codec canvas), so the slow
+sequential scan costs milliseconds — the CUDA kernel buys nothing here.
+"""
+
+__version__ = "2.2.6.mage-vl-pure-torch"

mamba_ssm/models/mixer_seq_simple.py

@@ -0,0 +1,173 @@
+"""Pure-PyTorch `create_block` — Mamba-1 mixer + residual block, no CUDA kernels.
+
+Mirrors `mamba_ssm.modules.mamba_simple.Mamba` and `mamba_ssm.modules.block.Block`
+for the argument set `create_block()` is called with by Mage-VL's StreamMind gate
+(`create_block(d_model, d_intermediate=0, layer_idx=i)`), i.e. the upstream
+defaults: ssm_cfg={} -> d_state=16, d_conv=4, expand=2, dt_rank=ceil(d_model/16),
+rms_norm=False -> nn.LayerNorm, fused_add_norm=False, residual_in_fp32=False.
+
+Parameter names/shapes match upstream exactly, so the released
+`streammind_gate.safetensors` loads with strict=True:
+
+    mixer.in_proj.weight    (2*d_inner, d_model)
+    mixer.conv1d.{weight,bias}  (d_inner, 1, d_conv) / (d_inner,)
+    mixer.x_proj.weight     (dt_rank + 2*d_state, d_inner)
+    mixer.dt_proj.{weight,bias}  (d_inner, dt_rank) / (d_inner,)
+    mixer.A_log             (d_inner, d_state)
+    mixer.D                 (d_inner,)
+    mixer.out_proj.weight   (d_model, d_inner)
+    norm.{weight,bias}      (d_model,)
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def selective_scan_ref(u, delta, A, B, C, D=None, z=None, delta_bias=None,
+                       delta_softplus=False, return_last_state=False):
+    """Upstream reference implementation (single-group, non-complex path).
+
+    u, delta, z: (b, d, l) · A: (d, n) · B, C: (b, n, l) · D: (d,)
+    """
+    dtype_in = u.dtype
+    u = u.float()
+    delta = delta.float()
+    if delta_bias is not None:
+        delta = delta + delta_bias[..., None].float()
+    if delta_softplus:
+        delta = F.softplus(delta)
+    batch, dim = u.shape[0], A.shape[0]
+    dstate = A.shape[1]
+    B = B.float()
+    C = C.float()
+    x = A.new_zeros((batch, dim, dstate))
+    deltaA = torch.exp(torch.einsum("bdl,dn->bdln", delta, A))
+    deltaB_u = torch.einsum("bdl,bnl,bdl->bdln", delta, B, u)
+    ys = []
+    last_state = None
+    for i in range(u.shape[2]):
+        x = deltaA[:, :, i] * x + deltaB_u[:, :, i]
+        ys.append(torch.einsum("bdn,bn->bd", x, C[:, :, i]))
+        if i == u.shape[2] - 1:
+            last_state = x
+    y = torch.stack(ys, dim=2)
+    out = y if D is None else y + u * D.unsqueeze(-1)
+    if z is not None:
+        out = out * F.silu(z)
+    out = out.to(dtype=dtype_in)
+    return (out, last_state) if return_last_state else out
+
+
+class Mamba(nn.Module):
+    def __init__(self, d_model, d_state=16, d_conv=4, expand=2, dt_rank="auto",
+                 conv_bias=True, bias=False, layer_idx=None, **kwargs):
+        super().__init__()
+        self.d_model = d_model
+        self.d_state = d_state
+        self.d_conv = d_conv
+        self.expand = expand
+        self.d_inner = int(expand * d_model)
+        self.dt_rank = math.ceil(d_model / 16) if dt_rank == "auto" else dt_rank
+        self.layer_idx = layer_idx
+
+        self.in_proj = nn.Linear(self.d_model, self.d_inner * 2, bias=bias)
+        self.conv1d = nn.Conv1d(
+            in_channels=self.d_inner, out_channels=self.d_inner, bias=conv_bias,
+            kernel_size=d_conv, groups=self.d_inner, padding=d_conv - 1,
+        )
+        self.activation = "silu"
+        self.act = nn.SiLU()
+        self.x_proj = nn.Linear(self.d_inner, self.dt_rank + self.d_state * 2, bias=False)
+        self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True)
+        self.A_log = nn.Parameter(torch.zeros(self.d_inner, self.d_state))
+        self.D = nn.Parameter(torch.ones(self.d_inner))
+        self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=bias)
+
+    def forward(self, hidden_states, inference_params=None, **kwargs):
+        batch, seqlen, _ = hidden_states.shape
+        xz = self.in_proj(hidden_states).transpose(1, 2)          # (b, 2*d_inner, l)
+        A = -torch.exp(self.A_log.float())                        # (d_inner, d_state)
+        x, z = xz.chunk(2, dim=1)
+        x = self.act(self.conv1d(x)[..., :seqlen])
+        x_dbl = self.x_proj(x.transpose(1, 2).reshape(batch * seqlen, self.d_inner))
+        dt, B, C = torch.split(
+            x_dbl, [self.dt_rank, self.d_state, self.d_state], dim=-1
+        )
+        dt = (self.dt_proj.weight @ dt.t()).view(self.d_inner, batch, seqlen)
+        dt = dt.permute(1, 0, 2).contiguous()                     # (b, d_inner, l)
+        B = B.view(batch, seqlen, self.d_state).transpose(1, 2).contiguous()
+        C = C.view(batch, seqlen, self.d_state).transpose(1, 2).contiguous()
+        y = selective_scan_ref(
+            x, dt, A, B, C, self.D.float(), z=z,
+            delta_bias=self.dt_proj.bias.float(), delta_softplus=True,
+        )
+        return self.out_proj(y.transpose(1, 2))
+
+
+class Block(nn.Module):
+    """`mamba_ssm.modules.block.Block` with fused_add_norm=False, mlp=Identity."""
+
+    def __init__(self, dim, mixer_cls, norm_cls=nn.LayerNorm, mlp_cls=nn.Identity,
+                 fused_add_norm=False, residual_in_fp32=False):
+        super().__init__()
+        self.residual_in_fp32 = residual_in_fp32
+        self.fused_add_norm = fused_add_norm
+        self.norm = norm_cls(dim)
+        self.mixer = mixer_cls(dim)
+        if mlp_cls is not nn.Identity:
+            self.norm2 = norm_cls(dim)
+            self.mlp = mlp_cls(dim)
+        else:
+            self.mlp = None
+
+    def forward(self, hidden_states, residual=None, inference_params=None, **kwargs):
+        residual = (hidden_states + residual) if residual is not None else hidden_states
+        hidden_states = self.norm(residual.to(dtype=self.norm.weight.dtype))
+        if self.residual_in_fp32:
+            residual = residual.to(torch.float32)
+        hidden_states = self.mixer(hidden_states, inference_params=inference_params)
+        if self.mlp is not None:
+            residual = hidden_states + residual
+            hidden_states = self.norm2(residual.to(dtype=self.norm2.weight.dtype))
+            hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+
+def create_block(d_model, d_intermediate=0, ssm_cfg=None, attn_layer_idx=None,
+                 attn_cfg=None, norm_epsilon=1e-5, rms_norm=False,
+                 residual_in_fp32=False, fused_add_norm=False, layer_idx=None,
+                 device=None, dtype=None, **kwargs):
+    if d_intermediate:
+        raise NotImplementedError(
+            "This pure-PyTorch mamba_ssm stand-in only supports d_intermediate=0 "
+            "(the configuration used by Mage-VL's StreamMind gate)."
+        )
+    if attn_layer_idx and layer_idx in attn_layer_idx:
+        raise NotImplementedError(
+            "Attention blocks are not supported by this mamba_ssm stand-in."
+        )
+    if rms_norm:
+        raise NotImplementedError(
+            "rms_norm=True is not supported by this mamba_ssm stand-in."
+        )
+    factory_kwargs = {"device": device, "dtype": dtype}
+    ssm_cfg = dict(ssm_cfg or {})
+    ssm_cfg.pop("layer", None)
+
+    def mixer_cls(dim):
+        return Mamba(dim, layer_idx=layer_idx, **ssm_cfg)
+
+    def norm_cls(dim):
+        return nn.LayerNorm(dim, eps=norm_epsilon)
+
+    block = Block(
+        d_model, mixer_cls, norm_cls=norm_cls, mlp_cls=nn.Identity,
+        fused_add_norm=fused_add_norm, residual_in_fp32=residual_in_fp32,
+    )
+    block.layer_idx = layer_idx
+    if factory_kwargs["device"] is not None or factory_kwargs["dtype"] is not None:
+        block = block.to(**{k: v for k, v in factory_kwargs.items() if v is not None})
+    return block