Spaces:
Running on Zero
Running on Zero
| """PyTorch Mamba-3 recurrence for CPU and CUDA inference. | |
| Adapted by Continuous Intelligence from Mamba's module and | |
| SISO/MIMO reference calculations: | |
| https://github.com/state-spaces/mamba/tree/a14b1dff0454a3bc27d9eb31355dc01e4b2490ec | |
| RMS normalisation is adapted from xLSTM Large: | |
| https://github.com/NX-AI/xlstm/blob/ab22eadbd245f293dd8dec38ed29963d73758a12/xlstm/xlstm_large/components.py | |
| Copyright (c) 2026, Dao AI Lab, Goombalab. | |
| Copyright 2024 NXAI GmbH and its affiliates 2024 | |
| Upstream portions are licensed under Apache-2.0; see NOTICE. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| import torch | |
| import torch.nn.functional as F | |
| from einops import rearrange | |
| from torch import Tensor, nn | |
| class RMSNorm(nn.Module): | |
| """RMS norm with optional float32 reductions.""" | |
| def __init__( | |
| self, | |
| num_features: int, | |
| eps: float = 1e-06, | |
| force_float32_reductions: bool = True, | |
| use_bias: bool = False, | |
| ) -> None: | |
| super().__init__() | |
| self.eps = eps | |
| self.force_float32_reductions = force_float32_reductions | |
| self.weight = nn.Parameter(torch.ones(num_features)) | |
| self.bias = nn.Parameter(torch.zeros(num_features)) if use_bias else None | |
| def forward(self, x: Tensor) -> Tensor: | |
| in_dtype = x.dtype | |
| if self.force_float32_reductions: | |
| x = x.float() | |
| x = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) | |
| x = x.to(in_dtype) | |
| x = x * self.weight | |
| if self.bias is not None: | |
| x = x + self.bias | |
| return x | |
| def _rotate_pairs(x: Tensor, cos: Tensor, sin: Tensor, interleaved: bool) -> Tensor: | |
| """Rotate pairs of the last axis; angles beyond ``cos`` leave their pairs unrotated.""" | |
| if interleaved: | |
| pairs = x.view(*x.shape[:-1], -1, 2) | |
| first, second = (pairs[..., 0], pairs[..., 1]) | |
| else: | |
| half = x.shape[-1] // 2 | |
| first, second = (x[..., :half], x[..., half:]) | |
| if cos.shape[-1] < first.shape[-1]: | |
| pad = first.shape[-1] - cos.shape[-1] | |
| cos = F.pad(cos, (0, pad), value=1.0) | |
| sin = F.pad(sin, (0, pad), value=0.0) | |
| rotated_first = first * cos - second * sin | |
| rotated_second = first * sin + second * cos | |
| if interleaved: | |
| return torch.stack([rotated_first, rotated_second], dim=-1).view_as(x) | |
| return torch.cat([rotated_first, rotated_second], dim=-1) | |
| class Mamba3Reference(nn.Module): | |
| """Mamba-3 block evaluated step by step in plain PyTorch.""" | |
| def __init__( | |
| self, | |
| d_model: int, | |
| d_state: int = 128, | |
| expand: int = 2, | |
| headdim: int = 64, | |
| is_mimo: bool = False, | |
| mimo_rank: int = 4, | |
| ) -> None: | |
| super().__init__() | |
| if type(d_state) is not int or d_state < 4 or d_state % 2: | |
| raise ValueError("mamba_state_dim must be an even integer of at least 4") | |
| self.d_state = d_state | |
| self.is_mimo = is_mimo | |
| self.mimo_rank = mimo_rank if is_mimo else 1 | |
| self.d_inner = expand * d_model | |
| if self.d_inner % headdim != 0: | |
| raise ValueError(f"expand * d_model {self.d_inner} not divisible by headdim {headdim}") | |
| self.headdim = headdim | |
| self.nheads = self.d_inner // headdim | |
| split = d_state // 2 - d_state // 2 % 2 | |
| self.num_rope_angles = split // 2 | |
| bc = self.d_state * self.mimo_rank | |
| d_in_proj = 2 * self.d_inner + 2 * bc + 3 * self.nheads + self.num_rope_angles | |
| self.in_proj = nn.Linear(d_model, d_in_proj, bias=False) | |
| dt = torch.exp( | |
| torch.rand(self.nheads) * (math.log(0.1) - math.log(0.001)) + math.log(0.001) | |
| ).clamp(min=1e-4) | |
| self.dt_bias = nn.Parameter(dt + torch.log(-torch.expm1(-dt))) | |
| self.B_bias = nn.Parameter(torch.ones(self.nheads, self.mimo_rank, d_state)) | |
| self.C_bias = nn.Parameter(torch.ones(self.nheads, self.mimo_rank, d_state)) | |
| self.B_norm = RMSNorm(d_state, eps=1e-05) | |
| self.C_norm = RMSNorm(d_state, eps=1e-05) | |
| if is_mimo: | |
| shape = (self.nheads, self.mimo_rank, headdim) | |
| self.mimo_x = nn.Parameter(torch.ones(shape) / self.mimo_rank) | |
| self.mimo_z = nn.Parameter(torch.ones(shape)) | |
| self.mimo_o = nn.Parameter(torch.ones(shape) / self.mimo_rank) | |
| self.D = nn.Parameter(torch.ones(self.nheads)) | |
| self.out_proj = nn.Linear(self.d_inner, d_model, bias=False) | |
| def forward(self, u: Tensor) -> Tensor: | |
| batch, seq, _ = u.shape | |
| bc = self.d_state * self.mimo_rank | |
| z, x, B, C, dd_dt, dd_A, trap, angles = torch.split( | |
| self.in_proj(u), | |
| [ | |
| self.d_inner, | |
| self.d_inner, | |
| bc, | |
| bc, | |
| self.nheads, | |
| self.nheads, | |
| self.nheads, | |
| self.num_rope_angles, | |
| ], | |
| dim=-1, | |
| ) | |
| z = rearrange(z, "b l (h p) -> b l h p", p=self.headdim).to(torch.float32) | |
| x = rearrange(x, "b l (h p) -> b l h p", p=self.headdim).to(torch.float32) | |
| B = rearrange(B, "b l (r g n) -> b l r g n", r=self.mimo_rank, g=1) | |
| C = rearrange(C, "b l (r g n) -> b l r g n", r=self.mimo_rank, g=1) | |
| A = torch.clamp(-F.softplus(dd_A.to(torch.float32)), max=-1e-4) | |
| DT = F.softplus(dd_dt.to(torch.float32) + self.dt_bias) | |
| alpha = torch.exp(A * DT) | |
| gate = torch.sigmoid(trap.to(torch.float32)) | |
| beta = (1.0 - gate) * DT * alpha | |
| gamma = gate * DT | |
| k = self.B_norm(B).expand(-1, -1, -1, self.nheads, -1).to(torch.float32) | |
| q = self.C_norm(C).expand(-1, -1, -1, self.nheads, -1).to(torch.float32) | |
| k = k + rearrange(self.B_bias, "h r n -> r h n") | |
| q = q + rearrange(self.C_bias, "h r n -> r h n") | |
| turns = torch.tanh(angles.to(torch.float32))[:, :, None, :] * DT[..., None] * math.pi | |
| cumulative = torch.cumsum(turns, dim=1) | |
| cos, sin = (torch.cos(cumulative)[:, :, None], torch.sin(cumulative)[:, :, None]) | |
| q = _rotate_pairs(q, cos, sin, interleaved=not self.is_mimo) | |
| k = _rotate_pairs(k, cos, sin, interleaved=not self.is_mimo) | |
| if self.is_mimo: | |
| v = torch.einsum("blhp,hrp->blrhp", x, self.mimo_x.to(torch.float32)) | |
| gates = torch.einsum("blhp,hrp->blrhp", z, self.mimo_z.to(torch.float32)) | |
| else: | |
| v, gates = (x[:, :, None], z[:, :, None]) | |
| state = u.new_zeros((batch, self.nheads, self.headdim, self.d_state), dtype=torch.float32) | |
| previous = torch.zeros_like(state) | |
| outputs = [] | |
| for t in range(seq): | |
| current = torch.einsum("brhn,brhp->bhpn", k[:, t], v[:, t]) | |
| state = ( | |
| alpha[:, t, :, None, None] * state | |
| + beta[:, t, :, None, None] * previous | |
| + gamma[:, t, :, None, None] * current | |
| ) | |
| previous = current | |
| step = torch.einsum("bhpn,brhn->brhp", state, q[:, t]) | |
| step = step + self.D[None, None, :, None] * v[:, t] | |
| step = step * F.silu(gates[:, t]) | |
| if self.is_mimo: | |
| step = torch.einsum("brhp,hrp->bhp", step, self.mimo_o.to(torch.float32)) | |
| else: | |
| step = step[:, 0] | |
| outputs.append(step) | |
| y = rearrange(torch.stack(outputs, dim=1), "b l h p -> b l (h p)") | |
| return self.out_proj(y.to(u.dtype)) | |
| class Mamba3(Mamba3Reference): | |
| """Checkpoint-compatible cell with fused CUDA inference when configured by the loader.""" | |
| fused = False | |
| def forward(self, u: Tensor) -> Tensor: | |
| if self.fused and u.is_cuda and not torch.is_grad_enabled(): | |
| from ._cuda import mamba3_forward | |
| return mamba3_forward(self, u) | |
| return super().forward(u) | |