modeling/mamba2_block.py

"""
Mamba2 Block wrapper for hybrid architecture.

Mamba2 improves on Mamba1 with:
- Multi-head SSM (similar to multi-head attention)
- Larger effective state size
- Hardware-efficient SSD (Structured State Space Duality) algorithm
- 2-8x faster than Mamba1

For embedding (bidirectional), we run Mamba2 in both directions
and combine the outputs.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from .utils import RMSNorm

try:
    from mamba_ssm import Mamba2 as Mamba2Core
    MAMBA2_AVAILABLE = True
except ImportError:
    MAMBA2_AVAILABLE = False


class Mamba2Fallback(nn.Module):
    """
    Pure PyTorch fallback when mamba_ssm is not installed.
    Implements simplified SSM: y = SSM(Conv1d(Linear(x)))
    Not as fast as CUDA kernel but functionally equivalent.
    """

    def __init__(
        self,
        d_model: int,
        d_state: int = 128,
        d_conv: int = 4,
        expand: int = 2,
        headdim: int = 64,
        ngroups: int = 8,
    ):
        super().__init__()
        self.d_model = d_model
        self.d_inner = d_model * expand
        self.d_state = d_state
        self.d_conv = d_conv
        self.headdim = headdim
        self.nheads = self.d_inner // headdim
        self.ngroups = ngroups

        assert self.d_inner % headdim == 0, (
            f"d_inner ({self.d_inner}) must be divisible by headdim ({headdim})"
        )
        assert self.nheads % ngroups == 0, (
            f"nheads ({self.nheads}) must be divisible by ngroups ({ngroups})"
        )

        # Input projection: x → (z, x_ssm)
        self.in_proj = nn.Linear(d_model, self.d_inner * 2, bias=False)

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

        # SSM parameters
        self.dt_proj = nn.Linear(self.d_inner, self.nheads, bias=True)
        # A must be negative for SSM stability: A = -exp(A_log)
        # Init A_log ~ log(uniform(1, 16)) per Mamba paper
        A_init = torch.log(torch.rand(self.nheads) * 15 + 1)  # log(U(1,16))
        self.A_log = nn.Parameter(A_init)
        self.D = nn.Parameter(torch.ones(self.nheads))

        # B, C projections (input-dependent)
        self.B_proj = nn.Linear(self.d_inner, self.ngroups * d_state, bias=False)
        self.C_proj = nn.Linear(self.d_inner, self.ngroups * d_state, bias=False)

        # Output projection
        self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)
        self.norm = nn.LayerNorm(self.d_inner)

    def _ssm_scan(self, x, dt, A, B, C, D):
        """
        Parallel Associative Scan implementation of SSD (Structured State Space Duality).
        Bypasses the slow O(T) Python loop and avoids Blackwell Triton hangs.
        """
        orig_dtype = x.dtype
        x = x.float()
        dt = dt.float()
        A = A.float()
        B = B.float()
        C = C.float()

        B_seq, T, d_inner = x.shape
        nheads = self.nheads
        headdim = self.headdim

        x = x.view(B_seq, T, nheads, headdim)
        dt = F.softplus(dt)  # (B, T, nheads)
        A = -torch.exp(A)    # (nheads,)

        # 1. Compute dA: A_bar[t] = dt[t] * A
        log_dA = dt.unsqueeze(-1) * A.view(1, 1, nheads, 1) # (B, T, H, 1)
        
        # 2. Compute cumulative product of dA using cumsum of logs
        # M[t, i] = prod_{j=i+1}^t exp(log_dA[j]) = exp(cumsum(log_dA)[t] - cumsum(log_dA)[i])
        cum_log_dA = torch.cumsum(log_dA, dim=1) # (B, T, H, 1)
        
        # clg: (B, H, T, 1)
        clg = cum_log_dA.permute(0, 2, 1, 3) 
        # diff[t, i] = cum_log_dA[t] - cum_log_dA[i] -> (B, H, T, T)
        diff = clg - clg.transpose(-1, -2)
        decay_matrix = torch.exp(diff)
        
        # Causal mask
        m_idx = torch.arange(T, device=x.device)
        mask = (m_idx.view(-1, 1) >= m_idx.view(1, -1)).to(x.dtype)
        decay_matrix = decay_matrix * mask.view(1, 1, T, T)

        # 3. Compute B_bar * x: (B, T, H, S) * (B, T, H, D) -> states
        B = B.view(B_seq, T, self.ngroups, self.d_state)
        heads_per_group = nheads // self.ngroups
        B = B.repeat_interleave(heads_per_group, dim=2)  # (B, T, H, S)
        
        # dB_x: (B, T, H, S, D)
        dB = dt.unsqueeze(-1) * B # (B, T, H, S)
        bx = torch.einsum("bths,bthd->bthsd", dB, x) 
        
        # 4. Apply decay matrix to states: h_t = sum_{i=0}^t M[t, i] * (dB_i * x_i)
        # bx_flat: (B, H, T, S*D)
        bx_flat = bx.permute(0, 2, 1, 3, 4).reshape(B_seq, nheads, T, -1)
        h_flat = torch.einsum("bhti,bhis->bhts", decay_matrix, bx_flat)
        h = h_flat.view(B_seq, nheads, T, self.d_state, headdim)
        
        # 5. Output: y_t = C_t * h_t
        C = C.view(B_seq, T, self.ngroups, self.d_state)
        C = C.repeat_interleave(heads_per_group, dim=2) # (B, T, H, S)
        
        # h_perm: (B, T, H, S, D)
        h_perm = h.permute(0, 2, 1, 3, 4)
        # y: (B, T, H, D)
        y = torch.einsum("bths,bthsd->bthd", C, h_perm)
        
        # Add D skip connection
        y = y + D.view(1, 1, nheads, 1) * x
        
        return y.reshape(B_seq, T, d_inner).to(orig_dtype)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not hasattr(self, '_warned') and x.shape[1] > 512:
            import warnings
            warnings.warn(
                f"Mamba2Fallback: sequential SSM scan with seq_len={x.shape[1]} will be slow. "
                f"Install mamba_ssm for CUDA-accelerated scan.",
                stacklevel=2,
            )
            self._warned = True

        B, T, D = x.shape

        # Input projection
        xz = self.in_proj(x)  # (B, T, 2*d_inner)
        x_ssm, z = xz.chunk(2, dim=-1)

        # Conv1d
        x_conv = x_ssm.transpose(1, 2)  # (B, d_inner, T)
        x_conv = self.conv1d(x_conv)[:, :, :T]  # trim padding
        x_conv = x_conv.transpose(1, 2)  # (B, T, d_inner)
        x_conv = F.silu(x_conv)

        # SSM parameters from input
        dt = self.dt_proj(x_conv)
        B_param = self.B_proj(x_conv)
        C_param = self.C_proj(x_conv)

        # SSM scan
        y = self._ssm_scan(x_conv, dt, self.A_log, B_param, C_param, self.D)

        # Gate with z
        y = y * F.silu(z)
        y = self.out_proj(self.norm(y))

        return y


class Mamba2Block(nn.Module):
    """
    Bidirectional Mamba2 block for embedding.
    Runs Mamba2 forward + backward, combines outputs.
    """

    def __init__(
        self,
        d_model: int,
        d_state: int = 128,
        d_conv: int = 4,
        expand: int = 2,
        headdim: int = 64,
        ngroups: int = 8,
        bidirectional: bool = True,
    ):
        super().__init__()
        self.bidirectional = bidirectional
        
        mamba_cls = Mamba2Core if MAMBA2_AVAILABLE else Mamba2Fallback
        mamba_kwargs = dict(
            d_model=d_model,
            d_state=d_state,
            d_conv=d_conv,
            expand=expand,
            headdim=headdim,
            ngroups=ngroups,
        )

        self.forward_mamba = mamba_cls(**mamba_kwargs)

        if bidirectional:
            self.backward_mamba = mamba_cls(**mamba_kwargs)
            self.merge_proj = nn.Linear(d_model * 2, d_model, bias=False)

    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
        mask = None
        if attention_mask is not None:
            mask = attention_mask.unsqueeze(-1).to(x.dtype)
            x = x * mask

        fwd_out = self.forward_mamba(x)

        if self.bidirectional:
            x_flip = x.flip(dims=[1])
            if mask is not None:
                x_flip = x_flip * mask.flip(dims=[1])
            bwd_out = self.backward_mamba(x_flip).flip(dims=[1])
            out = self.merge_proj(torch.cat([fwd_out, bwd_out], dim=-1))
        else:
            out = fwd_out

        # Re-apply mask: SSM hidden state decay để lại artifact tại padding positions
        if mask is not None:
            out = out * mask

        return out