DoL-1.2-Base / Model /layers /mamba3_layer.py

Add DoL 1.2 model card and inference runtime

431ffa3 verified 3 months ago

14.8 kB

# -*- coding: utf-8 -*-

from __future__ import annotations

import math

import torch
import torch.nn as nn
import torch.nn.functional as F

from Model.config import MIN_OFFICIAL_MAMBA3_D_STATE
from Model.layers.rmsnorm import GroupedRMSNorm


try:
    from mamba_ssm.modules.mamba3 import Mamba3 as OfficialMamba3
except ImportError:
    OfficialMamba3 = None


def official_available() -> bool:
    return OfficialMamba3 is not None


def _validate_official_cfg(cfg) -> None:
    if cfg.mamba_d_state < MIN_OFFICIAL_MAMBA3_D_STATE:
        raise ValueError(
            "official Mamba3 requires mamba_d_state >= "
            f"{MIN_OFFICIAL_MAMBA3_D_STATE}; got {cfg.mamba_d_state}"
        )


def _init_dt_bias(
    nfeatures: int,
    dt_min: float,
    dt_max: float,
    dt_init_floor: float,
) -> torch.Tensor:
    """Upstream Mamba-3 dt_bias initialization.
    Samples ``dt`` uniformly in log space between ``dt_min`` and ``dt_max``,
    floors it, and converts to the additive bias ``dt + log(-expm1(-dt))``
    so that ``softplus(dt_bias) ≈ dt`` at init.
    """

    log_lo = math.log(dt_min)
    log_hi = math.log(dt_max)
    dt = torch.exp(torch.rand(nfeatures) * (log_hi - log_lo) + log_lo)
    dt = torch.clamp(dt, min=dt_init_floor)
    return dt + torch.log(-torch.expm1(-dt))


