Proact-VL-Washing/mage_demo/mamba_ssm/models/mixer_seq_simple.py

"""Pure-PyTorch `create_block` — Mamba-1 mixer + residual block, no CUDA kernels.
“Pure-PyTorch `create_block` — Mamba-1 混合器 + 残余块，没有 CUDA 内核。
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