modeling_mamba3.py

"""
Mamba3 SISO — pure-PyTorch reference implementation.

Architecture per: https://goombalab.github.io/blog/2026/mamba3-part2/
  - Exponential-trapezoidal discretization with data-dependent lambda
  - Complex SSM via RoPE on B/C projections (data-dependent angles)
  - Data-dependent A via softplus gate from in_proj
  - BCNorm for training stability
  - in_proj layout: [z, x, B, C, dd_dt, dd_A, trap, angles]

No Triton/TileLang dependencies — sequential scan, CPU-compatible.
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass, field
from typing import Optional


@dataclass
class Mamba3Config:
    d_model: int = 2560
    n_layer: int = 64
    vocab_size: int = 50288
    d_state: int = 128
    expand: int = 2
    headdim: int = 64
    ngroups: int = 1
    rope_fraction: float = 0.5
    dt_min: float = 0.001
    dt_max: float = 0.1
    dt_init_floor: float = 1e-4
    is_safe_A: bool = True
    tie_embeddings: bool = True
    pad_vocab_size_multiple: int = 16


class Mamba3Mixer(nn.Module):
    def __init__(self, config: Mamba3Config, layer_idx: Optional[int] = None):
        super().__init__()
        self.d_model = config.d_model
        self.d_state = config.d_state
        self.headdim = config.headdim
        self.ngroups = config.ngroups
        self.is_safe_A = config.is_safe_A
        self.layer_idx = layer_idx

        self.d_inner = int(config.expand * config.d_model)
        assert self.d_inner % self.headdim == 0
        self.nheads = self.d_inner // self.headdim
        self.heads_per_group = self.nheads // self.ngroups

        # num_rope_angles: half the rotary d_state → each angle controls one (cos,sin) pair
        rope_dim = int(config.d_state * config.rope_fraction)
        if rope_dim % 2 != 0:
            rope_dim -= 1
        self.num_rope_angles = rope_dim // 2

        # in_proj layout: [z, x, B, C, dd_dt, dd_A, trap, angles]
        self.d_in_proj = (
            2 * self.d_inner
            + 2 * self.d_state * self.ngroups
            + 3 * self.nheads
            + self.num_rope_angles
        )
        self.in_proj = nn.Linear(self.d_model, self.d_in_proj, bias=False)

        # Learned dt_bias per head (softplus-parameterized)
        _dt = torch.exp(
            torch.rand(self.nheads) * (math.log(config.dt_max) - math.log(config.dt_min))
            + math.log(config.dt_min)
        ).clamp(min=config.dt_init_floor)
        self.dt_bias = nn.Parameter(_dt + torch.log(-torch.expm1(-_dt)))
        self.dt_bias._no_weight_decay = True

        # B/C additive biases (nheads × d_state), init=1 per reference
        self.B_bias = nn.Parameter(torch.ones(self.nheads, self.d_state))
        self.C_bias = nn.Parameter(torch.ones(self.nheads, self.d_state))

        # BCNorm: RMSNorm applied per-group on d_state dimension
        self.B_norm = nn.RMSNorm(self.d_state, eps=1e-5)
        self.C_norm = nn.RMSNorm(self.d_state, eps=1e-5)

        # D skip connection per head
        self.D = nn.Parameter(torch.ones(self.nheads))
        self.D._no_weight_decay = True

        self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=False)

    @staticmethod
    def _apply_rope(bc: torch.Tensor, cum_angles: torch.Tensor) -> torch.Tensor:
        """
        Apply rotary embedding to B or C.
        bc: (B, L, G, d_state)
        cum_angles: (B, L, num_rope_angles) — cumulative rotation angles
        Rotates the first 2*nr dimensions of d_state, leaves the rest.
        """
        nr = cum_angles.shape[-1]
        bc_rot, bc_static = bc[..., :2 * nr], bc[..., 2 * nr:]

        x1, x2 = bc_rot[..., :nr], bc_rot[..., nr:]
        cos_a = torch.cos(cum_angles).unsqueeze(2)  # (B, L, 1, nr)
        sin_a = torch.sin(cum_angles).unsqueeze(2)

        return torch.cat([x1 * cos_a - x2 * sin_a,
                          x1 * sin_a + x2 * cos_a,
                          bc_static], dim=-1)

    def forward(self, hidden_states: torch.Tensor, inference_params=None) -> torch.Tensor:
        batch, seqlen, _ = hidden_states.shape

        proj = self.in_proj(hidden_states)
        z, x, B_raw, C_raw, dd_dt, dd_A, trap_raw, angles_raw = torch.split(
            proj,
            [
                self.d_inner, self.d_inner,
                self.d_state * self.ngroups,
                self.d_state * self.ngroups,
                self.nheads, self.nheads, self.nheads,
                self.num_rope_angles,
            ],
            dim=-1,
        )

        # Data-dependent A: A < 0 enforced by softplus; is_safe_A shifts further negative
        A = -F.softplus(dd_A.float())
        if self.is_safe_A:
            A = A - 1.0
        DT = F.softplus(dd_dt.float() + self.dt_bias)  # (B, L, H)
        ADT = A * DT  # log(decay_factor) per head per token
        alpha = torch.exp(ADT)  # (B, L, H)

        # Data-dependent trapezoidal lambda ∈ (0, 1)
        trap = torch.sigmoid(trap_raw.float())  # (B, L, H)

        # BCNorm + bias
        B_raw = B_raw.view(batch, seqlen, self.ngroups, self.d_state)
        C_raw = C_raw.view(batch, seqlen, self.ngroups, self.d_state)
        B_raw = self.B_norm(B_raw)
        C_raw = self.C_norm(C_raw)
        # Bias: (nheads, d_state) → take first entry per group for ngroups=1
        B_bias_g = self.B_bias.view(self.ngroups, self.heads_per_group, self.d_state)[:, 0, :]
        C_bias_g = self.C_bias.view(self.ngroups, self.heads_per_group, self.d_state)[:, 0, :]
        B_raw = B_raw + B_bias_g  # broadcast over (B, L)
        C_raw = C_raw + C_bias_g

        # RoPE: cumulative per-token angles (B, L, num_rope_angles)
        cum_angles = torch.cumsum(angles_raw.float(), dim=1)

        # B gets positive rotation, C gets conjugate (negative rotation)
        B = self._apply_rope(B_raw, cum_angles)   # (B, L, G, d_state)
        C = self._apply_rope(C_raw, -cum_angles)  # conjugate

        # Expand groups → heads
        B = B.repeat_interleave(self.heads_per_group, dim=2)  # (B, L, H, d_state)
        C = C.repeat_interleave(self.heads_per_group, dim=2)

        x = x.view(batch, seqlen, self.nheads, self.headdim)  # (B, L, H, P)

        # Sequential trapezoidal SSM scan
        # h_t = α_t·h_{t-1} + (1-λ_t)·Δ_t·α_t·(B_{t-1}⊗x_{t-1}) + λ_t·Δ_t·(B_t⊗x_t)
        h = torch.zeros(batch, self.nheads, self.headdim, self.d_state,
                        dtype=torch.float32, device=hidden_states.device)
        x_prev = torch.zeros(batch, self.nheads, self.headdim,
                             dtype=torch.float32, device=hidden_states.device)
        B_prev = torch.zeros(batch, self.nheads, self.d_state,
                             dtype=torch.float32, device=hidden_states.device)

        outputs = []
        for t in range(seqlen):
            alpha_t = alpha[:, t].view(batch, self.nheads, 1, 1)
            DT_t    = DT[:, t].view(batch, self.nheads, 1, 1)
            trap_t  = trap[:, t].view(batch, self.nheads, 1, 1)

            B_t = B[:, t].float()  # (B, H, d_state)
            C_t = C[:, t].float()
            x_t = x[:, t].float()  # (B, H, P)

            bx_curr = torch.einsum("bhp,bhs->bhps", x_t, B_t)
            bx_prev = torch.einsum("bhp,bhs->bhps", x_prev, B_prev)

            h = (alpha_t * h
                 + (1.0 - trap_t) * DT_t * alpha_t * bx_prev
                 + trap_t * DT_t * bx_curr)

            # y_t = Cᵀh + D·x
            y_t = (torch.einsum("bhps,bhs->bhp", h, C_t)
                   + self.D.view(1, self.nheads, 1) * x_t)

            outputs.append(y_t.to(hidden_states.dtype))
            x_prev = x_t
            B_prev = B_t

        y = torch.stack(outputs, dim=1)  # (B, L, H, P)

        # Headwise gate: RMSNorm(y) * SiLU(z)
        z = z.view(batch, seqlen, self.nheads, self.headdim)
        y = F.rms_norm(y.float(), [self.headdim]).to(y.dtype) * F.silu(z.float().to(y.dtype))

        y = y.reshape(batch, seqlen, self.d_inner)
        return self.out_proj(y)


class Mamba3Block(nn.Module):
    def __init__(self, config: Mamba3Config, layer_idx: Optional[int] = None):
        super().__init__()
        self.norm = nn.RMSNorm(config.d_model, eps=1e-5)
        self.mixer = Mamba3Mixer(config, layer_idx=layer_idx)

    def forward(self, hidden_states: torch.Tensor, residual=None, inference_params=None):
        residual = hidden_states if residual is None else residual
        return self.mixer(self.norm(hidden_states), inference_params=inference_params) + residual


class Mamba3CausalLM(nn.Module):
    def __init__(self, config: Mamba3Config):
        super().__init__()
        self.config = config
        vocab = (math.ceil(config.vocab_size / config.pad_vocab_size_multiple)
                 * config.pad_vocab_size_multiple)
        self.embeddings = nn.Embedding(vocab, config.d_model)
        self.layers = nn.ModuleList(
            [Mamba3Block(config, layer_idx=i) for i in range(config.n_layer)]
        )
        self.norm_f = nn.RMSNorm(config.d_model, eps=1e-5)
        self.lm_head = nn.Linear(config.d_model, vocab, bias=False)
        if config.tie_embeddings:
            self.lm_head.weight = self.embeddings.weight

    def forward(self, input_ids: torch.Tensor, inference_params=None):
        hidden_states = self.embeddings(input_ids)
        residual = None
        for layer in self.layers:
            hidden_states = layer(hidden_states, residual=residual,
                                  inference_params=inference_params)
            residual = hidden_states
        hidden_states = self.norm_f(hidden_states)
        return {"logits": self.lm_head(hidden_states)}