class NaiveSSM(nn.Module):
    def __init__(
        self,
        d_model: int,
        d_state: int = 16,
        expand: int = 2,
        headdim: int = 64,
        d_conv: int = 4,
        dt_min: float = 0.001,
        dt_max: float = 0.1,
        dt_init_floor: float = 1e-4,
    ):
        super().__init__()

        if d_model <= 0:
            raise ValueError("d_model must be positive")
        if d_state <= 0:
            raise ValueError("d_state must be positive")
        if expand <= 0:
            raise ValueError("expand must be positive")
        if headdim <= 0:
            raise ValueError("headdim must be positive")
        if d_conv <= 0:
            raise ValueError("d_conv must be positive")

        self.d_model = d_model
        self.d_inner = d_model * expand
        self.d_state = d_state
        self.headdim = headdim
        self.d_conv = d_conv

        if self.d_inner % headdim != 0:
            raise ValueError("d_model * expand must be divisible by headdim")

        self.nheads = self.d_inner // headdim

        self.in_proj = nn.Linear(d_model, 2 * self.d_inner, bias=False)

        self.conv1d = nn.Conv1d(
            self.d_inner,
            self.d_inner,
            kernel_size=d_conv,
            groups=self.d_inner,
            padding=d_conv - 1,
            bias=True,
        )

        self.x_proj = nn.Linear(
            self.d_inner,
            self.nheads + 2 * d_state,
            bias=False,
        )
        self.dt_proj = nn.Linear(self.nheads, self.d_inner, bias=True)

        with torch.no_grad():
            self.dt_proj.bias.copy_(
                _init_dt_bias(self.d_inner, dt_min, dt_max, dt_init_floor)
            )
        self.dt_proj.bias._no_weight_decay = True  # type: ignore[attr-defined]

        self.A_log = nn.Parameter(torch.log(torch.rand(self.nheads, d_state) + 0.5))
        self.A_log._no_weight_decay = True  # type: ignore[attr-defined]
        self.D = nn.Parameter(torch.ones(self.nheads))
        self.D._no_weight_decay = True  # type: ignore[attr-defined]
        self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)

    def forward(
        self,
        x: torch.Tensor,
        attn_mask: torch.Tensor | None = None,
        cache=None,
    ) -> torch.Tensor:
        if x.ndim != 3:
            raise ValueError("x must have shape [B, L, d_model]")
        if x.shape[-1] != self.d_model:
            raise ValueError(f"expected d_model={self.d_model}, got {x.shape[-1]}")

        if cache is not None:
            if attn_mask is not None and not bool(attn_mask.all()):
                raise ValueError(
                    "NaiveSSM cached decode requires an all-ones attn_mask "
                    "(no padding); pad-free generation only"
                )
            return self._forward_cached(x, cache)

        bsz, seq_len, _ = x.shape
        dtype = x.dtype

        mask = None
        if attn_mask is not None:
            if attn_mask.shape != (bsz, seq_len):
                raise ValueError("attn_mask must have shape [B, L]")
            mask = attn_mask.to(device=x.device, dtype=torch.float32)

        xz = self.in_proj(x)
        u, z = xz.chunk(2, dim=-1)

        u = self.conv1d(u.transpose(1, 2))[..., :seq_len].transpose(1, 2)
        u = F.silu(u)

        if mask is not None:
            u = u * mask.unsqueeze(-1).to(dtype=u.dtype)
            z = z * mask.unsqueeze(-1).to(dtype=z.dtype)

        params = self.x_proj(u)
        dt, b_param, c_param = params.split(
            [self.nheads, self.d_state, self.d_state],
            dim=-1,
        )
        dt = F.softplus(self.dt_proj(dt))

        u_h = u.reshape(bsz, seq_len, self.nheads, self.headdim)
        dt_h = dt.reshape(bsz, seq_len, self.nheads, self.headdim)

        a = -torch.exp(self.A_log.float())

        state = torch.zeros(
            bsz,
            self.nheads,
            self.headdim,
            self.d_state,
            device=x.device,
            dtype=torch.float32,
        )

        ys: list[torch.Tensor] = []

        for t in range(seq_len):
            dt_t = dt_h[:, t].float()
            u_t = u_h[:, t].float()
            b_t = b_param[:, t].float().view(bsz, 1, 1, self.d_state)
            c_t = c_param[:, t].float().view(bsz, 1, 1, self.d_state)

            da = torch.exp(dt_t.unsqueeze(-1) * a.view(1, self.nheads, 1, self.d_state))
            bu = dt_t.unsqueeze(-1) * b_t * u_t.unsqueeze(-1)

            next_state = state * da + bu
            if mask is not None:
                active = mask[:, t].view(bsz, 1, 1, 1).bool()
                state = torch.where(active, next_state, state)
            else:
                state = next_state

            y_t = (state * c_t).sum(dim=-1)
            y_t = y_t + self.D.view(1, self.nheads, 1) * u_t
            ys.append(y_t)

        y = torch.stack(ys, dim=1)
        y = y.reshape(bsz, seq_len, self.d_inner).to(dtype)

        if mask is not None:
            y = y * mask.unsqueeze(-1).to(dtype=y.dtype)

        y = y * F.silu(z)

        return self.out_proj(y)

    def _forward_cached(self, x: torch.Tensor, cache) -> torch.Tensor:
        """Bit-exact incremental scan over ``x`` (``[B, m, d_model]``).
        Reuses ``cache.conv_window`` (last ``d_conv - 1`` pre-conv columns) and
        ``cache.ssm_state`` so that processing the prompt in one call followed
        by single-token steps reproduces the full :meth:`forward` scan exactly.
        Both fields are updated in place.
        """

        bsz, m, _ = x.shape
        dtype = x.dtype
        k = self.d_conv

        xz = self.in_proj(x)
        u_pre, z = xz.chunk(2, dim=-1)

        u_t = u_pre.transpose(1, 2)  # [B, d_inner, m]

        if k > 1:
            if cache.conv_window is None:
                window = u_t.new_zeros(bsz, self.d_inner, k - 1)
            else:
                window = cache.conv_window
            conv_in = torch.cat([window, u_t], dim=2)  # [B, d_inner, (k-1)+m]
            cache.conv_window = conv_in[..., -(k - 1):]
        else:
            conv_in = u_t

        conv_out = F.conv1d(
            conv_in,
            self.conv1d.weight,
            self.conv1d.bias,
            padding=0,
            groups=self.d_inner,
        )  # [B, d_inner, m]
        u = F.silu(conv_out.transpose(1, 2))  # [B, m, d_inner]

        params = self.x_proj(u)
        dt, b_param, c_param = params.split(
            [self.nheads, self.d_state, self.d_state],
            dim=-1,
        )
        dt = F.softplus(self.dt_proj(dt))

        u_h = u.reshape(bsz, m, self.nheads, self.headdim)
        dt_h = dt.reshape(bsz, m, self.nheads, self.headdim)

        a = -torch.exp(self.A_log.float())

        if cache.ssm_state is None:
            state = torch.zeros(
                bsz,
                self.nheads,
                self.headdim,
                self.d_state,
                device=x.device,
                dtype=torch.float32,
            )
        else:
            state = cache.ssm_state

        ys: list[torch.Tensor] = []
        for t in range(m):
            dt_t = dt_h[:, t].float()
            u_tt = u_h[:, t].float()
            b_t = b_param[:, t].float().view(bsz, 1, 1, self.d_state)
            c_t = c_param[:, t].float().view(bsz, 1, 1, self.d_state)

            da = torch.exp(dt_t.unsqueeze(-1) * a.view(1, self.nheads, 1, self.d_state))
            bu = dt_t.unsqueeze(-1) * b_t * u_tt.unsqueeze(-1)

            state = state * da + bu

            y_t = (state * c_t).sum(dim=-1)
            y_t = y_t + self.D.view(1, self.nheads, 1) * u_tt
            ys.append(y_t)

        cache.ssm_state = state

        y = torch.stack(ys, dim=1)
        y = y.reshape(bsz, m, self.d_inner).to(dtype)
        y = y * F.silu(z)

        return self.out_proj(y)


