""" 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)}