model.py

import  math

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

from enum import Enum
from dataclasses import dataclass, field
from causal_conv1d.causal_conv1d_varlen import causal_conv1d_varlen_states
from mamba_ssm.ops.triton.selective_state_update import selective_state_update
from mamba_ssm.ops.triton.ssd_combined import mamba_split_conv1d_scan_combined

from .causal_conv1d_compilable import causal_conv1d_fn, causal_conv1d_update
from .ssm_compilable import mamba_chunk_scan_combined
from .norms import build_norm


class InitStdFactor(Enum):
    DISABLED = "disabled"            # Init std is divided by 1.0
    GLOBAL_DEPTH = "global_depth"    # Init std is divided by sqrt(2*num_layers)
    CURRENT_DEPTH = "current_depth"  # Init std is divided by sqrt(2*depth)
    DIM_RATIO = "dim_ratio"          # Init std is divided by model_dim/4096


@dataclass
class InitConfig:
    dt_max: float = 0.1
    dt_min: float = 0.001

    dt_init_floor: float = 1e-4

    A_init_min: float = 1
    A_init_max: float = 16


DEFAULT_INIT_CONFIG = InitConfig()


@dataclass
class BaseMambaConfig:
    """
    Configuration for the Mamba family of models.
    """
    dim: int = 512
    num_layers: int = 8
    num_heads: int = 8

    state_dim: int = 128
    num_groups: int = 1
    conv_size: int | None = 4

    bias: bool = False      # Linear bias
    conv_bias: bool = True  # Convolutional bias
    dt_bias: bool = False
    D_has_head_dim: bool = False
    learnable_init_states: bool = False

    ffn_dim_multiplier: float = 2.0
    multiple_of: int = 256  # Enforce that MLP hidden layer size is multiple of a large power of 2

    norm_eps: float = 1e-6
    norm_type: str = "rmsnorm"
    
    # CUDA-related items
    ssm_chunk_size: int = 256
    use_mem_eff_path: bool = False

    # Initialization-related items
    init_use_depth: bool = False
    init_base_std: float | None = None
    init_std_factor: str = "disabled"  # e.g. "global_depth"
    init_config: InitConfig = field(default_factory=InitConfig)