class Mamba3Layer(nn.Module):
    def __init__(
        self,
        cfg,
        layer_idx: int | None = None,
        num_groups: int = 8,
    ):
        super().__init__()

        self.cfg = cfg
        self.layer_idx = layer_idx

        self.norm = GroupedRMSNorm(
            cfg.d_model,
            num_groups=num_groups,
            eps=cfg.rmsnorm_eps,
        )

        use_official = bool(cfg.use_official_mamba and OfficialMamba3 is not None)

        if cfg.use_official_mamba and OfficialMamba3 is None:
            raise RuntimeError(
                "use_official_mamba=True but mamba_ssm is not importable. "
                "Install with `pip install mamba-ssm` (requires CUDA) or set "
                "cfg.use_official_mamba=False to use the NaiveSSM fallback."
            )

        if use_official:
            _validate_official_cfg(cfg)
            self.mamba = self._build_official(cfg, layer_idx)
            self.backend = "official"
        else:
            self.mamba = NaiveSSM(
                d_model=cfg.d_model,
                d_state=cfg.mamba_d_state,
                expand=cfg.mamba_expand,
                headdim=cfg.mamba_headdim,
                d_conv=cfg.mamba_d_conv,
                dt_min=cfg.mamba_dt_min,
                dt_max=cfg.mamba_dt_max,
                dt_init_floor=cfg.mamba_dt_init_floor,
            )
            self.backend = "fallback"

    def _build_official(self, cfg, layer_idx: int | None):
        """Build the upstream Mamba-3 module with cfg-aligned kwargs.
        Upstream `Mamba3.__init__` does NOT accept ``d_conv`` (it bakes the
        causal short conv into the fused kernel). We forward the parameters
        it does accept; mismatches will surface as a hard ``TypeError`` so
        we never silently degrade.
        """

        return OfficialMamba3(
            d_model=cfg.d_model,
            d_state=cfg.mamba_d_state,
            expand=cfg.mamba_expand,
            headdim=cfg.mamba_headdim,
            chunk_size=cfg.mamba_chunk_size,
            dt_min=cfg.mamba_dt_min,
            dt_max=cfg.mamba_dt_max,
            dt_init_floor=cfg.mamba_dt_init_floor,
            layer_idx=layer_idx,
        )

    def forward(
        self,
        x: torch.Tensor,
        attn_mask: torch.Tensor | None = None,
        cache=None,
        **kwargs,
    ) -> torch.Tensor:
        if x.ndim != 3:
            raise ValueError("x must have shape [B, L, d_model]")

        if cache is not None:
            if not isinstance(self.mamba, NaiveSSM):
                raise NotImplementedError(
                    "incremental decode cache is only supported with the "
                    "NaiveSSM fallback backend; the official Mamba kernels are "
                    "not steppable here"
                )
            residual = x
            mamba_input = self.norm(residual)
            y = self.mamba(mamba_input, attn_mask=attn_mask, cache=cache)
            return residual + y

        mask = None
        if attn_mask is not None:
            if attn_mask.shape != x.shape[:2]:
                raise ValueError("attn_mask must have shape [B, L]")
            if (
                not isinstance(self.mamba, NaiveSSM)
                and not bool(attn_mask.all())
                and not _is_right_padding_mask(attn_mask)
            ):
                raise ValueError(
                    "official Mamba backend only supports all-ones or right-padded "
                    "attention masks; use fallback Mamba for left/interior padding"
                )
            mask = attn_mask.to(device=x.device, dtype=x.dtype).unsqueeze(-1)

        residual = x if mask is None else x * mask
        mamba_input = self.norm(residual)

        if isinstance(self.mamba, NaiveSSM):
            y = self.mamba(mamba_input, attn_mask=attn_mask)
        else:
            y = self.mamba(mamba_input)
            if mask is not None:
                y = y * mask

        out = residual + y
        if mask is not None:
            out = out * mask
        return out


