added configs

__init__.py

@@ -0,0 +1,2 @@
+from .configuration_minimamba import MiniMambaConfig
+from .modeling_minimamba import MiniMamba

added_tokens.json

@@ -0,0 +1,3 @@
+{
+    "<|endofprompt|>": 200018
+}

causal_conv1d_compilable.py

@@ -0,0 +1,215 @@
+from typing import Optional, Tuple
+import torch
+import causal_conv1d_cuda
+
+# Causal Conv1D Forward Function
+@torch.library.custom_op(
+    "mamba_causal_conv1d::causal_conv1d_fwd",
+    mutates_args=(),
+    device_types="cuda",
+)
+def causal_conv1d_fwd(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    seq_idx: Optional[torch.Tensor] = None,
+    activation: Optional[str] = None,
+) -> torch.Tensor:
+    # Ensure activation is valid
+    if activation not in [None, "silu", "swish"]:
+        raise NotImplementedError("activation must be None, silu, or swish")
+
+    # Ensure x is contiguous
+    if x.stride(2) != 1 and x.stride(1) != 1:
+        x = x.contiguous()
+
+    # Make bias and seq_idx contiguous if they exist
+    bias = bias.contiguous() if bias is not None else None
+    seq_idx = seq_idx.contiguous() if seq_idx is not None else None
+
+    # Translate activation to bool for custom CUDA kernel
+    use_activation = activation in ["silu", "swish"]
+
+    # Call custom CUDA kernel for forward pass
+    out = causal_conv1d_cuda.causal_conv1d_fwd(
+        x, weight, bias, seq_idx, None, None, use_activation
+    )
+    return out
+
+# Register a fake forward pass for tracing
+@causal_conv1d_fwd.register_fake
+def _causal_conv1d_fwd_fake(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    seq_idx: Optional[torch.Tensor] = None,
+    activation: Optional[str] = None,
+) -> torch.Tensor:
+    torch._check(x.shape[-2] == weight.shape[0])
+    return torch.empty_like(x)
+
+# Causal Conv1D Backward Function
+@torch.library.custom_op(
+    "mamba_causal_conv1d::causal_conv1d_bwd", 
+    mutates_args=(),
+    device_types="cuda",
+)
+def causal_conv1d_bwd(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    dout: torch.Tensor,
+    seq_idx: Optional[torch.Tensor],
+    activation: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    # Ensure dout is contiguous
+    if dout.stride(2) != 1 and dout.stride(1) != 1:
+        dout = dout.contiguous()
+
+    # Call custom CUDA kernel for backward pass
+    dx, dweight, dbias, _ = causal_conv1d_cuda.causal_conv1d_bwd(
+        x, weight, bias, dout, seq_idx, None, None, None, False, activation
+    )
+
+    # Handle optional bias gradient
+    dbias = dbias if bias is not None else torch.empty((0,), device=dout.device)
+    
+    return dx, dweight, dbias
+
+# Register a fake backward pass for tracing
+@causal_conv1d_bwd.register_fake
+def _causal_conv1d_bwd_fake(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    dout: torch.Tensor,
+    seq_idx: Optional[torch.Tensor],
+    activation: bool,
+):
+    return (
+        torch.empty_like(x),
+        torch.empty_like(weight),
+        torch.empty_like(bias) if bias is not None else None,
+    )
+
+# Setup context for autograd
+def causal_conv1d_setup_context(ctx, inputs, output):
+    x, weight, bias, seq_idx, activation = inputs
+    ctx.activation = activation in ["silu", "swish"]
+    ctx.save_for_backward(x, weight, bias, seq_idx)
+
+# Bridge for backward pass in autograd
+def causal_conv1d_bwd_bridge(ctx, dout):
+    x, weight, bias, seq_idx = ctx.saved_tensors
+    dx, dweight, dbias = causal_conv1d_bwd(x, weight, bias, dout, seq_idx, ctx.activation)
+    
+    # Handle None return values
+    dbias = dbias if bias is not None else None
+    return dx, dweight, dbias, None, None
+
+# Register custom autograd function
+torch.library.register_autograd(
+    "mamba_causal_conv1d::causal_conv1d_fwd",
+    causal_conv1d_bwd_bridge,
+    setup_context=causal_conv1d_setup_context,
+)
+
+# Define a higher-level function to invoke the custom op
+def causal_conv1d_fn(x, weight, bias=None, seq_idx=None, activation=None):
+    return causal_conv1d_fwd(x, weight, bias, seq_idx, activation)
+
+
+@torch.library.custom_op(
+    "mamba_causal_conv1d::causal_conv1d_update",
+    mutates_args=(),
+    device_types="cuda",
+)
+def causal_conv1d_update_fwd(
+    x: torch.Tensor, 
+    conv_state: torch.Tensor,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    activation: Optional[str] = None,
+    cache_seqlens: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    x: (batch, dim) or (batch, dim, seqlen)
+    conv_state: (batch, dim, state_len), where state_len >= width - 1
+    weight: (dim, width)
+    bias: (dim,)
+    cache_seqlens: (batch,), dtype int32.
+        If not None, the conv_state is treated as a circular buffer.
+        The conv_state will be updated by copying x to the conv_state starting at the index
+        @cache_seqlens % state_len.
+
+    out: (batch, dim) or (batch, dim, seqlen)
+    """
+    if activation not in [None, "silu", "swish"]:
+        raise NotImplementedError("activation must be None, silu, or swish")
+    activation = activation in ["silu", "swish"]
+    unsqueeze = x.dim() == 2
+    if unsqueeze:
+        x = x.unsqueeze(-1)
+    out = causal_conv1d_cuda.causal_conv1d_update(
+        x, conv_state, weight, bias, activation, cache_seqlens
+    )
+    if unsqueeze:
+        out = out.squeeze(-1)
+    return out
+
+@causal_conv1d_update_fwd.register_fake
+def _causal_conv1d_update_fwd(
+    x: torch.Tensor, 
+    conv_state: torch.Tensor,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    activation: Optional[str] = None,
+    cache_seqlens: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    return torch.empty_like(x)
+
+def causal_conv1d_update(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None):
+    return causal_conv1d_update_fwd(x, conv_state, weight, bias, activation, cache_seqlens)
+
+# Test the implementation
+if __name__ == "__main__":
+    from causal_conv1d import causal_conv1d_fn as causal_conv1d_fn_ref
+
+    torch.manual_seed(0)
+
+    x = torch.randn(8, 32, 16, device="cuda", requires_grad=True)
+    weight = torch.randn(32, 3, device="cuda", requires_grad=True)
+    bias = None#torch.randn(32, device="cuda", requires_grad=True)
+
+    # Test the forward and backward pass
+    print("Custom Implementation")
+    out = causal_conv1d_fn(x, weight, bias, activation="silu")
+    out.sum().backward()
+
+    print(out.min(), out.max(), out.mean(), out.std())
+    print(x.grad.min(), x.grad.max(), x.grad.mean(), x.grad.std())
+    print(weight.grad.min(), weight.grad.max(), weight.grad.mean(), weight.grad.std())
+
+    # Try compiling the function using torch.compile
+    x.grad.zero_(), weight.grad.zero_()
+    compiled_conv1d = torch.compile(causal_conv1d_fn)
+    print(compiled_conv1d)
+
+    # Run the compiled function
+    print("Compiled Implementation")
+    out = compiled_conv1d(x, weight, bias, activation="silu")
+    out.sum().backward()
+
+    print(out.min(), out.max(), out.mean(), out.std())
+    print(x.grad.min(), x.grad.max(), x.grad.mean(), x.grad.std())
+    print(weight.grad.min(), weight.grad.max(), weight.grad.mean(), weight.grad.std())
+
+    print("Reference Implementation")
+    x.grad.zero_(), weight.grad.zero_()
+    out = causal_conv1d_fn_ref(x, weight, bias, activation="silu")
+    out.sum().backward()
+
+    print(out.min(), out.max(), out.mean(), out.std())
+    print(x.grad.min(), x.grad.max(), x.grad.mean(), x.grad.std())
+    print(weight.grad.min(), weight.grad.max(), weight.grad.mean(), weight.grad.std())
+

config.json

@@ -0,0 +1,72 @@
+{
+    "model_type": "minimamba",
+    "_name_or_path": "Mamba_500M",
+    "architectures": ["MiniMamba"],
+    "dim": 1024,
+    "num_layers": 54,
+    "num_heads": 32,
+    "state_dim": 128,
+    "num_groups": 1,
+    "conv_size": 4,
+    "use_mem_eff_path": true,
+    "dt_bias": true,
+    "D_has_head_dim": true,
+    "learnable_init_states": false,
+    "ssm_chunk_size": 256,
+    "vocab_size": 200064,
+    "ffn_dim_multiplier": 2.0,
+    "multiple_of": 256,
+    "norm_eps": 1e-05,
+    "init_use_depth": false,
+    "init_base_std": null,
+    "init_std_factor": "disabled",
+    "hidden_act": "silu",
+    "bias": false,
+    "torch_dtype": "bfloat16",
+    "seed": 1337,
+    "init_args": {
+      "dt_max": 0.1,
+      "dt_min": 0.001,
+      "dt_init_floor": 0.0001,
+      "A_init_min": 0.01,
+      "A_init_max": 16
+    },
+    "seq_len": 8192,
+    "weight_tying": true,
+    "dropout": 0.0,
+    "num_epochs": 1,
+    "global_bsz": 524288,
+    "bsz": 1,
+    "warmup_steps": 1907,
+    "eval_period": 50,
+    "save_period": 500,
+    "max_lr": 0.0003,
+    "min_lr": 3e-05,
+    "max_norm": 1.0,
+    "dilation": 1,
+    "fsdp": true,
+    "ddp": false,
+    "mixed_precision": true,
+    "cpu_offload": false,
+    "sharding_strategy": "full_shard",
+    "state_dict_type": "full",
+    "auto_wrap_policy": "partial",
+    "backward_prefetch": "backward_pre",
+    "forward_prefetch": false,
+    "sync_module_states": true,
+    "use_orig_params": true,
+    "device_id": null,
+    "precision": {
+      "param": "bfloat16",
+      "reduce": "bfloat16",
+      "buffer": "bfloat16"
+    },
+    "fsdp_modules": [
+      "MambaBlock"
+    ],
+    "use_activation_checkpointing": true,
+    "use_attn": false,
+    "softcap": 50.0,
+    "torch_compile": true
+  }
+  

configuration_minimamba.py

@@ -0,0 +1,155 @@
+from transformers import PretrainedConfig
+
+class MiniMambaConfig(PretrainedConfig):
+    """
+    Minimal or extended config class for MiniMamba.
+    Inherits from HF's PretrainedConfig so we can do:
+      model = MiniMamba.from_pretrained(...)
+    and it will load this config automatically.
+
+    This config includes all fields from the provided config.json.
+    """
+    model_type = "minimamba"
+
+    def __init__(
+        self,
+        # Standard HF fields:
+        model_type="minimamba",
+        _name_or_path="Mamba_500M",
+        architectures=["MiniMamba"],
+
+        # Key Mamba architecture hyperparameters:
+        dim=1024,
+        num_layers=54,
+        num_heads=32,
+        state_dim=128,
+        num_groups=1,
+        conv_size=4,
+        use_mem_eff_path=True,
+        dt_bias=True,
+        D_has_head_dim=True,
+        learnable_init_states=False,
+        ssm_chunk_size=256,
+        vocab_size=200064,
+        ffn_dim_multiplier=2.0,
+        multiple_of=256,
+        norm_eps=1e-5,
+        init_use_depth=False,
+        init_base_std=None,
+        init_std_factor="disabled",
+        hidden_act="silu",
+        bias=False,
+
+        # Torch / training:
+        torch_dtype="bfloat16",
+        seed=1337,
+
+        # The init_config block nested in JSON:
+        init_args=None,  # e.g. dict with dt_max, dt_min, dt_init_floor, ...
+        
+        # Additional Mamba or training fields:
+        seq_len=8192,
+        weight_tying=True,
+        dropout=0.0,
+        num_epochs=1,
+        global_bsz=524288,
+        bsz=1,
+        warmup_steps=1907,
+        eval_period=50,
+        save_period=500,
+        max_lr=0.0003,
+        min_lr=3e-5,
+        max_norm=1.0,
+        dilation=1,
+        fsdp=True,
+        ddp=False,
+        mixed_precision=True,
+        cpu_offload=False,
+        sharding_strategy="full_shard",
+        state_dict_type="full",
+        auto_wrap_policy="partial",
+        backward_prefetch="backward_pre",
+        forward_prefetch=False,
+        sync_module_states=True,
+        use_orig_params=True,
+        device_id=None,
+        precision=None,   # e.g. dict with param="bfloat16", reduce="bfloat16", buffer="bfloat16"
+        fsdp_modules=None,# e.g. ["MambaBlock"]
+        use_activation_checkpointing=True,
+        use_attn=False,
+        softcap=50.0,
+        torch_compile=True,
+
+        # Now accept arbitrary additional kwargs, to remain flexible:
+        **kwargs
+    ):
+        super().__init__(
+            # In HF, these common keys are typically passed to the parent:
+            model_type=model_type,
+            _name_or_path=_name_or_path,
+            architectures=architectures,
+            **kwargs
+        )
+
+        self.dim = dim
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.state_dim = state_dim
+        self.num_groups = num_groups
+        self.conv_size = conv_size
+        self.use_mem_eff_path = use_mem_eff_path
+        self.dt_bias = dt_bias
+        self.D_has_head_dim = D_has_head_dim
+        self.learnable_init_states = learnable_init_states
+        self.ssm_chunk_size = ssm_chunk_size
+        self.vocab_size = vocab_size
+        self.ffn_dim_multiplier = ffn_dim_multiplier
+        self.multiple_of = multiple_of
+        self.norm_eps = norm_eps
+        self.init_use_depth = init_use_depth
+        self.init_base_std = init_base_std
+        self.init_std_factor = init_std_factor
+        self.hidden_act = hidden_act
+        self.bias = bias
+
+        self.torch_dtype = torch_dtype
+        self.seed = seed
+
+        # Nested init_args (dt_max, dt_min, etc.).
+        # Could store it as a dict, or parse out the fields individually:
+        self.init_args = init_args or {}
+
+        self.seq_len = seq_len
+        self.weight_tying = weight_tying
+        self.dropout = dropout
+        self.num_epochs = num_epochs
+        self.global_bsz = global_bsz
+        self.bsz = bsz
+        self.warmup_steps = warmup_steps
+        self.eval_period = eval_period
+        self.save_period = save_period
+        self.max_lr = max_lr
+        self.min_lr = min_lr
+        self.max_norm = max_norm
+        self.dilation = dilation
+        self.fsdp = fsdp
+        self.ddp = ddp
+        self.mixed_precision = mixed_precision
+        self.cpu_offload = cpu_offload
+        self.sharding_strategy = sharding_strategy
+        self.state_dict_type = state_dict_type
+        self.auto_wrap_policy = auto_wrap_policy
+        self.backward_prefetch = backward_prefetch
+        self.forward_prefetch = forward_prefetch
+        self.sync_module_states = sync_module_states
+        self.use_orig_params = use_orig_params
+        self.device_id = device_id
+        self.precision = precision
+        self.fsdp_modules = fsdp_modules
+        self.use_activation_checkpointing = use_activation_checkpointing
+        self.use_attn = use_attn
+        self.softcap = softcap
+        self.torch_compile = torch_compile
+
+        # If you want to store any leftover kwargs:
+        self.extra_args = kwargs

model.py

@@ -0,0 +1,848 @@
+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()
+

modeling_minimamba.py

@@ -0,0 +1,222 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutput
+
+from .configuration_minimamba import MiniMambaConfig
+from .model import Mamba2, Mamba2Config
+
+
+
+class MiniMamba(PreTrainedModel):
+    """
+    A Hugging Face–style wrapper around a Mamba2 model, providing:
+      • forward(...) returning a CausalLMOutput
+      • support for HF training loops
+      • a naive generate(...) method with top-k/top-p sampling
+    """
+    config_class = MiniMambaConfig  # Tells HF which config class to use
+
+    def __init__(self, config: MiniMambaConfig) -> None:
+        """
+        Initialize the MiniMamba model, bridging Mamba2 with HF's PreTrainedModel.
+        """
+        super().__init__(config)
+
+        # If your config includes Mamba2-like parameters, you can build a Mamba2Config from it:
+        mamba2_args = Mamba2Config(
+            dim=config.dim,
+            num_layers=config.num_layers,
+            num_heads=config.num_heads,
+            state_dim=config.state_dim,
+            num_groups=config.num_groups,
+            conv_size=config.conv_size,
+            use_mem_eff_path=config.use_mem_eff_path,
+            dt_bias=config.dt_bias,
+            D_has_head_dim=config.D_has_head_dim,
+            learnable_init_states=config.learnable_init_states,
+            ssm_chunk_size=config.ssm_chunk_size,
+            vocab_size=config.vocab_size,
+            ffn_dim_multiplier=config.ffn_dim_multiplier,
+            multiple_of=config.multiple_of,
+            norm_eps=config.norm_eps,
+            init_use_depth=config.init_use_depth,
+            init_base_std=config.init_base_std,
+            init_std_factor=config.init_std_factor,
+            bias=config.bias,
+        
+            # Torch / training:
+            seed=config.seed,
+        
+            # The init_config block nested in JSON:
+            
+            
+            # Additional Mamba or training fields:
+            
+            
+
+            
+            weight_tying=config.weight_tying if hasattr(config, "weight_tying") else False,
+            torch_dtype=getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype,
+        )
+
+        # Internally hold a Mamba2 model
+        self.mamba = Mamba2(config=mamba2_args)
+
+        # Because HF wants the final linear to be part of this top-level model,
+        # you *can* rely on Mamba2’s built-in embedding + output if you prefer.
+        # Mamba2 already has self.tok_emb and self.output. 
+        # So we typically do NOT need a separate embedding or lm_head here.
+        #
+        # We only do so if we want the “HF standard” tie-weights approach:
+        #    self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd)
+        #    self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        #    self.lm_head.weight = self.tok_emb.weight
+        # 
+        # But Mamba2 does that internally if config.weight_tying == True.
+
+        # This is optional: store any device or dtype you might want
+        self.device_ = 'cuda' if torch.cuda.is_available() else 'cpu'
+        if isinstance(config.torch_dtype, str):
+            self.dtype_ = getattr(torch, config.torch_dtype)
+        else:
+            self.dtype_ = config.torch_dtype
+        
+        # Parameter initialization (HF calls them with self._init_weights in some flows).
+        self.apply(self._init_weights)
+
+        print("MiniMamba Model Parameter Count: %.2fM\n" % (self._get_num_params() / 1e6,))
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        labels: torch.LongTensor = None,
+        **kwargs
+    ) -> CausalLMOutput:
+        """
+        Forward pass for causal language modeling. 
+        Returns a CausalLMOutput that includes loss (if labels is provided) and logits.
+        """
+        # Mamba2's forward expects (x: torch.Tensor, target: torch.Tensor|None, ...)
+        # but we only need the logits from the simple call:
+        logits = self.mamba(input_ids)  # shape: [batch, seq_len, vocab_size]
+
+        loss = None
+        if labels is not None:
+            # By default, huggingface GPT-like models shift the logits by one
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(
+                shift_logits.view(-1, shift_logits.size(-1)), 
+                shift_labels.view(-1)
+            )
+
+        return CausalLMOutput(
+            loss=loss,
+            logits=logits,
+        )
+
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.LongTensor,
+        max_new_tokens: int = 50,
+        temperature: float = 0.5,
+        top_k: int = 50,
+        top_p: float = 0.95,
+        eos_token_id: int = None,
+        pad_token_id: int = 0,
+        **kwargs
+    ):
+        """
+        A naive token-by-token generation loop (greedy + top-k/top-p + temperature).
+        """
+        # We'll accumulate new tokens in generated_ids
+        generated_ids = input_ids.clone()
+
+        for _ in range(max_new_tokens):
+            # Forward pass to get logits for the last token
+            outputs = self.forward(generated_ids)
+            logits = outputs.logits[:, -1, :]  # shape: (batch_size, vocab_size)
+
+            # Scale by temperature
+            if temperature != 1.0:
+                logits = logits / temperature
+
+            # Filter
+            logits = self.top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
+
+            # Sample next token
+            probs = F.softmax(logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)  # shape: (batch, 1)
+
+            # Append
+            generated_ids = torch.cat([generated_ids, next_token], dim=1)
+
+            # If we have an EOS token, we can break early if all sequences have ended
+            if eos_token_id is not None and (next_token == eos_token_id).all():
+                break
+
+        return generated_ids
+
+    @staticmethod
+    def top_k_top_p_filtering(
+        logits: torch.Tensor,
+        top_k: int = 50,
+        top_p: float = 0.95,
+        filter_value: float = float("-inf"),
+    ):
+        """
+        Filters logits using top-k and/or nucleus (top-p) filtering.
+        """
+        # top_k
+        if top_k > 0:
+            top_k = min(top_k, logits.size(-1))
+            indices_to_remove = logits < torch.topk(logits, top_k, dim=-1).values[:, -1, None]
+            logits[indices_to_remove] = filter_value
+
+        # top_p (nucleus)
+        if 0 < top_p < 1.0:
+            sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+            cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+            # Remove tokens with cumulative probability above the threshold
+            sorted_indices_to_remove = cumulative_probs > top_p
+
+            # Shift right to keep also the first token above threshold
+            sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
+            sorted_indices_to_remove[:, 0] = False
+
+            # Scatter to get back to original indexing
+            indices_to_remove = sorted_indices_to_remove.scatter(
+                dim=1, index=sorted_indices, src=sorted_indices_to_remove
+            )
+            logits[indices_to_remove] = filter_value
+
+        return logits
+
+    def _init_weights(self, module):
+        """
+        HF calls _init_weights to initialize parameters. 
+        If you prefer Mamba’s own init approach, you can call model.mamba.init_weights().
+        """
+        # As an example, we just call Mamba2's init routine for the entire submodel,
+        # or do some standard PyTorch inits for linear layers, embeddings, etc.
+        if isinstance(module, Mamba2):
+            module.init_weights()  # Mamba2’s internal init
+        elif isinstance(module, nn.Linear):
+            # e.g. standard xavier or normal init
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        # If needed, do your specialized inits for other modules
+
+    def _get_num_params(self):
+        # Count trainable params, subtract duplicates if tying weights, etc.
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)

norms.py

@@ -0,0 +1,357 @@
+"""Adapted from https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/norms.py"""
+
+import math
+
+from functools import partial
+
+import torch
+import torch.nn as nn
+
+import triton
+import triton.language as tl
+
+from torch.distributed._tensor import Partial, Replicate, Shard
+from torch.distributed._tensor.experimental import local_map
+from torch._utils import _get_available_device_type, _get_device_module
+
+
+def get_device_info():
+    device_type = _get_available_device_type()
+
+    if device_type is None:
+        device_type = "cuda"  # Default to CUDA
+
+    device_module = _get_device_module(device_type)
+    return device_type, device_module
+
+device_type, device_module = get_device_info()
+
+def build_norm(norm_type: str, dim: int, eps: float = 1e-6):
+    """
+    Builds the specified normalization layer based on the norm_type.
+
+    Args:
+        norm_type (str): The type of normalization layer to build.
+            Supported types: layernorm, np_layernorm, rmsnorm, fused_rmsnorm
+        dim (int): The dimension of the normalization layer.
+        eps (float, optional): The epsilon value for numerical stability. Defaults to 1e-6.
+
+    Returns:
+        The built normalization layer.
+
+    Raises:
+        NotImplementedError: If an unknown norm_type is provided.
+    """
+    norm_type = norm_type.lower()  # Normalize to lowercase
+
+    if norm_type == "layernorm":
+        return nn.LayerNorm(dim, eps=eps, bias=False)
+    elif norm_type == "np_layernorm":
+        return nn.LayerNorm(dim, eps=eps, elementwise_affine=False, bias=False)
+    elif norm_type == "rmsnorm":
+        return RMSNorm(dim, eps=eps)
+    elif norm_type == "fused_rmsnorm":
+        return FusedRMSNorm(dim, eps=eps)
+    else:
+        raise NotImplementedError(f"Unknown norm_type: '{norm_type}'")
+
+
+class FusedRMSNorm(nn.Module):
+    """Fused RMS Norm, wraps a fused Triton Kernel"""
+
+    def __init__(
+        self,
+        dim: int,
+        eps: float = 1e-6,
+    ):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.fused_rms_norm_fn = fused_rms_norm_fn
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """leverages Triton Fused RMS Norm kernel"""
+        return self.fused_rms_norm_fn(
+            x,
+            self.weight,
+            eps=self.eps,
+        )
+
+    def reset_parameters(self):
+        torch.nn.init.ones_(self.weight)  # type: ignore
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        """
+        Initialize the RMSNorm normalization layer.
+
+        Args:
+            dim (int): The dimension of the input tensor.
+            eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
+
+        Attributes:
+            eps (float): A small value added to the denominator for numerical stability.
+            weight (nn.Parameter): Learnable scaling parameter.
+
+        """
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        """
+        Apply the RMSNorm normalization to the input tensor.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The normalized tensor.
+
+        """
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        """
+        Forward pass through the RMSNorm layer.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The output tensor after applying RMSNorm.
+
+        """
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+    def reset_parameters(self):
+        torch.nn.init.ones_(self.weight)  # type: ignore
+
+
+# FusedRMSNorm in Triton
+
+# Credit
+# Tri Dao's Triton LayerNorm: https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/ops/triton/layer_norm.py
+# Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=1),
+        triton.Config({}, num_warps=2),
+        triton.Config({}, num_warps=4),
+        triton.Config({}, num_warps=8),
+        triton.Config({}, num_warps=16),
+        triton.Config({}, num_warps=32),
+    ],
+    key=["N"],
+)
+@triton.jit
+def _rms_norm_fwd_kernel(
+    X,
+    stride_x,
+    Y,
+    stride_y,
+    W,
+    Rstd,
+    eps,
+    M,  # num rows
+    N,  # num cols
+    block_N: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, block_N)
+
+    # Load input data and weights
+    mask = cols < N
+    x = tl.load(X + row * stride_x + cols, mask=mask, other=0.0).to(tl.float32)
+    w = tl.load(W + cols, mask=mask, other=0.0).to(tl.float32)
+
+    # Compute mean and variance
+    xbar = tl.where(cols < N, x, 0.0)
+    var = tl.sum(xbar * xbar, axis=0) / N
+    rstd = 1 / tl.sqrt(var + eps)
+
+    # Store the reciprocal standard deviation
+    tl.store(Rstd + row, rstd)
+
+    # Normalize and apply linear transformation
+    x_hat = x * rstd
+    y = x_hat * w
+
+    # Write output
+    tl.store(Y + row * stride_y + cols, y, mask=mask)
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=1),
+        triton.Config({}, num_warps=2),
+        triton.Config({}, num_warps=4),
+        triton.Config({}, num_warps=8),
+        triton.Config({}, num_warps=16),
+        triton.Config({}, num_warps=32),
+    ],
+    key=["N"],
+)
+@triton.jit
+def _rms_norm_bwd_kernel_sm(
+    X,
+    stride_x,
+    W,
+    DY,
+    stride_dy,
+    DX,
+    stride_dx,
+    Rstd,
+    DW,
+    eps,
+    M,  # num rows
+    N,  # num cols
+    rows_per_program,
+    block_N: tl.constexpr,
+):
+    row_block_id = tl.program_id(0)
+    row_start = row_block_id * rows_per_program
+    cols = tl.arange(0, block_N)
+    mask = cols < N
+
+    # Load weights
+    w = tl.load(W + cols, mask=mask, other=0.0).to(tl.float32)
+
+    # Accumulate gradients for weights
+    dw = tl.zeros((block_N,), dtype=tl.float32)
+
+    row_end = min(row_start + rows_per_program, M)
+    for row in range(row_start, row_end):
+        # Load input, output gradient, and reciprocal standard deviation
+        x = tl.load(X + row * stride_x + cols, mask=mask, other=0.0).to(tl.float32)
+        dy = tl.load(DY + row * stride_dy + cols, mask=mask, other=0.0).to(tl.float32)
+        rstd = tl.load(Rstd + row)
+
+        # Compute normalized input and gradients
+        x_hat = x * rstd
+        wdy = w * dy
+        dw += dy * x_hat
+        c1 = tl.sum(x_hat * wdy, axis=0) / N
+        dx = (wdy - x_hat * c1) * rstd
+
+        # Store input gradient
+        tl.store(DX + row * stride_dx + cols, dx, mask=mask)
+
+    # Store weight gradients
+    tl.store(DW + row_block_id * N + cols, dw, mask=mask)
+
+
+class TritonFusedRMSNorm(torch.autograd.Function):
+    @partial(
+        local_map,
+        out_placements=[Shard(1)],
+        in_placements=(None, [Shard(1)], [Replicate()], None),
+    )
+    @staticmethod
+    def forward(ctx, x, weight, eps):
+        x_shape_start = x.shape
+
+        # Flatten input
+        x = x.view(-1, x.shape[-1])
+        if x.stride(-1) != 1:
+            x = x.contiguous()
+        if weight.stride(-1) != 1:
+            weight = weight.contiguous()
+
+        M, N = x.shape
+        y = torch.empty_like(x)
+        rstd = torch.empty((M,), dtype=torch.float32, device=x.device)
+
+        max_size = 65536 // x.element_size()
+        block_N = min(max_size, triton.next_power_of_2(N))
+
+        if N > block_N:
+            raise ValueError(f"N {N} must be <= {block_N=}")
+
+        grid = lambda meta: (M,)
+        _rms_norm_fwd_kernel[grid](
+            x,
+            x.stride(0),
+            y,
+            y.stride(0),
+            weight,
+            rstd,
+            eps,
+            M,
+            N,
+            block_N,
+        )
+
+        ctx.eps = eps
+        ctx.save_for_backward(x, weight, rstd)
+        ctx.x_shape_start = x_shape_start
+
+        y = y.reshape(x_shape_start)
+        return y
+
+    @partial(
+        local_map,
+        out_placements=([Shard(1)], [Partial()], None),
+        in_placements=(None, [Shard(1)]),
+    )
+    @staticmethod
+    def backward(ctx, dy):
+        x, weight, rstd = ctx.saved_tensors
+        eps = ctx.eps
+        x_shape_start = ctx.x_shape_start
+
+        # Flatten input and output gradients
+        dy = dy.view(-1, dy.shape[-1])
+        if dy.stride(-1) != 1:
+            dy = dy.contiguous()
+
+        M, N = dy.shape
+        dx = torch.empty_like(x)
+
+        sm_count = device_module.get_device_properties(x.device).multi_processor_count
+        _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device)
+
+        max_size = 65536 // x.element_size()
+        block_N = min(max_size, triton.next_power_of_2(N))
+        rows_per_sm = math.ceil(M / sm_count)
+
+        if N > block_N:
+            raise ValueError(f"N {N} must be <= {block_N=}")
+
+        grid = lambda meta: (sm_count,)
+        _rms_norm_bwd_kernel_sm[grid](
+            x,
+            x.stride(0),
+            weight,
+            dy,
+            dy.stride(0),
+            dx,
+            dx.stride(0),
+            rstd,
+            _dw,
+            eps,
+            M,
+            N,
+            rows_per_sm,
+            block_N,
+        )
+        dw = _dw.sum(0).to(weight.dtype)
+        dx = dx.view(x_shape_start)
+        return dx, dw, None
+
+
+# expose fusedRMSNorm as a function
+def fused_rms_norm_fn(
+    x,
+    weight,
+    eps=1e-6,
+):
+    return TritonFusedRMSNorm.apply(
+        x,
+        weight,
+        eps,
+    )

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

ssm_compilable.py

@@ -0,0 +1,222 @@
+from typing import List, Optional, Tuple
+import torch
+
+from mamba_ssm.ops.triton.ssd_combined import _mamba_chunk_scan_combined_fwd, _mamba_chunk_scan_combined_bwd
+
+
+@torch.compile(options={"triton.cudagraphs": True}, fullgraph=True)
+def _compiled_mamba_chunk_scan_combined_fwd(x, dt, A, B, C, chunk_size, D=None, z=None, dt_bias=None, initial_states=None, seq_idx=None, cu_seqlens=None, dt_softplus=False, dt_limit=None):
+    return _mamba_chunk_scan_combined_fwd(x, dt, A, B, C, chunk_size, D=D, z=z, dt_bias=dt_bias, initial_states=initial_states, seq_idx=seq_idx, cu_seqlens=cu_seqlens, dt_softplus=dt_softplus, dt_limit=dt_limit)
+
+@torch.compile(options={"triton.cudagraphs": True}, fullgraph=True)
+def _compiled_mamba_chunk_scan_combined_bwd(dout, x, dt, A, B, C, out, chunk_size, D=None, z=None, dt_bias=None, initial_states=None, dfinal_states=None, seq_idx=None, dt_softplus=False, dt_limit=None):
+    return _mamba_chunk_scan_combined_bwd(dout, x, dt, A, B, C, out, chunk_size, D=D, z=z, dt_bias=dt_bias, initial_states=initial_states, dfinal_states=dfinal_states, seq_idx=seq_idx, dt_softplus=dt_softplus, dt_limit=dt_limit)
+
+
+@torch.library.custom_op(
+    "mamba_ssm::ssm_chunk_scan_combined_fwd",
+    mutates_args=(),
+    device_types="cuda",
+)
+def ssm_chunk_scan_combined_fwd(
+    x: torch.Tensor, 
+    dt: torch.Tensor, 
+    A: torch.Tensor, 
+    B: torch.Tensor, 
+    C: torch.Tensor, 
+    chunk_size: int, 
+    D: Optional[torch.Tensor] = None, 
+    z: Optional[torch.Tensor] = None,  
+    dt_bias: Optional[torch.Tensor] = None, 
+    initial_states: Optional[torch.Tensor] = None,  
+    seq_idx: Optional[torch.Tensor] = None,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    dt_softplus: bool = False,
+    dt_limit: Optional[List[float]] = None
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    out, out_x, dt_out, dA_cumsum, states, final_states, *rest = _mamba_chunk_scan_combined_fwd(x, dt, A, B, C, chunk_size, D=D, z=z, dt_bias=dt_bias, initial_states=initial_states, seq_idx=seq_idx, cu_seqlens=cu_seqlens, dt_softplus=dt_softplus, dt_limit=dt_limit)
+
+    return out, out_x if out_x is not None else out.new_empty(0), rest[0] if cu_seqlens is not None else out.new_empty(0)
+
+@ssm_chunk_scan_combined_fwd.register_fake
+def _ssm_chunk_scan_combined_fwd_fake(
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    chunk_size: int,
+    D: Optional[torch.Tensor] = None,
+    z: Optional[torch.Tensor] = None,
+    dt_bias: Optional[torch.Tensor] = None,
+    initial_states: Optional[torch.Tensor] = None,
+    seq_idx: Optional[torch.Tensor] = None,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    dt_softplus: bool = False,
+    dt_limit: Optional[List[float]] = None
+):
+    _, _, n_heads, head_dim = x.shape
+    return (
+        torch.empty_like(x), 
+        torch.empty_like(x) if z is not None else None, 
+        x.new_empty((cu_seqlens.size(0)-1, n_heads, head_dim, B.size(0))) if cu_seqlens is not None else None,
+    )
+
+@torch.library.custom_op(
+    "mamba_ssm::ssm_chunk_scan_combined_bwd", 
+    mutates_args=(),
+    device_types="cuda",
+)
+def ssm_chunk_scan_combined_bwd(
+    dout: torch.Tensor,
+    x: torch.Tensor, 
+    dt: torch.Tensor, 
+    A: torch.Tensor, 
+    B: torch.Tensor, 
+    C: torch.Tensor,
+    out: torch.Tensor,
+    chunk_size: int, 
+    D: Optional[torch.Tensor] = None,
+    z: Optional[torch.Tensor] = None,
+    dt_bias: Optional[torch.Tensor] = None,
+    initial_states: Optional[torch.Tensor] = None,
+    seq_idx: Optional[torch.Tensor] = None,
+    dt_softplus: bool = False,
+    dt_limit: Optional[List[float]] = None
+)-> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    dx, ddt, dA, dB, dC, dD, dz, ddt_bias, dinitial_states = _mamba_chunk_scan_combined_bwd(dout, x, dt, A, B, C, out, chunk_size, D=D, z=z, dt_bias=dt_bias, initial_states=initial_states, dfinal_states=None, seq_idx=seq_idx, dt_softplus=dt_softplus, dt_limit=dt_limit)
+    return (
+        dx,
+        ddt,
+        dA,
+        dB,
+        dC,
+        dD if dD is not None else dx.new_empty(0),
+        dz if dz is not None else dx.new_empty(0),
+        ddt_bias if ddt_bias is not None else dx.new_empty(0),
+        dinitial_states if dinitial_states is not None else dx.new_empty(0)
+    )
+
+@ssm_chunk_scan_combined_bwd.register_fake
+def _ssm_chunk_scan_combined_bwd_fake(
+    dout: torch.Tensor,
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    out: torch.Tensor,
+    chunk_size: int,
+    D: Optional[torch.Tensor] = None,
+    z: Optional[torch.Tensor] = None,
+    dt_bias: Optional[torch.Tensor] = None,
+    initial_states: Optional[torch.Tensor] = None,
+    seq_idx: Optional[torch.Tensor] = None,
+    dt_softplus: bool = False,
+    dt_limit: Optional[List[float]] = None
+):
+    return (
+        torch.empty_like(x),
+        torch.empty_like(dt),
+        torch.empty_like(A),
+        torch.empty_like(B),
+        torch.empty_like(C),
+        torch.empty_like(D) if D is not None else None,
+        torch.empty_like(z) if z is not None else None,
+        torch.empty_like(dt_bias) if dt_bias is not None else None,
+        torch.empty_like(initial_states) if initial_states is not None else None,
+    )
+
+
+def ssm_chunk_scan_combined_setup_context(ctx, inputs, output):
+    x, dt, A, B, C, chunk_size, D, z, dt_bias, initial_states, seq_idx, cu_seqlens, dt_softplus, dt_limit = inputs
+    out, out_x, state_varlen = output
+
+    ctx.save_for_backward(out if z is None else out_x, x, dt, A, B, C, D, z, dt_bias, initial_states, seq_idx)
+    ctx.dt_softplus = dt_softplus
+    ctx.chunk_size = chunk_size
+    ctx.dt_limit = dt_limit
+
+def ssm_chunk_scan_combined_bridge(ctx, dout, dout_x, dout_state_varlen):
+    out, x, dt, A, B, C, D, z, dt_bias, initial_states, seq_idx = ctx.saved_tensors
+
+    dx, ddt, dA, dB, dC, dD, dz, ddt_bias, dinitial_states = ssm_chunk_scan_combined_bwd(dout, x, dt, A, B, C, out, ctx.chunk_size, D=D, z=z, dt_bias=dt_bias, initial_states=initial_states, seq_idx=seq_idx, dt_softplus=ctx.dt_softplus, dt_limit=ctx.dt_limit)
+
+    return (
+        dx,
+        ddt,
+        dA,
+        dB,
+        dC, 
+        None,
+        dD if D is not None else None, 
+        dz if z is not None else None,
+        ddt_bias if dt_bias is not None else None, 
+        dinitial_states if initial_states is not None else None, 
+        None,
+        None,
+        None,
+        None,
+    )
+
+# Register custom autograd function
+torch.library.register_autograd(
+    "mamba_ssm::ssm_chunk_scan_combined_fwd",
+    ssm_chunk_scan_combined_bridge,
+    setup_context=ssm_chunk_scan_combined_setup_context,
+)
+
+def mamba_chunk_scan_combined(x, dt, A, B, C, chunk_size, D=None, z=None, dt_bias=None, initial_states=None, seq_idx=None, cu_seqlens=None, dt_softplus=False, dt_limit=(0.0, float("inf"))):
+    """
+    Argument:
+        x: (batch, seqlen, nheads, headdim)
+        dt: (batch, seqlen, nheads)
+        A: (nheads)
+        B: (batch, seqlen, ngroups, dstate)
+        C: (batch, seqlen, ngroups, dstate)
+        chunk_size: int
+        D: (nheads, headdim) or (nheads,)
+        z: (batch, seqlen, nheads, headdim)
+        dt_bias: (nheads,)
+        initial_states: (batch, nheads, headdim, dstate)
+        seq_idx: (batch, seqlen)
+        cu_seqlens: (num_sequences + 1) or None
+        dt_softplus: Whether to apply softplus to dt
+    Return:
+        out: (batch, seqlen, nheads, headdim)
+    """
+    
+    out, _, varlen_states  = ssm_chunk_scan_combined_fwd(x, dt, A, B, C, chunk_size, D=D, z=z, dt_bias=dt_bias, initial_states=initial_states, seq_idx=seq_idx, cu_seqlens=cu_seqlens, dt_softplus=dt_softplus, dt_limit=dt_limit)
+    if cu_seqlens is not None:
+        return out, varlen_states
+    return out
+
+if __name__ == "__main__":
+    from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined as mamba_chunk_scan_combined_ref
+
+    torch.manual_seed(0)
+    torch.cuda.manual_seed(0)
+
+    x = torch.randn(2, 3, 4, 5).cuda()
+    dt = torch.randn(2, 3, 4).cuda()
+    A = torch.randn(4).cuda()
+    B = torch.randn(2, 3, 4, 5).cuda()
+    C = torch.randn(2, 3, 4, 5).cuda()
+    chunk_size = 2
+    D = torch.randn(4, 5).cuda()
+    z = torch.randn(2, 3, 4, 5).cuda()
+    dt_bias = torch.randn(4).cuda()
+
+    out = mamba_chunk_scan_combined(x, dt, A, B, C, chunk_size, D=D, z=z, dt_bias=dt_bias)
+
+    print(out.min(), out.max(), out.mean(), out.std())
+
+    compiled_mamba_chunk_scan_combined = torch.compile(mamba_chunk_scan_combined)
+    out = compiled_mamba_chunk_scan_combined(x, dt, A, B, C, chunk_size, D=D, z=z, dt_bias=dt_bias)
+
+    print(out.min(), out.max(), out.mean(), out.std())
+
+    out_ref = mamba_chunk_scan_combined_ref(x, dt, A, B, C, chunk_size, D=D, z=z, dt_bias=dt_bias)
+
+    print(out_ref.min(), out_ref.max(), out_ref.mean(), out_ref.std())
+

tokenizer_config.json

@@ -0,0 +1,27 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "199999": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200018": {
+      "content": "<|endofprompt|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 128000,
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}