class SSM(nn.Module):
    """
    State Space Model (SSM) implementation with selective state updates and convolution.

    Implements the core SSM computation with support for both training and inference modes.
    During inference, uses cached states for efficient token-by-token generation.
    """
    def __init__(self, config: BaseMambaConfig) -> None:
        """Initialize SSM parameters and layers.
        Args:
            config: Configuration containing model hyperparameters
        """
        super().__init__()
        self.config = config
        vars(self).update(vars(config))

        assert self.dim > 0,        "Model dimension (config.dim) must be positive"
        assert self.num_heads > 0,  "Number of heads (config.num_heads) must be positive"
        assert self.state_dim > 0,  "State dimension (config.state_dim) must be positive"

        if self.ffn_dim_multiplier is None:
            raise ValueError(
                "ffn_dim_multiplier must be set to a valid float (e.g. 2.0) "
                "to determine hidden_dim in SSM."
            )
        assert self.ffn_dim_multiplier > 0, "ffn_dim_multiplier must be > 0"

        self.hidden_dim = int(self.ffn_dim_multiplier * self.dim)
        self.hidden_dim = config.multiple_of * ( # Round up to multiple_of
            (self.hidden_dim + self.multiple_of - 1) // self.multiple_of
        )
        
        assert self.hidden_dim % self.num_heads == 0, (
            f"Hidden dim {self.hidden_dim} not divisible by num_heads={self.num_heads}."
        )

        self.head_dim = self.hidden_dim // self.num_heads

        self.dt_limit_kwargs = {}
        dt_limit = (self.init_config.dt_min, self.init_config.dt_max)
        if dt_limit != (0.0, float("inf")):
            self.dt_limit_kwargs = dict(dt_limit=dt_limit)

        # Order: [z, x, B, C, dt]
        d_input = (
            2 * self.hidden_dim
            + 2 * self.num_groups * self.state_dim
            + self.num_heads
        )

        self.input = nn.Linear(self.dim, d_input, bias=self.bias)

        # Only create Conv1d if self.conv_size is specified
        if self.conv_size is not None:
            conv_dim = self.hidden_dim + 2 * self.num_groups * self.state_dim

            # Depthwise-ish conv (groups = out_channels)
            # TODO: Check that this is used if causal_conv1d_fn and causal_conv1d_update cannot be imported
            self.conv1d = nn.Conv1d(
                in_channels=conv_dim,
                out_channels=conv_dim,
                kernel_size=self.conv_size,
                groups=conv_dim,
                bias=self.conv_bias,  # <- This is a boolean in your config, so pass that or True/False
                padding=self.conv_size - 1  # for "causal" style
            )

        if config.dt_bias:
            self.dt_bias = nn.Parameter(torch.empty(self.num_heads))
        else:
            self.dt_bias = nn.Parameter(torch.zeros(self.num_heads), requires_grad=False)

        self.A_log = nn.Parameter(torch.empty(self.num_heads))

        if config.D_has_head_dim:
            self.D = nn.Parameter(torch.ones(self.num_heads, self.head_dim))
        else:
            self.D = nn.Parameter(torch.ones(self.num_heads))
        
        if self.learnable_init_states:
            self.init_states = nn.Parameter(torch.zeros(self.num_heads, self.head_dim, self.state_dim))

        # Can also just use nn.RMSNorm
        self.norm = build_norm(config.norm_type, dim=self.hidden_dim, eps=self.norm_eps)
        
        self.output = nn.Linear(self.hidden_dim, self.dim, bias=self.bias)

    def _causal_conv(
        self, 
        zxbcdt: torch.Tensor,
        tok_idx: torch.Tensor | None = None,
        cu_seqlens: torch.Tensor | None = None,
        ssm_impl: str = "ssm"
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        # TODO: Make slightly less verbose
        """Processes input through causal convolution path, handling both full sequence and incremental cases.

        This function implements two processing modes:
        1. Full sequence ("ssm"): Used during training and initial prompt processing.
        2. Incremental ("ssm_update"): Used during token-by-token generation.

        Args:
            zxbcdt: Input tensor containing concatenated [z, x, B, C, dt] components
            tok_idx: Token indices for sequence processing. Required for "ssm" mode.
                Defaults to None.
            cu_seqlens: Cumulative sequence lengths for variable length processing.
                Used only in "ssm" mode with caching. Defaults to None.
            ssm_impl: Implementation mode, either "ssm" for full sequence processing
                or "ssm_update" for incremental generation. Defaults to "ssm".

        Returns:
            tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
                Tuple containing separated components (z, x, B, C, dt), where:
                - z: Gating branch
                - x: Main branch
                - B, C: SSM state matrices (analogous to K, Q in attention)
                - dt: Time delta values

        Notes:
            - When using "ssm" mode during inference, a cache should be pre-initialized
            externally. This design allows for flexible caching strategies without
            modifying model code.
            - The "ssm_update" mode requires a cache to exist and will use it for
            incremental state updates during generation.
            - B, C components correspond to Key, Query in the SSM/attention duality.
        """
        # Split input into components
        z, xBC, dt = torch.split(
            zxbcdt,
            [
                self.hidden_dim,
                self.hidden_dim + 2 * self.num_groups * self.state_dim,
                self.num_heads,
            ],
            dim=-1,
        )

        if ssm_impl == "ssm":
            if hasattr(self, "cache"):
                conv_varlen_states = causal_conv1d_varlen_states(
                    xBC.squeeze(0),
                    cu_seqlens,
                    state_len=self.cache.conv_cache.shape[-1],
                )
                self.cache.conv_cache.copy_(conv_varlen_states)

            xBC = causal_conv1d_fn(
                x=xBC.transpose(1, 2),
                weight=self.conv1d.weight.squeeze(1),
                bias=self.conv1d.bias,
                activation="silu",
                seq_idx=tok_idx,
            ).transpose(1, 2)
        elif ssm_impl == "ssm_update":
            xBC = causal_conv1d_update(
                x=xBC.squeeze(0),
                conv_state=self.cache.conv_cache,
                weight=self.conv1d.weight.squeeze(1),
                bias=self.conv1d.bias,
                activation="silu",
            ).unsqueeze(0)
        else:
            raise NotImplementedError(f"SSM implementation {ssm_impl} not supported")

        # Split processed tensor into components
        x, B, C = torch.split(
            xBC,
            [
                self.hidden_dim,
                self.num_groups * self.state_dim,
                self.num_groups * self.state_dim,
            ],
            dim=-1,
        )

        return z, x, B, C, dt

    def _non_causal_conv(self, zxbcdt: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        z, x, B, C, dt = torch.split(
            zxbcdt,
            [
                self.hidden_dim,
                self.hidden_dim,
                self.num_groups * self.state_dim,
                self.num_groups * self.state_dim,
                self.num_heads,
            ],
            dim=-1,
        )
        return z, x, B, C, dt

    def _fwd(self, x, dt, A, B, C, tok_idx, cu_seqlens, initial_states):
        """
        For training
        
        Returns:
            (bsz, seq_len, num_heads, head_dim)
        """
        y = mamba_chunk_scan_combined(
            x,
            dt,
            A,
            B,
            C,
            dt_bias=self.dt_bias,
            dt_softplus=True,
            chunk_size=self.ssm_chunk_size,
            D=self.D,
            z=None,
            seq_idx=tok_idx,
            cu_seqlens=cu_seqlens,
            initial_states=initial_states,
            **self.dt_limit_kwargs,
        )

        if hasattr(self, "cache"):
            y, varlen_states = y
            self.cache.state_cache.copy_(varlen_states)

        return y
    
    def _step(self, x, seq_len, dt, A, B, C):
        """
        For inference / generation.
        """
        x = x.squeeze(0)
        A = A[..., None, None].expand(self.num_heads, self.head_dim, self.state_dim)
        dt = dt.permute(1, 2, 0).expand(seq_len, self.num_heads, self.head_dim)
        D = self.D
        if D is not None and D.dim() == 1:
            D = D.unsqueeze(1).expand(self.num_heads, self.head_dim)
        B, C = B.squeeze(0), C.squeeze(0)
        y = selective_state_update(
            self.cache.state_cache,
            x,
            dt,
            A,
            B,
            C,
            D,
            z=None,
            dt_bias=(
                torch.zeros(self.num_heads, self.head_dim).to(x)
                if self.dt_bias is None
                else self.dt_bias.unsqueeze(1).expand(self.num_heads, self.head_dim)
            ),
            dt_softplus=True,
        ).unsqueeze(0)
        
        return y

    def forward(
        self,
        x: torch.Tensor,
        tok_idx: torch.Tensor | None = None,
        cu_seqlens: torch.Tensor | None = None,
        ssm_impl: str = "ssm",
    ) -> torch.Tensor:
        bsz, seq_len, _ = x.shape

        zxbcdt = self.input(x)

        A = -torch.exp(self.A_log.float())
        initial_states = (
            self.init_states.expand(bsz, -1, -1, -1)
            if self.learnable_init_states else None
        )

        # Causal conv path
        if self.conv_size is not None:
            
            # Memory-efficient Triton kernel path
            if self.use_mem_eff_path:
                out = mamba_split_conv1d_scan_combined(
                    zxbcdt,
                    self.conv1d.weight.squeeze(1),
                    self.conv1d.bias,
                    self.dt_bias,
                    A,
                    D=self.D,
                    chunk_size=self.ssm_chunk_size,
                    seq_idx=tok_idx,
                    activation="silu",
                    rmsnorm_weight=self.norm.weight,
                    rmsnorm_eps=self.norm.eps,
                    outproj_weight=self.output.weight,
                    outproj_bias=self.output.bias,
                    headdim=self.head_dim,
                    ngroups=self.num_groups,
                    norm_before_gate=False, # Post-norm, y = self.norm(y * F.silu(z))
                    initial_states=initial_states,
                    **self.dt_limit_kwargs,
                )
                return out
            else:
                # CUDA kernel path
                z, x, B, C, dt = self._causal_conv(zxbcdt)
        else:
            # Non-causal conv path
            z, x, B, C, dt = self._non_causal_conv(zxbcdt)

        x = x.view(bsz, seq_len, self.num_heads, self.head_dim)
        B = B.view(bsz, seq_len, self.num_groups, self.state_dim)
        C = C.view(bsz, seq_len, self.num_groups, self.state_dim)

        # Chunked SSM scan
        if ssm_impl == "ssm":
            # (bsz, seq_len, num_heads, head_dim)
            y = self._fwd(x, dt, A, B, C, tok_idx, cu_seqlens, initial_states)
        elif ssm_impl == "ssm_update":
            y = self._step(x, seq_len, dt, A, B, C)
        else:
            raise NotImplementedError(f"SSM implementation {ssm_impl} not supported")

        y = y.view(bsz, seq_len, self.hidden_dim)

        # Could be different activation function, including None.
        # Mamba people post_norm here also (sometimes norm(z)*y or norm(z*y))
        # y = self.norm(y) * F.silu(z)
        y = self.norm(y * F.silu(z))
        out = self.output(y)

        return out

    @torch.inference_mode()
    def reset_parameters(self, init_std, factor) -> None:
        config = self.config
        init_config = config.init_config
        if init_config is None:
            init_config = DEFAULT_INIT_CONFIG

        # Linear layers
        in_init_std = init_std or (self.dim ** (-0.5))
        out_init_std = init_std or (self.hidden_dim ** (-0.5))
        out_init_std = out_init_std / factor

        nn.init.trunc_normal_(
            self.input.weight,
            mean=0.0,
            std=in_init_std,
            a=-3 * in_init_std,
            b=3 * in_init_std,
        )

        nn.init.trunc_normal_(
            self.output.weight,
            mean=0.0,
            std=out_init_std,
            a=-3 * out_init_std,
            b=3 * out_init_std,
        )

        # SSM
        if self.dt_bias is not None and self.dt_bias.requires_grad:
            log_dt_min = math.log(init_config.dt_min)
            log_dt_max = math.log(init_config.dt_max)
            
            # Sample log_dt ~ Uniform[log_dt_min, log_dt_max]
            log_dt = torch.rand(self.num_heads, device=self.dt_bias.device) * (log_dt_max - log_dt_min) + log_dt_min
            dt = torch.exp(log_dt)
            dt = torch.clamp(dt, min=init_config.dt_init_floor)

            # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
            inv_dt = dt + torch.log(-torch.expm1(-dt))
            self.dt_bias.copy_(inv_dt)

        elif self.dt_bias is not None:
            # If dt_bias is not trainable, we can just keep it zero or set to any constant
            self.dt_bias.fill_(0.0)

        # Convolution
        if self.conv_size is not None:
            conv_std = init_std or (self.conv_size ** (-0.5))
            nn.init.trunc_normal_(
                self.conv1d.weight,
                mean=0.0,
                std=conv_std,
                a=-3 * conv_std,
                b=3 * conv_std,
            )
            if self.conv1d.bias is not None:
                nn.init.zeros_(self.conv1d.bias)

        # Learnable init states
        if self.learnable_init_states:
            self.init_states.zero_()

        # Initialize A_log ~ log( Uniform(A_init_min, A_init_max) )
        self.A_log.uniform_(init_config.A_init_min, init_config.A_init_max)
        self.A_log.log_()

        if self.D is not None:
            self.D.data.fill_(1.0)

        # Reset norm parameters
        self.norm.reset_parameters()


class MambaBlock(nn.Module):
    def __init__(self, config: BaseMambaConfig):
        super().__init__()
        self.norm = build_norm(config.norm_type, dim=config.dim, eps=config.norm_eps)
        self.ssm = SSM(config)

    def forward(
        self,
        x: torch.Tensor,
        tok_idx: torch.Tensor | None,
        cu_seqlens: torch.Tensor | None,
        ssm_impl: str = "ssm",
    ) -> torch.Tensor:
        x = x + self.ssm(self.norm(x), tok_idx=tok_idx, cu_seqlens=cu_seqlens, ssm_impl=ssm_impl)
        return x

    @torch.inference_mode()
    def init_weights(self, init_std=None, factor=1.0):
        self.norm.reset_parameters()
        self.ssm.reset_parameters(init_std, factor)


class BaseMamba(nn.Module):
    def __init__(self, config: BaseMambaConfig):
        super().__init__()
        self.model_dim = config.dim
        self.init_base_std = config.init_base_std

        self.init_config = config.init_config
        self.init_std_factor = InitStdFactor(config.init_std_factor)

        self.layers = nn.ModuleList()
        for _ in range(config.num_layers):
            self.layers.append(MambaBlock(config))

    def forward(
        self,
        h: torch.Tensor,
        tok_idx: torch.Tensor | None,
        cu_seqlens: torch.Tensor | None,
        ssm_impl: str = "ssm",
    ) -> torch.Tensor:
        for layer in self.layers:
            h = layer(h, tok_idx=tok_idx, cu_seqlens=cu_seqlens, ssm_impl=ssm_impl)
        return h

    @torch.inference_mode()
    def reset_parameters(self):
        pass

    @torch.inference_mode()
    def init_weights(self):
        self.reset_parameters()
        for depth, layer in enumerate(self.layers):
            factor = {
                InitStdFactor.CURRENT_DEPTH: (2 * (depth + 1)) ** 0.5,
                InitStdFactor.GLOBAL_DEPTH: (2 * (len(self.layers) + 1)) ** 0.5,
                InitStdFactor.DIM_RATIO: self.model_dim / 4096,
                InitStdFactor.DISABLED: 1.0,
            }[self.init_std_factor]

            layer.init_weights(self.init_base_std, factor)


@dataclass
class Mamba2Config(BaseMambaConfig):
    seed: int = 1337

    vocab_size: int = -1 # Will error if unchanged, makes you double check!
    weight_tying: bool = False
    torch_dtype: torch.dtype = torch.bfloat16

    loss_reduction: str = "mean"

    use_attn: bool = False
    softcap: float = 50.0


class Mamba2(BaseMamba):
    def __init__(self, config: Mamba2Config) -> None:
        super().__init__(config)
        self.weight_tying = config.weight_tying
        self.loss_reduction = config.loss_reduction

        assert config.vocab_size > 0, "vocab_size must be set and > 0"

        self.tok_emb = torch.nn.Embedding(config.vocab_size, config.dim)

        self.norm = nn.RMSNorm(config.dim, eps=config.norm_eps)

        self.output = nn.Linear(
            config.dim,
            config.vocab_size,
            bias=False,
        )

        if config.weight_tying:
            self.output.weight = self.tok_emb.weight

        print("Model Parameter Count: %.2fM\n" % (self._get_num_params() / 1e6,))

    def _get_num_params(self):
        n_params = sum(p.numel() for p in self.parameters())
        if hasattr(self, "pos_emb") and self.pos_emb is not None:
            n_params -= self.pos_emb.weight.numel()
        if self.tok_emb.weight is not self.output.weight:
            n_params -= self.tok_emb.weight.numel()
        return n_params

    def forward(
        self,
        x: torch.Tensor,
        target: torch.Tensor | None = None,
        tok_idx: torch.Tensor | None = None,
        cu_seqlens: torch.Tensor | None = None,
        ssm_impl: str = "ssm",
    ) -> torch.Tensor:
        h = self.tok_emb(x)
        h = super().forward(h, tok_idx=tok_idx, cu_seqlens=cu_seqlens, ssm_impl=ssm_impl)
        logits = self.output(self.norm(h))
        return logits

    @torch.inference_mode()
    def reset_parameters(self, init_std=None):
        # Either use fixed base std or sqrt model dim
        super().reset_parameters()
        init_std = init_std or (self.model_dim ** (-0.5))
        self.norm.reset_parameters()
        nn.init.trunc_normal_(
            self.tok_emb.weight,
            mean=0.0,
            std=init_std,
            a=-3 * init_std,
            b=3 * init_std,
        )
        if not self.weight_tying:
            nn.init.trunc_normal_(
                self.output.weight,
                mean=0.0,
                std=init_std,
                a=-3 * init_std,
                b=3 * init_std,
            )

    @torch.inference_mode()
    def init_weights(self, buffer_device: torch.device = None):
        """
        Initialize model parameters and optionally compute buffers on a specific device.
        
        Args:
            buffer_device (torch.device, optional): If provided, any large or precomputed
                buffers (like RoPE frequency tensors) will be allocated or re-created on
                this device during initialization. This can avoid overhead from transferring
                buffers between CPU and GPU after creation. If None, buffers default to the
                device of the first parameter or CPU.

        Usage:
            - Pass a GPU device (e.g., ``torch.device('cuda')``) when you want to ensure
            buffers are created directly on GPU, preventing extra transfers.
            - Pass a CPU device (e.g., ``torch.device('cpu')``) if you want to keep
            large buffers in CPU memory (common in CPU-offload or pipeline-parallel setups).
            - Leave it as ``None`` to rely on the model’s existing parameter device or
            the default PyTorch device context.

        When / Why:
            - Useful in distributed or pipeline-parallel training where parameters may
            initially live on CPU, but you still need certain buffers on GPU to avoid
            overhead during forward passes.
            - Prevents large re-allocations or re-copies when big buffers (like RoPE
            frequency tables) are needed per rank.
        """
        super().init_weights()

    @classmethod
    def from_model_args(cls, config: Mamba2Config) -> "Mamba2":
        """
        Initialize a Mamba model from a MambaConfig object.

        Args:
            config (MambaConfig): Mamba configuration arguments.

        Returns:
            Mamba: Mamba-2 model.
        """
        return cls(config)


def get_mamba2_flops(
    seq_len: int,
    dim: int,
    num_layers: int,
    vocab_size: int,
    ffn_multiplier: float = 2.0,
    state_dim: int = 128,
    conv_size: int = 4,
    num_heads: int = 8,
    num_groups: int = 1,
    multiple_of: int = 256,
    include_input_embedding: bool = True,
    include_output_logits: bool = True,
    forward_backward_multiplier: float = 1.0,
) -> int:
    """
    Estimate the FLOPs for a Mamba-2 style model using a "Chinchilla-like" shape-based approach.

    By default, this returns the forward-pass cost. If you want a rough
    forward+backward estimate, set `forward_backward_multiplier=3.0` (common
    rule-of-thumb for these models).

    What gets counted:
    • Hidden dimension is rounded up to 'multiple_of' = 256 (as in Mamba).
    • Per-layer:
        1) Input Linear: [dim → 2*hidden_dim + 2*(groups*state_dim) + num_heads]
        2) Depthwise Conv1D: 2*(conv_dim * conv_size), where conv_dim=hidden_dim + 2*groups*state_dim
        3) SSM selective scan: ~9*(dim*state_dim) (from Mamba dev discussion)
        4) Output Linear: [hidden_dim → dim]
    • Each layer’s cost is multiplied by (seq_len * num_layers).
    • Optionally adds:
        - The cost of the input embedding (treating it as a matmul: seq_len×vocab_size × vocab_size×dim).
        - The cost of the final projection [dim → vocab_size].
    • Finally scaled by `forward_backward_multiplier` if desired.

    Args:
        seq_len (int): Sequence length (number of tokens).
        dim (int): Model (embedding) dimension.
        num_layers (int): Number of Mamba layers.
        vocab_size (int): Vocabulary size for final logits projection.
        ffn_multiplier (float): FFN expansion ratio, e.g. 2.0 => hidden_dim=2×dim (rounded up).
        state_dim (int): SSM state dimension (commonly 128).
        conv_size (int): Kernel size for the depthwise conv1d (default=4).
        num_heads (int): Number of heads (slightly affects input-lin out_dim).
        num_groups (int): For "grouped" states in some Mamba variants (usually 1).
        multiple_of (int): Round hidden_dim up to this multiple (commonly 256).
        include_input_embedding (bool): If True, count the cost of an “embedding matmul”
                                        for the input tokens => shape-based approach.
        include_output_logits (bool): If True, count the cost of final [dim → vocab_size].
        forward_backward_multiplier (float): E.g. 1.0 for forward only, 2.0 or 3.0 for forward+backward.

    Returns:
        int: Approximate total FLOPs (multiply-adds) for the selected pass(es),
            as an integer.
    """
    # 0) Input embedding (optional)
    flops_embedding = 0
    if include_input_embedding:
        flops_embedding = 2 * (seq_len * vocab_size * dim)

    # 1) Round up hidden_dim
    raw_hidden_dim = int(ffn_multiplier * dim)
    hidden_dim = multiple_of * ((raw_hidden_dim + multiple_of - 1) // multiple_of)

    # 2) Per-layer forward cost
    out_dim_input = 2*hidden_dim + 2*(num_groups*state_dim) + num_heads
    flops_input_linear = 2 * (dim * out_dim_input)
    conv_dim = hidden_dim + 2*(num_groups*state_dim)
    flops_conv = 2 * (conv_dim * conv_size)
    flops_ssm = 9 * state_dim * dim
    flops_output_linear = 2 * (hidden_dim * dim)
    flops_layer = (flops_input_linear + flops_conv + flops_ssm + flops_output_linear)

    # Multiply by #layers and sequence length
    flops_layers = flops_layer * num_layers * seq_len

    # 3) Final projection [dim → vocab_size] (optional)
    flops_vocab = 0
    if include_output_logits:
        flops_vocab = 2 * (seq_len * dim * vocab_size)

    # 4) Total forward FLOPs
    flops_forward = flops_embedding + flops_layers + flops_vocab

    # 5) Scale for forward+backward if desired
    return int(flops_forward * forward_backward_multiplier)

def get_mamba2_flops_per_token(
    **kwargs
) -> float:
    """
    Estimate FLOPs per token for a Mamba-2 style model.

    This function extracts necessary parameters from kwargs and calculates the FLOPs per token.
    
    Args:
        **kwargs: Dictionary containing model configuration parameters.
    
    Returns:
        float: Approximate FLOPs per token.
    """
    defaults = {
        'ffn_dim_multiplier': 2.0,
        'state_dim': 128,
        'conv_size': 4,
        'num_heads': 8,
        'num_groups': 1,
        'multiple_of': 256,
        'include_input_embedding': True,
        'include_output_logits': True,
        'forward_backward_multiplier': 1.0,
    }
    # Merge defaults
    for k, v in defaults.items():
        kwargs.setdefault(k, v)
    # Mandatory keys
    for required in ['seq_len', 'dim', 'num_layers', 'vocab_size']:
        if required not in kwargs:
            raise ValueError(f"Missing required parameter: {required}")

    total_flops = get_mamba2_flops(
        seq_len=kwargs['seq_len'],
        dim=kwargs['dim'],
        num_layers=kwargs['num_layers'],
        vocab_size=kwargs['vocab_size'],
        ffn_multiplier=kwargs['ffn_dim_multiplier'],
        state_dim=kwargs['state_dim'],
        conv_size=kwargs['conv_size'],
        num_heads=kwargs['num_heads'],
        num_groups=kwargs['num_groups'],
        multiple_of=kwargs['multiple_of'],
        include_input_embedding=kwargs['include_input_embedding'],
        include_output_logits=kwargs['include_output_logits'],
        forward_backward_multiplier=kwargs['forward_backward_multiplier'],
    )
    flops_per_token = total_flops / kwargs['seq_len']

    return flops_per_token


# Optional policy for activation checkpointing. With None, we stick to the default (defined distributed.py: default_no_recompute_ops)
def get_no_recompute_ops():
    return {
        torch.ops.aten.mm.default,
        torch.ops.aten._scaled_mm.default,
        torch.ops.c10d_functional.reduce_scatter_tensor.default,
        torch.ops.mamba_ssm.ssm_chunk_scan_combined_fwd.default,

        # For low-precision training, it's useful to always save the result of max(abs(tensor))
        torch.ops.aten.abs.default,
        torch.ops.aten.max.default,
    }


def main():
    from mamba_ssm import Mamba2 as MambaRef

    x = torch.randn(2, 64, 192).cuda()

    # Create and run the first model
    model = MambaRef(
        d_model=192,
        expand=2,
        d_conv=4,
        d_state=64,
        headdim=48,
    ).cuda()
    y = model(x)
    print("Mamba reference output: ", y)
    print("Mean of MambaRef output: ", y.mean().item())
    print("Stddev of MambaRef output: ", y.std().item())

    # Create and run the second model
    config = Mamba2Config(vocab_size=200064, use_mem_eff_path=True)
    model2 = Mamba2(
        config=config,
    ).cuda()

    # Fix: Convert x to torch.LongTensor
    x_indices = torch.randint(0, config.vocab_size, (2, 64), dtype=torch.long).cuda()

    y2 = model2(x_indices)
    print("Mamba output: ", y2)
    print("Mean of Mamba output: ", y2.mean().item())
    print("Stddev of Mamba output: ", y2.std().item())

if __name__ == "__main__":
    main()