def _is_right_padding_mask(attn_mask: torch.Tensor) -> bool:
    if attn_mask.ndim != 2:
        return False
    mask = attn_mask.to(dtype=torch.long)
    if mask.numel() == 0:
        return True
    if not bool(((mask == 0) | (mask == 1)).all()):
        return False
    return bool((mask[:, 1:] <= mask[:, :-1]).all())


def _check() -> None:
    from Model.config import tiny_config

    torch.manual_seed(0)

    cfg = tiny_config()
    layer = Mamba3Layer(cfg, layer_idx=0)
    layer.eval()

    print("Mamba3Layer")
    print(f"  backend: {layer.backend}")
    print(f"  official_available: {official_available()}")

    bsz, seq_len = 2, 16
    x = torch.randn(bsz, seq_len, cfg.d_model)

    y = layer(x)

    print(f"  shape: {tuple(x.shape)} -> {tuple(y.shape)}")
    print(f"  params: {sum(p.numel() for p in layer.parameters()):,}")

    x2 = torch.randn(bsz, seq_len, cfg.d_model, requires_grad=True)
    loss = layer(x2).sum()
    loss.backward()

    print(f"  grad_norm: {x2.grad.norm().item():.6f}")

    x3 = torch.randn(1, 8, cfg.d_model)
    out_full = layer(x3)

    x3_mod = x3.clone()
    x3_mod[0, 5:] += 10.0
    out_mod = layer(x3_mod)

    diff = (out_full[0, :5] - out_mod[0, :5]).abs().max().item()
    print(f"  causal_diff_before_changed_span: {diff:.6e}")

    ignored = layer(
        x,
        word_pos=torch.zeros(bsz, seq_len, dtype=torch.long),
        morph_depth=torch.zeros(bsz, seq_len, dtype=torch.long),
        attn_mask=torch.ones(bsz, seq_len, dtype=torch.long),
    )

    print(f"  kwargs_ignored_shape: {tuple(ignored.shape)}")


if __name__ == "__main__":
    _check()

	# -- coding: utf-8 --

	from __future__ import annotations

	import math

	import torch
	import torch.nn as nn
	import torch.nn.functional as F

	from Model.config import MIN_OFFICIAL_MAMBA3_D_STATE
	from Model.layers.rmsnorm import GroupedRMSNorm


	try:
	from mamba_ssm.modules.mamba3 import Mamba3 as OfficialMamba3
	except ImportError:
	OfficialMamba3 = None


	def official_available() -> bool:
	return OfficialMamba3 is not None


	def _validate_official_cfg(cfg) -> None:
	if cfg.mamba_d_state < MIN_OFFICIAL_MAMBA3_D_STATE:
	raise ValueError(
	"official Mamba3 requires mamba_d_state >= "
	f"{MIN_OFFICIAL_MAMBA3_D_STATE}; got {cfg.mamba_d_state}"
	)


	def _init_dt_bias(
	nfeatures: int,
	dt_min: float,
	dt_max: float,
	dt_init_floor: float,
	) -> torch.Tensor:
	"""Upstream Mamba-3 dt_bias initialization.

	Samples ``dt`` uniformly in log space between ``dt_min`` and ``dt_max``,
	floors it, and converts to the additive bias ``dt + log(-expm1(-dt))``
	so that ``softplus(dt_bias) ≈ dt`` at init.
	"""

	log_lo = math.log(dt_min)
	log_hi = math.log(dt_max)
	dt = torch.exp(torch.rand(nfeatures) * (log_hi - log_lo) + log_lo)
	dt = torch.clamp(dt, min=dt_init_floor)
	return dt + torch.log(-torch.expm1(-dt))


	class NaiveSSM(nn.Module):
	def __init__(
	self,
	d_model: int,
	d_state: int = 16,
	expand: int = 2,
	headdim: int = 64,
	d_conv: int = 4,
	dt_min: float = 0.001,
	dt_max: float = 0.1,
	dt_init_floor: float = 1e-4,
	):
	super().__init__()

	if d_model <= 0:
	raise ValueError("d_model must be positive")
	if d_state <= 0:
	raise ValueError("d_state must be positive")
	if expand <= 0:
	raise ValueError("expand must be positive")
	if headdim <= 0:
	raise ValueError("headdim must be positive")
	if d_conv <= 0:
	raise ValueError("d_conv must be positive")

	self.d_model = d_model
	self.d_inner = d_model * expand
	self.d_state = d_state
	self.headdim = headdim
	self.d_conv = d_conv

	if self.d_inner % headdim != 0:
	raise ValueError("d_model * expand must be divisible by headdim")

	self.nheads = self.d_inner // headdim

	self.in_proj = nn.Linear(d_model, 2 * self.d_inner, bias=False)

	self.conv1d = nn.Conv1d(
	self.d_inner,
	self.d_inner,
	kernel_size=d_conv,
	groups=self.d_inner,
	padding=d_conv - 1,
	bias=True,
	)

	self.x_proj = nn.Linear(
	self.d_inner,
	self.nheads + 2 * d_state,
	bias=False,
	)
	self.dt_proj = nn.Linear(self.nheads, self.d_inner, bias=True)

	with torch.no_grad():
	self.dt_proj.bias.copy_(
	_init_dt_bias(self.d_inner, dt_min, dt_max, dt_init_floor)
	)
	self.dt_proj.bias._no_weight_decay = True # type: ignore[attr-defined]

	self.A_log = nn.Parameter(torch.log(torch.rand(self.nheads, d_state) + 0.5))
	self.A_log._no_weight_decay = True # type: ignore[attr-defined]
	self.D = nn.Parameter(torch.ones(self.nheads))
	self.D._no_weight_decay = True # type: ignore[attr-defined]
	self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)

	def forward(
	self,
	x: torch.Tensor,
	attn_mask: torch.Tensor \| None = None,
	cache=None,
	) -> torch.Tensor:
	if x.ndim != 3:
	raise ValueError("x must have shape [B, L, d_model]")
	if x.shape[-1] != self.d_model:
	raise ValueError(f"expected d_model={self.d_model}, got {x.shape[-1]}")

	if cache is not None:
	if attn_mask is not None and not bool(attn_mask.all()):
	raise ValueError(
	"NaiveSSM cached decode requires an all-ones attn_mask "
	"(no padding); pad-free generation only"
	)
	return self._forward_cached(x, cache)

	bsz, seq_len, _ = x.shape
	dtype = x.dtype

	mask = None
	if attn_mask is not None:
	if attn_mask.shape != (bsz, seq_len):
	raise ValueError("attn_mask must have shape [B, L]")
	mask = attn_mask.to(device=x.device, dtype=torch.float32)

	xz = self.in_proj(x)
	u, z = xz.chunk(2, dim=-1)

	u = self.conv1d(u.transpose(1, 2))[..., :seq_len].transpose(1, 2)
	u = F.silu(u)

	if mask is not None:
	u = u * mask.unsqueeze(-1).to(dtype=u.dtype)
	z = z * mask.unsqueeze(-1).to(dtype=z.dtype)

	params = self.x_proj(u)
	dt, b_param, c_param = params.split(
	[self.nheads, self.d_state, self.d_state],
	dim=-1,
	)
	dt = F.softplus(self.dt_proj(dt))

	u_h = u.reshape(bsz, seq_len, self.nheads, self.headdim)
	dt_h = dt.reshape(bsz, seq_len, self.nheads, self.headdim)

	a = -torch.exp(self.A_log.float())

	state = torch.zeros(
	bsz,
	self.nheads,
	self.headdim,
	self.d_state,
	device=x.device,
	dtype=torch.float32,
	)

	ys: list[torch.Tensor] = []

	for t in range(seq_len):
	dt_t = dt_h[:, t].float()
	u_t = u_h[:, t].float()
	b_t = b_param[:, t].float().view(bsz, 1, 1, self.d_state)
	c_t = c_param[:, t].float().view(bsz, 1, 1, self.d_state)

	da = torch.exp(dt_t.unsqueeze(-1) * a.view(1, self.nheads, 1, self.d_state))
	bu = dt_t.unsqueeze(-1) * b_t * u_t.unsqueeze(-1)

	next_state = state * da + bu
	if mask is not None:
	active = mask[:, t].view(bsz, 1, 1, 1).bool()
	state = torch.where(active, next_state, state)
	else:
	state = next_state

	y_t = (state * c_t).sum(dim=-1)
	y_t = y_t + self.D.view(1, self.nheads, 1) * u_t
	ys.append(y_t)

	y = torch.stack(ys, dim=1)
	y = y.reshape(bsz, seq_len, self.d_inner).to(dtype)

	if mask is not None:
	y = y * mask.unsqueeze(-1).to(dtype=y.dtype)

	y = y * F.silu(z)

	return self.out_proj(y)

	def _forward_cached(self, x: torch.Tensor, cache) -> torch.Tensor:
	"""Bit-exact incremental scan over ``x`` (``[B, m, d_model]``).

	Reuses ``cache.conv_window`` (last ``d_conv - 1`` pre-conv columns) and
	``cache.ssm_state`` so that processing the prompt in one call followed
	by single-token steps reproduces the full :meth:`forward` scan exactly.
	Both fields are updated in place.
	"""

	bsz, m, _ = x.shape
	dtype = x.dtype
	k = self.d_conv

	xz = self.in_proj(x)
	u_pre, z = xz.chunk(2, dim=-1)

	u_t = u_pre.transpose(1, 2) # [B, d_inner, m]

	if k > 1:
	if cache.conv_window is None:
	window = u_t.new_zeros(bsz, self.d_inner, k - 1)
	else:
	window = cache.conv_window
	conv_in = torch.cat([window, u_t], dim=2) # [B, d_inner, (k-1)+m]
	cache.conv_window = conv_in[..., -(k - 1):]
	else:
	conv_in = u_t

	conv_out = F.conv1d(
	conv_in,
	self.conv1d.weight,
	self.conv1d.bias,
	padding=0,
	groups=self.d_inner,
	) # [B, d_inner, m]
	u = F.silu(conv_out.transpose(1, 2)) # [B, m, d_inner]

	params = self.x_proj(u)
	dt, b_param, c_param = params.split(
	[self.nheads, self.d_state, self.d_state],
	dim=-1,
	)
	dt = F.softplus(self.dt_proj(dt))

	u_h = u.reshape(bsz, m, self.nheads, self.headdim)
	dt_h = dt.reshape(bsz, m, self.nheads, self.headdim)

	a = -torch.exp(self.A_log.float())

	if cache.ssm_state is None:
	state = torch.zeros(
	bsz,
	self.nheads,
	self.headdim,
	self.d_state,
	device=x.device,
	dtype=torch.float32,
	)
	else:
	state = cache.ssm_state

	ys: list[torch.Tensor] = []
	for t in range(m):
	dt_t = dt_h[:, t].float()
	u_tt = u_h[:, t].float()
	b_t = b_param[:, t].float().view(bsz, 1, 1, self.d_state)
	c_t = c_param[:, t].float().view(bsz, 1, 1, self.d_state)

	da = torch.exp(dt_t.unsqueeze(-1) * a.view(1, self.nheads, 1, self.d_state))
	bu = dt_t.unsqueeze(-1) * b_t * u_tt.unsqueeze(-1)

	state = state * da + bu

	y_t = (state * c_t).sum(dim=-1)
	y_t = y_t + self.D.view(1, self.nheads, 1) * u_tt
	ys.append(y_t)

	cache.ssm_state = state

	y = torch.stack(ys, dim=1)
	y = y.reshape(bsz, m, self.d_inner).to(dtype)
	y = y * F.silu(z)

	return self.out_proj(y)


	class Mamba3Layer(nn.Module):
	def __init__(
	self,
	cfg,
	layer_idx: int \| None = None,
	num_groups: int = 8,
	):
	super().__init__()

	self.cfg = cfg
	self.layer_idx = layer_idx

	self.norm = GroupedRMSNorm(
	cfg.d_model,
	num_groups=num_groups,
	eps=cfg.rmsnorm_eps,
	)

	use_official = bool(cfg.use_official_mamba and OfficialMamba3 is not None)

	if cfg.use_official_mamba and OfficialMamba3 is None:
	raise RuntimeError(
	"use_official_mamba=True but mamba_ssm is not importable. "
	"Install with `pip install mamba-ssm` (requires CUDA) or set "
	"cfg.use_official_mamba=False to use the NaiveSSM fallback."
	)

	if use_official:
	_validate_official_cfg(cfg)
	self.mamba = self._build_official(cfg, layer_idx)
	self.backend = "official"
	else:
	self.mamba = NaiveSSM(
	d_model=cfg.d_model,
	d_state=cfg.mamba_d_state,
	expand=cfg.mamba_expand,
	headdim=cfg.mamba_headdim,
	d_conv=cfg.mamba_d_conv,
	dt_min=cfg.mamba_dt_min,
	dt_max=cfg.mamba_dt_max,
	dt_init_floor=cfg.mamba_dt_init_floor,
	)
	self.backend = "fallback"

	def _build_official(self, cfg, layer_idx: int \| None):
	"""Build the upstream Mamba-3 module with cfg-aligned kwargs.

	Upstream `Mamba3.__init__` does NOT accept ``d_conv`` (it bakes the
	causal short conv into the fused kernel). We forward the parameters
	it does accept; mismatches will surface as a hard ``TypeError`` so
	we never silently degrade.
	"""

	return OfficialMamba3(
	d_model=cfg.d_model,
	d_state=cfg.mamba_d_state,
	expand=cfg.mamba_expand,
	headdim=cfg.mamba_headdim,
	chunk_size=cfg.mamba_chunk_size,
	dt_min=cfg.mamba_dt_min,
	dt_max=cfg.mamba_dt_max,
	dt_init_floor=cfg.mamba_dt_init_floor,
	layer_idx=layer_idx,
	)

	def forward(
	self,
	x: torch.Tensor,
	attn_mask: torch.Tensor \| None = None,
	cache=None,
	**kwargs,
	) -> torch.Tensor:
	if x.ndim != 3:
	raise ValueError("x must have shape [B, L, d_model]")

	if cache is not None:
	if not isinstance(self.mamba, NaiveSSM):
	raise NotImplementedError(
	"incremental decode cache is only supported with the "
	"NaiveSSM fallback backend; the official Mamba kernels are "
	"not steppable here"
	)
	residual = x
	mamba_input = self.norm(residual)
	y = self.mamba(mamba_input, attn_mask=attn_mask, cache=cache)
	return residual + y

	mask = None
	if attn_mask is not None:
	if attn_mask.shape != x.shape[:2]:
	raise ValueError("attn_mask must have shape [B, L]")
	if (
	not isinstance(self.mamba, NaiveSSM)
	and not bool(attn_mask.all())
	and not _is_right_padding_mask(attn_mask)
	):
	raise ValueError(
	"official Mamba backend only supports all-ones or right-padded "
	"attention masks; use fallback Mamba for left/interior padding"
	)
	mask = attn_mask.to(device=x.device, dtype=x.dtype).unsqueeze(-1)

	residual = x if mask is None else x * mask
	mamba_input = self.norm(residual)

	if isinstance(self.mamba, NaiveSSM):
	y = self.mamba(mamba_input, attn_mask=attn_mask)
	else:
	y = self.mamba(mamba_input)
	if mask is not None:
	y = y * mask

	out = residual + y
	if mask is not None:
	out = out * mask
	return out


	def _is_right_padding_mask(attn_mask: torch.Tensor) -> bool:
	if attn_mask.ndim != 2:
	return False
	mask = attn_mask.to(dtype=torch.long)
	if mask.numel() == 0:
	return True
	if not bool(((mask == 0) \| (mask == 1)).all()):
	return False
	return bool((mask[:, 1:] <= mask[:, :-1]).all())


	def _check() -> None:
	from Model.config import tiny_config

	torch.manual_seed(0)

	cfg = tiny_config()
	layer = Mamba3Layer(cfg, layer_idx=0)
	layer.eval()

	print("Mamba3Layer")
	print(f" backend: {layer.backend}")
	print(f" official_available: {official_available()}")

	bsz, seq_len = 2, 16
	x = torch.randn(bsz, seq_len, cfg.d_model)

	y = layer(x)

	print(f" shape: {tuple(x.shape)} -> {tuple(y.shape)}")
	print(f" params: {sum(p.numel() for p in layer.parameters()):,}")

	x2 = torch.randn(bsz, seq_len, cfg.d_model, requires_grad=True)
	loss = layer(x2).sum()
	loss.backward()

	print(f" grad_norm: {x2.grad.norm().item():.6f}")

	x3 = torch.randn(1, 8, cfg.d_model)
	out_full = layer(x3)

	x3_mod = x3.clone()
	x3_mod[0, 5:] += 10.0
	out_mod = layer(x3_mod)

	diff = (out_full[0, :5] - out_mod[0, :5]).abs().max().item()
	print(f" causal_diff_before_changed_span: {diff:.6e}")

	ignored = layer(
	x,
	word_pos=torch.zeros(bsz, seq_len, dtype=torch.long),
	morph_depth=torch.zeros(bsz, seq_len, dtype=torch.long),
	attn_mask=torch.ones(bsz, seq_len, dtype=torch.long),
	)

	print(f" kwargs_ignored_shape: {tuple(ignored.shape)}")


	if __name__ == "__main__":
	_check()