Upload JambaForCausalLM

config.json

@@ -0,0 +1,208 @@
+{
+  "architectures": [
+    "JambaForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "attn_hidden_size": -1,
+  "attn_implementation": "flash_attention_2",
+  "attn_implementation_new": "flash_attention_2",
+  "attn_layer_offset": 4,
+  "attn_layer_period": 8,
+  "attn_reuse_every_i_layer": -1,
+  "auto_map": {
+    "AutoConfig": "configuration_jamba.JambaConfig",
+    "AutoModelForCausalLM": "modeling_jamba.JambaForCausalLM"
+  },
+  "bos_token_id": 1,
+  "calc_logits_for_entire_prompt": false,
+  "compact_gating": false,
+  "compute_attn_mat": false,
+  "d_conv": 4,
+  "dense_public_ffn_structure": false,
+  "double_v_dim": false,
+  "enable_mod": false,
+  "eos_token_id": 2,
+  "expert_layer_offset": 1,
+  "expert_layer_period": 2,
+  "ffn_expand_ratio": 3,
+  "ffn_reuse_every_i_layer": -1,
+  "ffn_sharing_config": null,
+  "fully_parallel_jamba": false,
+  "fused_multihead_config": null,
+  "global_attn_idx": [],
+  "gradient_checkpoint_layer": null,
+  "hash_grid_config": null,
+  "hash_grid_config_mlp": null,
+  "hidden_act": "silu",
+  "hidden_size": 3072,
+  "hybrid_block_indices": [],
+  "hybrid_decoder_layer": "mamba",
+  "initializer_range": 0.02,
+  "intermediate_size": 0,
+  "kq_head_dim": -1,
+  "kq_norm": "none",
+  "kv_reuse_every_i_layer": -1,
+  "kv_reuse_group": null,
+  "kv_weight_reuse": false,
+  "layer_type": [
+    "m",
+    "a",
+    "m",
+    "a",
+    "a",
+    "a",
+    "m",
+    "a",
+    "m",
+    "a",
+    "m",
+    "a",
+    "a",
+    "a",
+    "m",
+    "a",
+    "m",
+    "a",
+    "m",
+    "a",
+    "a",
+    "a",
+    "m",
+    "a",
+    "m",
+    "a",
+    "m",
+    "a",
+    "m",
+    "a",
+    "m",
+    "a",
+    "m",
+    "a",
+    "m",
+    "a"
+  ],
+  "layer_types": [
+    "deltanet",
+    "f",
+    "m2",
+    "f",
+    "a",
+    "f",
+    "m2",
+    "f",
+    "deltanet",
+    "f",
+    "m2",
+    "f",
+    "a",
+    "f",
+    "m2",
+    "f",
+    "deltanet",
+    "f",
+    "m2",
+    "f",
+    "a",
+    "f",
+    "m2",
+    "f",
+    "deltanet",
+    "f",
+    "m2",
+    "f",
+    "deltanet",
+    "f",
+    "m2",
+    "f",
+    "deltanet",
+    "f",
+    "m2",
+    "f"
+  ],
+  "layerwise_memory_token": false,
+  "local_expand_ratio": 1,
+  "local_global_dual_branch": false,
+  "local_global_dual_branch_merge_op": "mean",
+  "lookback_mode": "",
+  "macro_arch": "",
+  "mamba2_headdim": 64,
+  "mamba_attnaug_config": null,
+  "mamba_conv_bias": true,
+  "mamba_d_conv": 4,
+  "mamba_d_state": 16,
+  "mamba_dt_rank": 192,
+  "mamba_expand": 2,
+  "mamba_inner_layernorms": true,
+  "mamba_latent_size": null,
+  "mamba_multihead_config": null,
+  "mamba_proj_bias": false,
+  "mamba_reuse_every_i_layer": -1,
+  "max_position_embeddings": 2048,
+  "memory_tokens_interspersed_every": 0,
+  "mlp_hidden_act": "silu",
+  "mod_topk": 2,
+  "model_type": "jamba",
+  "moe_config": null,
+  "nGPT_config": {
+    "extra_grad": false,
+    "gate_scaling": false,
+    "init_norm": false,
+    "learned_scaling": false,
+    "norm_bc": false,
+    "norm_gating": false,
+    "norm_ssm_input": false,
+    "post_norm": false,
+    "qk_norm": false,
+    "weight_norm": true
+  },
+  "nGPT_mode": null,
+  "new_seq_length": 2048,
+  "no_dt_bias": false,
+  "num_attention_heads": 24,
+  "num_attn_per_ffn": 3,
+  "num_experts": 1,
+  "num_experts_per_tok": 1,
+  "num_ffn": 1,
+  "num_hidden_layers": 36,
+  "num_key_value_heads": 6,
+  "num_mamba": 1,
+  "num_memory_tokens": 256,
+  "orig_max_position_embeddings": 2048,
+  "other_args": null,
+  "output_router_logits": false,
+  "pad_token_id": 0,
+  "public_ffn_structure": false,
+  "pure_linear_attn": false,
+  "reduce_attn_ratio": 0.5,
+  "reduce_method": "mean",
+  "repeat_ffn": null,
+  "rms_norm_eps": 1e-06,
+  "rope": true,
+  "rope_theta": 10000.0,
+  "rope_type": null,
+  "router_aux_loss_coef": 0.001,
+  "save_input_output": false,
+  "self_attn_type": null,
+  "seq_length": 2048,
+  "sequential_jamba": false,
+  "share_kv": false,
+  "shared_module_attn": "",
+  "shared_module_mamba": "",
+  "sliding_window": null,
+  "sliding_window_size": null,
+  "supernet_config": null,
+  "swa_full_head": false,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.48.2",
+  "use_cache": false,
+  "use_mamba2": false,
+  "use_mamba_kernels": true,
+  "use_nGPT": true,
+  "use_nemotron5": false,
+  "v_head_dim": -1,
+  "visual_attn": false,
+  "visual_entropy": false,
+  "vocab_size": 131072
+}

configuration_jamba.py

@@ -0,0 +1,429 @@
+# coding=utf-8
+# Copyright 2024 AI21 Labs Ltd. and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" Jamba model configuration"""
+import math
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class JambaConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`JambaModel`]. It is used to instantiate a
+    Jamba model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of the jamba-small architecture.
+
+    [ai21labs/jamba-small](https://huggingface.co/ai21labs/Jamba-v0.1)
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 65536):
+            Vocabulary size of the Jamba model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`JambaModel`]
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether the model's input and output word embeddings should be tied. Note that this is only relevant if the
+            model has a output word embedding layer.
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 14336):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 32):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        num_key_value_heads (`int`, *optional*, defaults to 8):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details checkout [this
+            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `8`.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        calc_logits_for_entire_prompt (`bool`, *optional*, defaults to `False`):
+            Whether or not to calculate logits for entire prompt during generation. If `False`, only the logits of the
+            last prompt token will be calculated, which are the only logits needed for generation. For long sequences,
+            the logits for the entire sequence may use a lot of memory so setting `calc_logits_for_entire_prompt=False`
+            will reduce memory footprint significantly.
+            Note: some generation features may not be available if this is set to `False`.
+        output_router_logits (`bool`, *optional*, defaults to `False`):
+            Whether or not the router logits should be returned by the model. Enabling this will also
+            allow the model to output the auxiliary loss. See [here]() for more details
+        router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
+            The aux loss factor for the total loss.
+        pad_token_id (`int`, *optional*, defaults to 0):
+            The id of the padding token.
+        bos_token_id (`int`, *optional*, defaults to 1):
+            The id of the "beginning-of-sequence" token.
+        eos_token_id (`int`, *optional*, defaults to 2):
+            The id of the "end-of-sequence" token.
+        sliding_window (`int`, *optional*):
+            Sliding window attention window size. If not specified, will default to `None`.
+        n_ctx (`int`, *optional*, defaults to 262144):
+            This value doesn't have any real effect. The maximum sequence length that this model is intended to be
+            used with. It can be used with longer sequences, but performance may degrade.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        num_experts_per_tok (`int`, *optional*, defaults to 2):
+            The number of experts to root per-token, can be also interpreted as the `top-p` routing
+            parameter
+        num_experts (`int`, *optional*, defaults to 16):
+            Number of experts per Sparse MLP layer.
+        expert_layer_period (`int`, *optional*, defaults to 2):
+            Once in this many layers, we will have an expert layer
+        expert_layer_offset (`int`, *optional*, defaults to 1):
+            The first layer index that contains an expert mlp layer
+        attn_layer_period (`int`, *optional*, defaults to 8):
+            Once in this many layers, we will have a vanilla attention layer
+        attn_layer_offset (`int`, *optional*, defaults to 4):
+            The first layer index that contains a vanilla attention mlp layer
+        use_mamba_kernels (`bool`, *optional*, defaults to `True`):
+            Flag indicating whether or not to use the fast mamba kernels. These are available only if `mamba-ssm` and
+            `causal-conv1d` are installed, and the mamba modules are running on a CUDA device. Raises ValueError if
+            `True` and kernels are not available
+        mamba_d_state (`int`, *optional*, defaults to 16):
+            The dimension the mamba state space latents
+        mamba_d_conv (`int`, *optional*, defaults to 4):
+            The size of the mamba convolution kernel
+        mamba_expand (`int`, *optional*, defaults to 2):
+            Expanding factor (relative to hidden_size) used to determine the mamba intermediate size
+        mamba_dt_rank (`Union[int,str]`, *optional*, defaults to `"auto"`):
+            Rank of the the mamba discretization projection matrix. `"auto"` means that it will default to `math.ceil(self.hidden_size / 16)`
+        mamba_conv_bias (`bool`, *optional*, defaults to `True`):
+            Flag indicating whether or not to use bias in the convolution layer of the mamba mixer block.
+        mamba_proj_bias (`bool`, *optional*, defaults to `False`):
+            Flag indicating whether or not to use bias in the input and output projections (["in_proj", "out_proj"]) of the mamba mixer block
+        mamba_inner_layernorms (`bool`, *optional*, defaults to `True`):
+            Flag indicating whether or not to apply layernorms to internal mamba activations
+
+    """
+
+    model_type = "jamba"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+            self,
+            vocab_size=65536,
+            tie_word_embeddings=False,
+            hidden_size=4096,
+            intermediate_size=14336,
+            num_hidden_layers=32,
+            num_attention_heads=32,
+            num_key_value_heads=8,
+            hidden_act="silu",
+            initializer_range=0.02,
+            rms_norm_eps=1e-6,
+            use_cache=True,
+            calc_logits_for_entire_prompt=False,
+            output_router_logits=False,
+            router_aux_loss_coef=0.001,
+            pad_token_id=0,
+            bos_token_id=1,
+            eos_token_id=2,
+            sliding_window=None,
+            max_position_embeddings=262144,
+            orig_max_position_embeddings=None,
+            attention_dropout=0.0,
+            num_experts_per_tok=2,
+            num_experts=16,
+            expert_layer_period=2,
+            expert_layer_offset=1,
+            attn_layer_period=8,
+            attn_layer_offset=4,
+            use_mamba_kernels=True,
+            mamba_d_state=16,
+            mamba_d_conv=4,
+            mamba_expand=2,
+            mamba_dt_rank="auto",
+            mamba_conv_bias=True,
+            mamba_proj_bias=False,
+            mamba_inner_layernorms=True,
+
+            hybrid_decoder_layer='mamba',
+            share_kv=False,
+            double_v_dim=False,
+            compact_gating=False,
+            kv_reuse_every_i_layer=-1,
+            kv_reuse_group=None,
+            kv_weight_reuse=False,
+
+            num_ffn=1,
+            ffn_reuse_every_i_layer=-1,
+            attn_reuse_every_i_layer=-1,
+            mamba_reuse_every_i_layer=-1,
+
+            macro_arch='',
+
+            lookback_mode='',
+
+            shared_module_attn='',
+            shared_module_mamba='',
+
+            ffn_sharing_config=None,
+
+            sliding_window_size=None,
+            global_attn_idx=None,
+
+            num_mamba=1,
+            mamba_latent_size=None,
+
+            public_ffn_structure=False,
+            num_attn_per_ffn=3,
+            dense_public_ffn_structure=False,
+
+            local_global_dual_branch=False,
+            local_expand_ratio=1,
+            local_global_dual_branch_merge_op='mean',
+
+            mamba_multihead_config=None,
+
+            moe_config=None,
+
+            enable_mod=False,
+            mod_topk=2,
+
+            sequential_jamba=False,
+            fully_parallel_jamba=False,
+
+            attn_implementation_new='sdpa',
+
+            fused_multihead_config=None,
+
+            compute_attn_mat=False,
+            visual_attn=False,
+            save_input_output=False,
+
+            use_mamba2=False,
+            mamba2_headdim=64,
+
+            swa_full_head=False,
+
+            gradient_checkpoint_layer=None,
+
+            rope_type=None,
+            
+            visual_entropy=False,
+            
+            use_nemotron5=False,
+            
+            use_nGPT=False,
+            nGPT_mode=None,
+            nGPT_config=None,
+            
+            mamba_attnaug_config=None,
+            
+            no_dt_bias=False,
+            
+            hash_grid_config=None,
+            
+            hash_grid_config_mlp=None,
+            
+            repeat_ffn=None,
+            
+            layer_types=None,
+            
+            supernet_config=None,
+            
+            pure_linear_attn=False,
+            
+            self_attn_type=None,
+            
+            other_args=None,
+            
+            ffn_expand_ratio=None,
+            
+            d_conv=4,
+            
+            layerwise_memory_token=False,
+            
+            **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.tie_word_embeddings = tie_word_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.sliding_window = sliding_window
+        self.max_position_embeddings = max_position_embeddings
+        self.orig_max_position_embeddings = orig_max_position_embeddings
+        self.attention_dropout = attention_dropout
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+
+        self.use_cache = use_cache
+        self.calc_logits_for_entire_prompt = calc_logits_for_entire_prompt
+        self.output_router_logits = output_router_logits
+        self.router_aux_loss_coef = router_aux_loss_coef
+
+        self.num_experts_per_tok = num_experts_per_tok
+        self.num_experts = num_experts
+        self.expert_layer_period = expert_layer_period
+        self.expert_layer_offset = expert_layer_offset
+        self.attn_layer_period = attn_layer_period
+        self.attn_layer_offset = attn_layer_offset
+
+        self.use_mamba_kernels = use_mamba_kernels
+        self.mamba_d_state = mamba_d_state
+        self.mamba_d_conv = mamba_d_conv
+        self.mamba_expand = mamba_expand
+        self.mamba_dt_rank = math.ceil(self.hidden_size / 16) if mamba_dt_rank == "auto" else mamba_dt_rank
+        self.mamba_conv_bias = mamba_conv_bias
+        self.mamba_proj_bias = mamba_proj_bias
+        self.mamba_inner_layernorms = mamba_inner_layernorms
+
+        # added by Xin 
+        self.reduce_method = kwargs.pop("reduce_method", "mean")
+        self.hybrid_block_indices = kwargs.pop("hybrid_block_indices", [])
+        self.reduce_attn_ratio = kwargs.pop("reduce_attn_ratio", 0.5)
+        self.attn_hidden_size = kwargs.pop("attn_hidden_size", -1)
+        self.kq_head_dim = kwargs.pop("kq_head_dim", -1)
+        self.v_head_dim = kwargs.pop("v_head_dim", -1)
+        self.kq_norm = kwargs.pop("kq_norm", None)
+        self.rope = kwargs.pop("rope", False)
+        self.rope_theta = kwargs.pop("rope_theta", 10000.0)
+        self.num_memory_tokens = kwargs.pop("num_memory_tokens", 0)
+        self.memory_tokens_interspersed_every = kwargs.pop("memory_tokens_interspersed_every", 0)
+
+        #! adhoc change
+        self.new_seq_length = 2048
+        self.visual_entropy = kwargs.pop("visual_entropy", False)
+
+        self.hybrid_decoder_layer = hybrid_decoder_layer
+        self.share_kv = share_kv
+        self.double_v_dim = double_v_dim
+        self.compact_gating = compact_gating
+        self.kv_reuse_every_i_layer = kv_reuse_every_i_layer
+        self.kv_reuse_group = kv_reuse_group
+        self.kv_weight_reuse = kv_weight_reuse
+
+        self.num_ffn = num_ffn
+        self.ffn_reuse_every_i_layer = ffn_reuse_every_i_layer
+        self.attn_reuse_every_i_layer = attn_reuse_every_i_layer
+        self.mamba_reuse_every_i_layer = mamba_reuse_every_i_layer
+
+        self.macro_arch = macro_arch
+
+        self.lookback_mode = lookback_mode
+
+        self.shared_module_attn = shared_module_attn
+        self.shared_module_mamba = shared_module_mamba
+
+        self.ffn_sharing_config = ffn_sharing_config
+
+        self.sliding_window_size = sliding_window_size
+        self.global_attn_idx = global_attn_idx
+
+        self.num_mamba = num_mamba
+
+        self.mamba_latent_size = mamba_latent_size
+
+        self.public_ffn_structure = public_ffn_structure
+        self.num_attn_per_ffn = num_attn_per_ffn
+        self.dense_public_ffn_structure = dense_public_ffn_structure
+
+        self.local_global_dual_branch = local_global_dual_branch
+        self.local_expand_ratio = local_expand_ratio
+        self.local_global_dual_branch_merge_op = local_global_dual_branch_merge_op
+
+        self.mamba_multihead_config = mamba_multihead_config
+
+        self.moe_config = moe_config
+
+        self.enable_mod = enable_mod
+        self.mod_topk = mod_topk
+
+        self.sequential_jamba = sequential_jamba
+        self.fully_parallel_jamba = fully_parallel_jamba
+
+        self.attn_implementation_new = attn_implementation_new
+
+        self.fused_multihead_config = fused_multihead_config
+
+        self.compute_attn_mat = compute_attn_mat
+        self.visual_attn = visual_attn
+        self.save_input_output = save_input_output
+
+        self.use_mamba2 = use_mamba2
+        self.mamba2_headdim = mamba2_headdim
+
+        self.swa_full_head = swa_full_head
+
+        self.gradient_checkpoint_layer = gradient_checkpoint_layer
+
+        self.rope_type = rope_type
+        
+        self.visual_entropy = visual_entropy
+        
+        self.use_nemotron5 = use_nemotron5
+        
+        self.use_nGPT = use_nGPT
+        self.nGPT_mode = nGPT_mode
+        
+        self.mamba_attnaug_config = mamba_attnaug_config
+        
+        self.no_dt_bias = no_dt_bias
+        
+        self.nGPT_config = nGPT_config
+        
+        self.hash_grid_config = hash_grid_config
+        
+        self.hash_grid_config_mlp = hash_grid_config_mlp
+        
+        self.repeat_ffn = repeat_ffn
+        
+        self.layer_types = layer_types
+        
+        self.supernet_config = supernet_config
+        
+        self.pure_linear_attn = pure_linear_attn
+        
+        self.self_attn_type = self_attn_type
+        
+        self.other_args = other_args
+        
+        self.ffn_expand_ratio = ffn_expand_ratio
+        
+        self.d_conv = d_conv
+        
+        self.layerwise_memory_token = layerwise_memory_token
+                
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )

delta_net.py

@@ -0,0 +1,333 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Dict, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from torch.nn import functional as F
+
+from fla.modules import FusedRMSNormSwishGate, RMSNorm, ShortConvolution
+from fla.ops.delta_rule import chunk_delta_rule, fused_recurrent_delta_rule
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+    from fla.models.utils import Cache
+
+
+def elu_p1(x):
+    return (F.elu(x, 1., False) + 1.).to(x)
+
+
+def sum_norm(x):
+    return (x / x.sum(-1, keepdim=True)).to(x)
+
+
+class DeltaNet(nn.Module):
+    r"""
+    The layer implementaion for [Parallelizing Linear Transformers with the Delta Rule over Sequence Length](https://arxiv.org/abs/2406.06484).  # noqa:
+    DeltaNet was originally proposed in [Linear Transformers Are Secretly Fast Weight Programmers](https://arxiv.org/abs/2102.11174). # noqa
+
+    Args:
+        mode (str, Optional):
+            Which DeltaNet kernel to use.
+            Currently available: `chunk`, `fused_recurrent`, and `fused_chunk`.
+            Default: `chunk`.
+        hidden_size (int, Optional):
+            The hidden size of the input. Default: 1024.
+        expand_k (float, Optional):
+            The expansion ratio for the key dim. Default: 1.0.
+        expand_v (float, Optional):
+            The expansion ratio for the value dim. Default: 1.0.
+        num_heads (int, Optional):
+            The number of heads. Default: 4.
+        use_beta (bool, Optional):
+            Whether to use beta. Default: `True`.
+        use_gate (bool, Optional):
+            Whether to use output gate. Default: `False`.
+        use_short_conv (bool, Optional):
+            Whether to use short convolutions. Default: `True`.
+        conv_size (int, Optional):
+            The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4.
+        conv_bias (bool, Optional):
+            Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`.
+        allow_neg_eigval (bool, Optional):
+            Allow negative eigenvalues. Default: `False`. If set to `True`, the beta will be multiplied by 2.
+            See reference: [Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues](https://arxiv.org/abs/2411.12537)
+        layer_idx (int, Optional):
+            The index of the layer. Default: None.
+        norm_eps (float, Optional):
+            The epsilon value for the layernorm/rmsnorm layer. Default: 1e-5.
+        qk_activation (str, Optional):
+            The activation function for the query and key. Default: `silu`.
+        qk_norm (str, Optional):
+            The normalization method for the query and key. Default: `l2`.
+    """
+
+    def __init__(
+        self,
+        mode: str = 'chunk',
+        d_model: int = None,
+        hidden_size: int = 1024,
+        expand_k: float = 1.0,
+        expand_v: float = 1.0,
+        num_heads: int = 4,
+        use_beta: bool = True,
+        use_gate: bool = False,
+        use_short_conv: bool = True,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        allow_neg_eigval: bool = False,
+        layer_idx: int = None,
+        qk_activation: str = 'silu',
+        qk_norm: str = 'l2',
+        norm_eps: float = 1e-5,
+        config = None,
+        **kwargs
+    ) -> DeltaNet:
+        super().__init__()
+
+        self.mode = mode
+        self.qk_activation = qk_activation
+        self.qk_norm = qk_norm
+
+        assert self.qk_activation in ['silu', 'relu', 'elu', 'identity']
+        assert self.qk_norm in ['l2', 'sum']
+        
+        self.config = config
+        if self.config is not None and self.config.use_nGPT and 'extra_grad' in self.config.nGPT_config and self.config.nGPT_config['extra_grad']:
+            self.weight_norm = True
+        else:
+            self.weight_norm = False
+
+        if d_model is not None:
+            hidden_size = d_model
+        self.hidden_size = hidden_size
+        self.expand_k = expand_k
+        self.expand_v = expand_v
+        self.num_heads = num_heads
+        self.use_gate = use_gate
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+        self.allow_neg_eigval = allow_neg_eigval
+
+        self.key_dim = int(hidden_size * expand_k)
+        self.value_dim = int(hidden_size * expand_v)
+        self.head_k_dim = self.key_dim // num_heads
+        self.head_v_dim = self.value_dim // num_heads
+        self.layer_idx = layer_idx
+
+        self.silu = nn.SiLU()
+        if mode == 'fused_chunk':
+            raise NotImplementedError("fused_chunk_delta_rule is now deprecated. Please use `chunk_delta_rule` instead.")
+        assert mode in ['chunk', 'fused_recurrent'], f"Not suppoerted mode `{mode}`."
+        assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}"
+        assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}"
+
+        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+
+        self.use_beta = use_beta
+        if self.use_beta:
+            self.b_proj = nn.Linear(hidden_size, self.num_heads, bias=False)
+        if use_short_conv:
+            self.conv_size = conv_size
+            self.q_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                activation='silu' if qk_activation == 'silu' else None
+            )
+            self.k_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                activation='silu' if qk_activation == 'silu' else None
+            )
+            self.v_conv1d = ShortConvolution(
+                hidden_size=self.value_dim,
+                kernel_size=conv_size,
+                activation='silu'
+            )
+        else:
+            raise UserWarning(
+                "ShortConvolution is crucial to the performance. "
+                "Do not turn it off, i.e., setting `use_short_conv=False` unless you know what you are doing."
+            )
+        if use_gate:
+            self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+            self.o_norm = FusedRMSNormSwishGate(self.head_v_dim, eps=norm_eps)
+        else:
+            self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps)
+
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+        self.apply(self._initialize_weights)
+
+    def _initialize_weights(self, module: nn.Module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        module._is_hf_initialized = True
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+        **kwargs: Unpack[Dict]
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        # change to inference mode.
+        mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode
+
+        last_state = None
+        if past_key_values is not None and len(past_key_values) > self.layer_idx:
+            last_state = past_key_values[self.layer_idx]
+
+        if self.use_short_conv:
+            conv_state_q, conv_state_k, conv_state_v = None, None, None
+            if last_state is not None:
+                conv_state_q, conv_state_k, conv_state_v = last_state['conv_state']
+            conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None
+            position_ids = kwargs.get('position_ids', None)
+            
+            q = self.q_proj(hidden_states)
+            if self.weight_norm:
+                q = q / self.q_proj.weight.norm(p=2, dim=1)
+            
+            q, conv_state_q = self.q_conv1d(x=q,
+                                            mask=conv_mask,
+                                            cache=conv_state_q,
+                                            output_final_state=use_cache,
+                                            seq_idx=position_ids)
+            
+            k = self.k_proj(hidden_states)
+            if self.weight_norm:
+                k = k / self.k_proj.weight.norm(p=2, dim=1)
+            k, conv_state_k = self.k_conv1d(x=k,
+                                            mask=conv_mask,
+                                            cache=conv_state_k,
+                                            output_final_state=use_cache,
+                                            seq_idx=position_ids)
+
+            v = self.v_proj(hidden_states)
+            if self.weight_norm:
+                v = v / self.v_proj.weight.norm(p=2, dim=1)
+            v, conv_state_v = self.v_conv1d(x=v,
+                                            mask=conv_mask,
+                                            cache=conv_state_v,
+                                            output_final_state=use_cache,
+                                            seq_idx=position_ids)
+        else:
+            q = self.q_proj(hidden_states)
+            k = self.k_proj(hidden_states)
+            v = self.v_proj(hidden_states)
+
+            if self.weight_norm:
+                q = q / self.q_proj.weight.norm(p=2, dim=1)
+                k = k / self.k_proj.weight.norm(p=2, dim=1)
+                v = v / self.v_proj.weight.norm(p=2, dim=1)
+                
+            if self.qk_activation == 'silu':
+                q, k = self.silu(q), self.silu(k)
+            
+            v = self.silu(v)
+
+        q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k))
+        v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim)
+        if self.qk_activation != 'silu':
+            if self.qk_activation == 'relu':
+                q, k = q.relu(), k.relu()
+            elif self.qk_activation == 'elu':
+                q, k = elu_p1(q), elu_p1(k)
+            elif self.qk_activation == 'identity':
+                pass
+            else:
+                raise NotImplementedError
+
+        if self.qk_norm == 'sum':
+            q = sum_norm(q).to(q)
+            k = sum_norm(k).to(k)
+
+        if self.use_beta:
+            beta = self.b_proj(hidden_states)
+            
+            if self.weight_norm:
+                beta = beta / self.b_proj.weight.norm(p=2, dim=1)
+            
+            beta = beta.sigmoid()
+        else:
+            beta = q.new_ones(q.shape[0], q.shape[1], q.shape[2])
+
+        if self.allow_neg_eigval:
+            beta = beta * 2.
+
+        # dealing with padding
+        if attention_mask is not None:
+            beta = beta.mul(attention_mask[:, -beta.shape[-2]:, None])
+
+        recurrent_state = last_state['recurrent_state'] if last_state is not None else None
+        cu_seqlens = kwargs.get('cu_seqlens', None)
+        if mode == 'fused_recurrent':
+            o, recurrent_state = fused_recurrent_delta_rule(
+                q=q,
+                k=k,
+                v=v,
+                beta=beta,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+                head_first=False,
+                use_qk_l2norm_in_kernel=True if self.qk_norm == 'l2' else False
+            )
+        elif mode == 'chunk':
+            o, recurrent_state = chunk_delta_rule(
+                q=q,
+                k=k,
+                v=v,
+                beta=beta,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+                head_first=False,
+                use_qk_l2norm_in_kernel=True if self.qk_norm == 'l2' else False
+            )
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        if past_key_values is not None:
+            past_key_values.update(
+                recurrent_state=recurrent_state,
+                conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None,
+                layer_idx=self.layer_idx,
+                offset=q.shape[1]
+            )
+
+        if self.use_gate:
+            g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim)
+            o = self.o_norm(o, g)
+        else:
+            o = self.o_norm(o)
+        o = rearrange(o, 'b t h d -> b t (h d)')
+        o = self.o_proj(o)
+
+        if self.weight_norm:
+            o = o / self.o_proj.weight.norm(p=2, dim=0)
+
+        return o, None, past_key_values

gated_deltanet.py

@@ -0,0 +1,333 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING, Dict, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from torch.nn import functional as F
+
+from fla.modules import FusedRMSNormSwishGate, RMSNorm, ShortConvolution
+from fla.ops.gated_delta_rule import (chunk_gated_delta_rule,
+                                      fused_recurrent_gated_delta_rule)
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import Unpack
+
+    from fla.models.utils import Cache
+
+
+def elu_p1(x):
+    return (F.elu(x, 1., False) + 1.).to(x)
+
+
+def sum_norm(x):
+    return (x / x.sum(-1, keepdim=True)).to(x)
+
+# https://github.com/IDSIA/recurrent-fwp/blob/master/algorithmic/layers.py#L86C1-L146C1
+
+
+class GatedDeltaNet(nn.Module):
+    """
+    The layer implementaion for [Gated Delta Networks: Improving Mamba2 with Delta Rule](https://arxiv.org/abs/2412.06464).  # noqa
+
+    Similar to Mamba2, each layer contains around 6*hidden_size*hidden_size parameters.
+    Parameter alloation when use_gate=True:
+        - 0.75 * hidden_size * hidden_size for the q_proj and k_proj each
+        - 1.5 * hidden_size * hidden_size for the v_proj, g_proj and o_proj each
+        - Others are ignorably small.
+        - In total = 0.75 * 2 + 1.5 * 3 = 6 * hidden_size * hidden_size
+    NOTE: num_heads * head_dim = 0.75 * hidden_size, please make sure to set the correct num_heads and head_dim.
+
+    Parameter allocation when use_gate=False:
+        - 1 * hidden_size * hidden_size for the q_proj and k_proj each
+        - 2 * hidden_size * hidden_size for the v_proj and o_proj each
+        - Others are ignorably small.
+        - In total = 1 * 2 + 2 * 2 = 6 * hidden_size * hidden_size
+
+    Args:
+        hidden_size (int, Optional):
+            The hidden size of the input. Default: 2048.
+        expand_v (float, Optional):
+            The expansion ratio for the value dim. Default: 2.0.
+        head_dim (int, Optional):
+            The dimension of each head. Default: 256.
+        num_heads (int, Optional):
+            The number of heads. Default: 4.
+        mode (str, Optional):
+            Which Gated DeltaNet kernel to use.
+            Currently available: `chunk` and `fused_recurrent`.
+            Default: `chunk`.
+        use_beta (bool, Optional):
+            Whether to use beta. Default: `True`.
+        use_gate (bool, Optional):
+            Whether to use output gate. Default: `True`.
+        use_short_conv (bool, Optional):
+            Whether to use short convolutions. Default: `True`.
+        conv_size (int, Optional):
+            The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4.
+        conv_bias (bool, Optional):
+            Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`.
+        layer_idx (int, Optional):
+            The index of the layer. Default: None.
+        norm_eps (float, Optional):
+            The epsilon value for the normalization layer. Default: 1e-5.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        expand_v: float = 2,
+        head_dim: int = 256,
+        num_heads: int = 6,
+        mode: str = 'chunk',
+        use_gate: bool = True,
+        use_short_conv: bool = True,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        layer_idx: int = None,
+        norm_eps: float = 1e-5,
+        config = None,
+        **kwargs
+    ) -> GatedDeltaNet:
+        super().__init__()
+        
+        self.config = config
+        if self.config is not None and self.config.use_nGPT and 'extra_grad' in self.config.nGPT_config and self.config.nGPT_config['extra_grad']:
+            self.weight_norm = True
+        else:
+            self.weight_norm = False
+
+        self.mode = mode
+
+        self.hidden_size = hidden_size
+        self.expand_v = expand_v
+
+        self.use_gate = use_gate
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+
+        self.head_dim = head_dim
+        self.num_heads = num_heads
+
+        self.key_dim = self.num_heads * self.head_dim
+        self.value_dim = self.key_dim * self.expand_v
+        self.head_k_dim = head_dim
+        self.head_v_dim = head_dim * self.expand_v
+        self.layer_idx = layer_idx
+        self.silu = nn.SiLU()
+
+        assert mode in ['chunk', 'fused_recurrent'], f"Not suppoerted mode `{mode}`."
+
+        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+        self.b_proj = nn.Linear(hidden_size, self.num_heads, bias=False)
+        self.a_proj = nn.Linear(hidden_size, self.num_heads, bias=False)
+        A = torch.empty(self.num_heads, dtype=torch.float32).uniform_(0, 16)
+        A_log = torch.log(A)
+        self.A_log = nn.Parameter(A_log)
+        self.A_log._no_weight_decay = True
+        self.D = nn.Parameter(torch.ones(self.num_heads))
+        self.D._no_weight_decay = True
+        # hard coded for now
+        dt_min = 0.001
+        dt_max = 0.1
+        dt_init_floor = 1e-4
+        dt = torch.exp(
+            torch.rand(self.num_heads) * (math.log(dt_max) - math.log(dt_min))
+            + math.log(dt_min)
+        )
+        dt = torch.clamp(dt, min=dt_init_floor)
+        # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
+        inv_dt = dt + torch.log(-torch.expm1(-dt))
+        self.dt_bias = nn.Parameter(inv_dt)
+        # Just to be explicit. Without this we already don't put wd on dt_bias because of the check
+        # name.endswith("bias") in param_grouping.py
+        self.dt_bias._no_weight_decay = True
+
+        if use_short_conv:
+            self.conv_size = conv_size
+            self.q_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                activation='silu'
+            )
+            self.k_conv1d = ShortConvolution(
+                hidden_size=self.key_dim,
+                kernel_size=conv_size,
+                activation='silu'
+            )
+            self.v_conv1d = ShortConvolution(
+                hidden_size=self.value_dim,
+                kernel_size=conv_size,
+                activation='silu'
+            )
+        else:
+            raise UserWarning(
+                "ShortConvolution is crucial to the performance. "
+                "Do not turn it off, i.e., setting `use_short_conv=False` unless you know what you are doing."
+            )
+        if use_gate:
+            self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+            self.o_norm = FusedRMSNormSwishGate(self.head_v_dim, eps=norm_eps)
+        else:
+            self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps)
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+        self.apply(self._initialize_weights)
+
+    def _initialize_weights(self, module: nn.Module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        module._is_hf_initialized = True
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+        **kwargs: Unpack[Dict]
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
+        if attention_mask is not None:
+            assert len(attention_mask.shape) == 2, (
+                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
+                "for padding purposes (0 indicating padding). "
+                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
+            )
+
+        mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode
+        if self.training:
+            assert mode == 'chunk', "Only chunk mode is supported in training."
+
+        last_state = None
+        if past_key_values is not None and len(past_key_values) > self.layer_idx:
+            last_state = past_key_values[self.layer_idx]
+
+        if self.use_short_conv:
+            conv_state_q, conv_state_k, conv_state_v = None, None, None
+            if last_state is not None:
+                conv_state_q, conv_state_k, conv_state_v = last_state['conv_state']
+            conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None
+            position_ids = kwargs.get('position_ids', None)
+
+            q = self.q_proj(hidden_states)
+            if self.weight_norm:
+                q = q / self.q_proj.weight.norm(p=2, dim=1)
+            q, conv_state_q = self.q_conv1d(x=q,
+                                            mask=conv_mask,
+                                            cache=conv_state_q,
+                                            output_final_state=use_cache,
+                                            seq_idx=position_ids)
+            
+            k = self.k_proj(hidden_states)
+            if self.weight_norm:
+                k = k / self.k_proj.weight.norm(p=2, dim=1)
+            k, conv_state_k = self.k_conv1d(x=k,
+                                            mask=conv_mask,
+                                            cache=conv_state_k,
+                                            output_final_state=use_cache,
+                                            seq_idx=position_ids)
+
+            v = self.v_proj(hidden_states)
+            if self.weight_norm:
+                v = v / self.v_proj.weight.norm(p=2, dim=1)
+            v, conv_state_v = self.v_conv1d(x=v,
+                                            mask=conv_mask,
+                                            cache=conv_state_v,
+                                            output_final_state=use_cache,
+                                            seq_idx=position_ids)
+            
+        else:
+            q = self.q_proj(hidden_states)
+            k = self.k_proj(hidden_states)
+            v = self.v_proj(hidden_states)
+
+            if self.weight_norm:
+                q = q / self.q_proj.weight.norm(p=2, dim=1)
+                k = k / self.k_proj.weight.norm(p=2, dim=1)
+                v = v / self.v_proj.weight.norm(p=2, dim=1)
+                
+            q, k, v = self.silu(q), self.silu(k), self.silu(v)
+            
+        q, k = map(lambda x: rearrange(x, 'b t (h d) -> b t h d', d=self.head_k_dim), (q, k))
+        v = rearrange(v, 'b t (h d) -> b t h d', d=self.head_v_dim)
+        
+        beta = self.b_proj(hidden_states)
+        if self.weight_norm:
+            beta = beta / self.b_proj.weight.norm(p=2, dim=1)
+        beta = beta.sigmoid()
+        
+        a_val = self.a_proj(hidden_states)
+        if self.weight_norm:
+            a_val = a_val / self.a_proj.weight.norm(p=2, dim=1)
+        g = -self.A_log.float().exp() * F.softplus(a_val.float() + self.dt_bias)
+
+        # dealing with padding
+        if attention_mask is not None:
+            beta = beta.mul(attention_mask[:, -beta.shape[-2]:, None])
+            g = g.mul(attention_mask[:, -g.shape[-2]:, None])
+
+        recurrent_state = last_state['recurrent_state'] if last_state is not None else None
+        cu_seqlens = kwargs.get('cu_seqlens', None)
+        if mode == 'chunk':
+            o, recurrent_state = chunk_gated_delta_rule(
+                q=q,
+                k=k,
+                v=v,
+                g=g,
+                beta=beta,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+                head_first=False,
+                use_qk_l2norm_in_kernel=True
+            )
+        elif mode == 'fused_recurrent':
+            o, recurrent_state = fused_recurrent_gated_delta_rule(
+                q=q,
+                k=k,
+                v=v,
+                g=g,
+                beta=beta,
+                initial_state=recurrent_state,
+                output_final_state=use_cache,
+                cu_seqlens=cu_seqlens,
+                head_first=False,
+                use_qk_l2norm_in_kernel=True
+            )
+        if past_key_values is not None:
+            past_key_values.update(
+                recurrent_state=recurrent_state,
+                conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None,
+                layer_idx=self.layer_idx,
+                offset=q.shape[1]
+            )
+
+        if self.use_gate:
+            gate_val = self.g_proj(hidden_states)
+            
+            # if self.weight_norm:
+            #     gate_val = gate_val / self.g_proj.weight.norm(p=2, dim=1)
+            
+            g = rearrange(gate_val, '... (h d) -> ... h d', d=self.head_v_dim)
+            o = self.o_norm(o, g)
+        else:
+            o = self.o_norm(o)
+        o = rearrange(o, 'b t h d -> b t (h d)')
+        o = self.o_proj(o)
+
+        if self.weight_norm:
+            o = o / self.o_proj.weight.norm(p=2, dim=0)
+            
+        return o, None, past_key_values

generation_config.json

@@ -0,0 +1,8 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 0,
+  "transformers_version": "4.48.2",
+  "use_cache": false
+}

mamba2.py

@@ -0,0 +1,1427 @@
+# Copyright (c) 2024, Tri Dao, Albert Gu.
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from einops import rearrange, repeat, pack, unpack
+
+try:
+    from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
+except ImportError:
+    causal_conv1d_fn, causal_conv1d_update = None, None
+
+try:
+    from causal_conv1d.causal_conv1d_varlen import causal_conv1d_varlen_states
+except ImportError:
+    causal_conv1d_varlen_states = None
+
+import sys
+# sys.path.insert(0, '/lustre/fsw/portfolios/nvr/users/yongganf/TLM/')
+
+from mamba_ssm.ops.triton.selective_state_update import selective_state_update
+from mamba_ssm.ops.triton.layernorm_gated import RMSNorm as RMSNormGated
+
+
+from mamba_ssm.distributed.tensor_parallel import ColumnParallelLinear, RowParallelLinear
+from mamba_ssm.distributed.distributed_utils import all_reduce, reduce_scatter
+
+from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined
+from mamba_ssm.ops.triton.ssd_combined import mamba_split_conv1d_scan_combined
+
+
+class Mamba2(nn.Module):
+    def __init__(
+        self,
+        config,
+        conv_init=None,
+        d_ssm=None,  # If not None, we only apply SSM on this many dimensions, the rest uses gated MLP
+        ngroups=1,
+        A_init_range=(1, 16),
+        D_has_hdim=False,
+        rmsnorm=True,
+        norm_before_gate=False,
+        dt_min=0.001,
+        dt_max=0.1,
+        dt_init_floor=1e-4,
+        dt_limit=(0.0, float("inf")),
+        bias=False,
+        conv_bias=True,
+        # Fused kernel and sharding options
+        chunk_size=256,
+        use_mem_eff_path=False, # True,
+        layer_idx=None,  # Absorb kwarg for general module
+        process_group=None,
+        sequence_parallel=True,
+        device=None,
+        dtype=None,
+    ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        
+        self.config = config
+        self.d_model = config.hidden_size
+        self.d_state = config.mamba_d_state
+        self.d_conv = config.mamba_d_conv
+
+        self.conv_init = conv_init
+        self.expand = config.mamba_expand
+        self.process_group = process_group
+        self.sequence_parallel = sequence_parallel
+        self.world_size = 1 if process_group is None else process_group.size()
+        self.local_rank = 0 if process_group is None else process_group.rank()
+        self.d_inner = (self.expand * self.d_model) // self.world_size
+        assert self.d_inner * self.world_size == self.expand * self.d_model
+        self.headdim = config.mamba2_headdim
+        self.d_ssm = self.d_inner if d_ssm is None else d_ssm // self.world_size
+        assert ngroups % self.world_size == 0
+        self.ngroups = ngroups // self.world_size
+        assert self.d_ssm % self.headdim == 0
+        self.nheads = self.d_ssm // self.headdim
+        self.D_has_hdim = D_has_hdim
+        self.rmsnorm = rmsnorm
+        self.norm_before_gate = norm_before_gate
+        self.dt_limit = dt_limit
+        self.activation = "silu"
+        self.chunk_size = chunk_size
+        self.use_mem_eff_path = use_mem_eff_path
+        self.layer_idx = layer_idx
+
+        assert (self.d_model * self.expand / self.headdim) % 8 == 0
+
+        # Order: [z, x, B, C, dt]
+        d_in_proj = 2 * self.d_inner + 2 * self.ngroups * self.d_state + self.nheads
+        if self.process_group is None:
+            self.in_proj = nn.Linear(self.d_model, d_in_proj, bias=bias, **factory_kwargs)
+        else:
+            self.in_proj = ColumnParallelLinear(self.d_model, d_in_proj * self.world_size, bias=bias,
+                                                process_group=self.process_group, sequence_parallel=self.sequence_parallel,
+                                                **factory_kwargs)
+
+        conv_dim = self.d_ssm + 2 * self.ngroups * self.d_state
+        self.conv1d = nn.Conv1d(
+            in_channels=conv_dim,
+            out_channels=conv_dim,
+            bias=conv_bias,
+            kernel_size=self.d_conv,
+            groups=conv_dim,
+            padding=self.d_conv - 1,
+            **factory_kwargs,
+        )
+        if self.conv_init is not None:
+            nn.init.uniform_(self.conv1d.weight, -self.conv_init, self.conv_init)
+
+        self.act = nn.SiLU()
+
+        # Initialize log dt bias
+        dt = torch.exp(
+            torch.rand(self.nheads, **factory_kwargs) * (math.log(dt_max) - math.log(dt_min))
+            + math.log(dt_min)
+        )
+        dt = torch.clamp(dt, min=dt_init_floor)
+        # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
+        inv_dt = dt + torch.log(-torch.expm1(-dt))
+        
+        if config.no_dt_bias:
+            self.dt_bias = None
+        else:
+            self.dt_bias = nn.Parameter(inv_dt)
+            # Just to be explicit. Without this we already don't put wd on dt_bias because of the check
+            # name.endswith("bias") in param_grouping.py
+            self.dt_bias._no_weight_decay = True
+
+        assert A_init_range[0] > 0 and A_init_range[1] >= A_init_range[0]
+        A = torch.empty(self.nheads, dtype=torch.float32, device=device).uniform_(*A_init_range)
+        A_log = torch.log(A).to(dtype=dtype)
+        self.A_log = nn.Parameter(A_log)
+        self.A_log._no_weight_decay = True
+
+        # D "skip" parameter
+        self.D = nn.Parameter(torch.ones(self.d_ssm if self.D_has_hdim else self.nheads, device=device))
+        self.D._no_weight_decay = True
+
+        if self.rmsnorm:
+            assert RMSNormGated is not None
+            self.norm = RMSNormGated(self.d_ssm, eps=1e-5, norm_before_gate=self.norm_before_gate,
+                                     group_size=self.d_ssm // ngroups, **factory_kwargs)
+
+        if self.process_group is None:
+            self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=bias, **factory_kwargs)
+        else:
+            self.out_proj = RowParallelLinear(self.d_inner * self.world_size, self.d_model, bias=bias,
+                                              process_group=self.process_group, sequence_parallel=self.sequence_parallel,
+                                              **factory_kwargs)
+
+        self.mamba_multihead_config = config.mamba_multihead_config
+        if self.mamba_multihead_config is not None:
+            assert self.mamba_multihead_config['alpha_mode'] == 'sparsity' or self.mamba_multihead_config['alpha_mode'] == 'cummax'
+            
+            if self.mamba_multihead_config['alpha_mode'] == 'cummax':
+                self.learned_dt_scale = nn.Parameter(torch.ones(1, device=device))
+            
+            if self.mamba_multihead_config['alpha_mode'] == 'sparsity':
+                if 'use_learned_thres' in self.mamba_multihead_config and self.mamba_multihead_config['use_learned_thres']:
+                    self.learned_thres = nn.Parameter(torch.zeros(self.nheads, device=device))
+                    self.smooth_factor = self.mamba_multihead_config['smooth_factor']
+                    self.detach_dt = self.mamba_multihead_config['detach_dt']
+                    
+                    if 'use_cummax' in self.mamba_multihead_config and self.mamba_multihead_config['use_cummax']:
+                        self.use_cummax = True
+                        self.cummax_lower_bound = self.mamba_multihead_config['cummax_lower_bound']
+                    else:
+                        self.use_cummax = False
+                    
+                else:
+                    self.learned_thres = None
+                    self.smooth_factor = None
+                    self.detach_dt = None
+
+                    self.sparsity_split = self.mamba_multihead_config['sparsity_split']
+                    self.sparsity_ratio = self.mamba_multihead_config['sparsity_ratio']
+
+        if self.config.layerwise_memory_token:
+            assert self.config.num_memory_tokens > 0
+            self.memory_tokens = nn.Parameter(torch.randn(self.config.num_memory_tokens, self.config.hidden_size))  
+        else:
+            self.memory_tokens = None
+            
+            
+    def forward(self, hidden_states, attention_mask=None, past_key_value=None, seqlen=None, seq_idx=None, cu_seqlens=None, inference_params=None):
+        """
+        hidden_states: (batch, seqlen, hidden_dim) if seqlen=None.
+            If seqlen is not None, hidden_states is (batch * seqlen, hidden_dim). This is so that when we
+            split hidden_states during sequence parallel, we split the batch * seqlen dimension
+            (in case batch is small).
+        Returns: same shape as u
+        """
+        # assert past_key_value is None, "Not implemented yet!!!"
+
+        if self.memory_tokens is not None:
+            hidden_states = hidden_states[:,self.config.num_memory_tokens:,...]
+            mem = repeat(self.memory_tokens, 'n d -> b n d', b = hidden_states.shape[0]) # prepend the memory to every segment of m by repeating the memory tokens
+            hidden_states, mem_packed_shape = pack((mem, hidden_states), 'b * d')  
+            
+        seqlen_og = seqlen
+        if seqlen is None:
+            batch, seqlen, dim = hidden_states.shape
+        else:
+            batch_seqlen, dim = hidden_states.shape
+            batch = batch_seqlen // seqlen
+
+        conv_state, ssm_state = None, None
+        if inference_params is not None:
+            inference_batch = cu_seqlens.shape[0] - 1 if cu_seqlens is not None else batch
+            conv_state, ssm_state = self._get_states_from_cache(inference_params, inference_batch)
+            if inference_params.seqlen_offset > 0:
+                # The states are updated inplace
+                out, _, _ = self.step(hidden_states, conv_state, ssm_state)
+                return out
+
+        zxbcdt = self.in_proj(hidden_states)  # (B, L, d_in_proj) or (B * L, d_in_proj)
+
+        if self.config.use_nGPT and 'extra_grad' in self.config.nGPT_config and self.config.nGPT_config['extra_grad']:
+            zxbcdt = zxbcdt / self.in_proj.weight.norm(p=2, dim=1)
+            
+        if seqlen_og is not None:
+            zxbcdt = rearrange(zxbcdt, "(b l) d -> b l d", l=seqlen)
+        # If the model is loaded in fp16, without the .float() here, A might be -inf
+        A = -torch.exp(self.A_log.float())  # (nheads) or (d_inner, d_state)
+        dt_limit_kwargs = {} if self.dt_limit == (0.0, float("inf")) else dict(dt_limit=self.dt_limit)
+        if self.use_mem_eff_path and inference_params is None:
+            out = mamba_split_conv1d_scan_combined(
+                zxbcdt,
+                rearrange(self.conv1d.weight, "d 1 w -> d w"),
+                self.conv1d.bias,
+                self.dt_bias,
+                A,
+                D=rearrange(self.D, "(h p) -> h p", p=self.headdim) if self.D_has_hdim else self.D,
+                chunk_size=self.chunk_size,
+                seq_idx=seq_idx,
+                activation=self.activation,
+                rmsnorm_weight=self.norm.weight if self.rmsnorm else None,
+                rmsnorm_eps=self.norm.eps if self.rmsnorm else 1e-6,
+                outproj_weight=self.out_proj.weight,
+                outproj_bias=self.out_proj.bias,
+                headdim=None if self.D_has_hdim else self.headdim,
+                ngroups=self.ngroups,
+                norm_before_gate=self.norm_before_gate,
+                **dt_limit_kwargs,
+            )
+            if seqlen_og is not None:
+                out = rearrange(out, "b l d -> (b l) d")
+            if self.process_group is not None:
+                reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce
+                out = reduce_fn(out, self.process_group)
+        else:
+            d_mlp = (zxbcdt.shape[-1] - 2 * self.d_ssm - 2 * self.ngroups * self.d_state - self.nheads) // 2
+            z0, x0, z, xBC, dt = torch.split(
+                zxbcdt,
+                [d_mlp, d_mlp, self.d_ssm, self.d_ssm + 2 * self.ngroups * self.d_state, self.nheads],
+                dim=-1
+            )
+            if conv_state is not None:
+                if cu_seqlens is None:
+                    # If we just take xBC[:, :, -self.d_conv :], it will error if seqlen < self.d_conv
+                    # Instead F.pad will pad with zeros if seqlen < self.d_conv, and truncate otherwise.
+                    xBC_t = rearrange(xBC, "b l d -> b d l")
+                    conv_state.copy_(F.pad(xBC_t, (self.d_conv - xBC_t.shape[-1], 0)))  # Update state (B D W)
+                else:
+                    assert causal_conv1d_varlen_states is not None, "varlen inference requires causal_conv1d package"
+                    assert batch == 1, "varlen inference only supports batch dimension 1"
+                    conv_varlen_states = causal_conv1d_varlen_states(
+                        xBC.squeeze(0), cu_seqlens, state_len=conv_state.shape[-1]
+                    )
+                    conv_state.copy_(conv_varlen_states)
+            assert self.activation in ["silu", "swish"]
+            if causal_conv1d_fn is None or self.activation not in ["silu", "swish"]:
+                assert seq_idx is None, "varlen conv1d requires the causal_conv1d package"
+                xBC = self.act(
+                    self.conv1d(xBC.transpose(1, 2)).transpose(1, 2)[:, -(self.dconv - 1):]
+                )  # (B, L, self.d_ssm + 2 * ngroups * d_state)
+            else:
+                xBC = causal_conv1d_fn(
+                    xBC.transpose(1, 2),
+                    rearrange(self.conv1d.weight, "d 1 w -> d w"),
+                    bias=self.conv1d.bias,
+                    activation=self.activation,
+                    # seq_idx=seq_idx,
+                ).transpose(1, 2)
+            x, B, C = torch.split(xBC, [self.d_ssm, self.ngroups * self.d_state, self.ngroups * self.d_state], dim=-1)
+            
+            no_dt_bias = False
+            if self.mamba_multihead_config is not None and self.mamba_multihead_config['alpha_mode'] == 'cummax':  ### todo: implement this in the fused kernel
+                dt = dt + self.dt_bias
+                dt = torch.nn.functional.softmax(dt, dim=-1)
+                dt = torch.cumsum(dt, dim=-1)
+                dt = dt * self.learned_dt_scale
+                
+                no_dt_bias = True
+        
+            if self.mamba_multihead_config is not None and self.mamba_multihead_config['alpha_mode'] == 'sparsity':
+                dt = dt + self.dt_bias
+                
+                if self.learned_thres is not None:
+                    dt = self.sparsify_learned_thres(dt)
+                else:
+                    dt = self.split_and_sparsify(dt, self.sparsity_split, self.sparsity_ratio)
+                    
+                no_dt_bias = True
+
+
+            y = mamba_chunk_scan_combined(
+                rearrange(x, "b l (h p) -> b l h p", p=self.headdim),
+                dt,
+                A,
+                rearrange(B, "b l (g n) -> b l g n", g=self.ngroups),
+                rearrange(C, "b l (g n) -> b l g n", g=self.ngroups),
+                chunk_size=self.chunk_size,
+                D=rearrange(self.D, "(h p) -> h p", p=self.headdim) if self.D_has_hdim else self.D,
+                z=rearrange(z, "b l (h p) -> b l h p", p=self.headdim) if not self.rmsnorm else None,
+                dt_bias=self.dt_bias if not no_dt_bias else None,
+                dt_softplus=True,
+                seq_idx=seq_idx,
+                cu_seqlens=cu_seqlens,
+                **dt_limit_kwargs,
+                return_final_states=ssm_state is not None,
+                return_varlen_states=cu_seqlens is not None and inference_params is not None,
+            )
+            if ssm_state is not None:
+                y, last_state, *rest = y
+                if cu_seqlens is None:
+                    ssm_state.copy_(last_state)
+                else:
+                    varlen_states = rest[0]
+                    ssm_state.copy_(varlen_states)
+            y = rearrange(y, "b l h p -> b l (h p)")
+            if self.rmsnorm:
+                y = self.norm(y, z)
+            if d_mlp > 0:
+                y = torch.cat([F.silu(z0) * x0, y], dim=-1)
+            if seqlen_og is not None:
+                y = rearrange(y, "b l d -> (b l) d")
+                
+            if self.config.use_nGPT and 'extra_grad' in self.config.nGPT_config and self.config.nGPT_config['extra_grad']:
+                y = y / self.out_proj.weight.norm(p=2, dim=0)
+                
+            out = self.out_proj(y)
+                
+        return out, past_key_value
+
+
+    def sparsify_learned_thres(self, dt):
+        """
+        Args:
+            dt: Tensor of shape [bs, seq_len, nheads]
+        Returns:
+            pruned_dt: Pruned tensor with the same shape as dt
+        """
+        # Compute sigmoid scores
+        
+        if self.use_cummax:
+            learned_thres = torch.nn.functional.softmax(self.learned_thres, dim=-1)
+            learned_thres = torch.cumsum(learned_thres, dim=-1) - self.cummax_lower_bound  ## keep the dt_normalized larger than 1 - self.cummax_lower_bound
+            
+            dt_normalized = (dt - dt.min(dim=-1, keepdim=True)[0]) / (dt.max(dim=-1, keepdim=True)[0] - dt.min(dim=-1, keepdim=True)[0])
+            
+            scores = torch.sigmoid((dt_normalized.detach() - self.learned_thres) / self.smooth_factor)
+                    
+        else:
+            if self.detach_dt:
+                scores = torch.sigmoid((dt.detach() - self.learned_thres) / self.smooth_factor)
+            else:
+                scores = torch.sigmoid((dt - self.learned_thres) / self.smooth_factor)
+        
+        # Generate binary mask for pruning (forward pass)
+        mask = (scores >= 0.5).float()
+        
+        # Apply mask in the forward pass and backward using sigmoid
+        pruned_dt = (dt * mask - dt * scores).detach() + dt * scores
+        
+        # print(pruned_dt.mean())
+        
+        return pruned_dt
+    
+    
+    def split_and_sparsify(self, dt, sparsity_split, sparsity_ratio):
+        """
+        dt:             a torch.Tensor of shape [bs, seq_len, dim]
+        sparsity_split: list of ratios (e.g., [0.4, 0.3, 0.3]) that sum to 1
+                        and define how to split dt along the last dimension
+        sparsity_ratio: list of ratios (e.g., [0.2, 0.5, 0.3]) that sum to 1
+                        and define how many time steps (along seq_len) to keep
+        """
+        bs, seq_len, dim = dt.shape
+        
+        assert sum(sparsity_split) == 1
+        
+        # Compute the exact split sizes (watching out for integer rounding)
+        split_sizes = [int(r * dim) for r in sparsity_split]
+        # Fix potential off-by-one rounding in the last split
+        split_sizes[-1] = dim - sum(split_sizes[:-1])
+
+        # Split the original tensor along the last dimension
+        splitted_tensors = torch.split(dt, split_sizes, dim=-1)
+
+        results = []
+        for i, sub_tensor in enumerate(splitted_tensors):
+            # sub_tensor has shape [bs, seq_len, split_dim_i]
+            k = int(sparsity_ratio[i] * seq_len)
+            
+            ### Strategy 1: keep at least one token
+            k = max(k, 1)
+            
+            ### Strategy 2: the #tokens is the same as training
+            # if self.config.orig_max_position_embeddings is not None:
+            #     k = int(self.config.orig_max_position_embeddings * self.sparsity_ratio[i])
+            # else:
+            #     assert self.config.max_position_embeddings is not None
+            #     k = int(self.config.max_position_embeddings * self.sparsity_ratio[i])
+            
+            # k = min(seq_len, k)
+            
+            # print(self.config.max_position_embeddings, sparsity_ratio[i], seq_len, k)
+                    
+            # 1) Average over the feature dimension (the last dim), 
+            #    resulting in shape [bs, seq_len]
+            averaged_values = sub_tensor.mean(dim=-1)  
+
+            # 2) Get top-k indices (along seq_len = dim=1)
+            topk_values, _ = torch.topk(averaged_values, k=k, dim=1)
+            # The smallest value among the top-k per batch element
+            threshold = topk_values[:, -1].unsqueeze(-1)  # shape [bs, 1]
+
+            # 3) Create a mask of shape [bs, seq_len] => True if >= threshold
+            averaged_mask = (averaged_values >= threshold)
+
+            # 4) Expand that mask back to [bs, seq_len, split_dim_i]
+            mask_3d = averaged_mask.unsqueeze(-1).expand_as(sub_tensor)
+
+            # 5) Zero out everything that is not in top-k
+            sparsified_sub = sub_tensor * mask_3d
+            
+            # print((sparsified_sub == 0).float().mean().item())
+            # input()
+
+            results.append(sparsified_sub)
+        
+        # Concatenate the results back along the last dimension
+        output = torch.cat(results, dim=-1)
+        return output
+
+    def step(self, hidden_states, conv_state, ssm_state):
+        dtype = hidden_states.dtype
+        assert hidden_states.shape[1] == 1, "Only support decoding with 1 token at a time for now"
+        zxbcdt = self.in_proj(hidden_states.squeeze(1))  # (B 2D)
+        d_mlp = (zxbcdt.shape[-1] - 2 * self.d_ssm - 2 * self.ngroups * self.d_state - self.nheads) // 2
+        z0, x0, z, xBC, dt = torch.split(
+            zxbcdt,
+            [d_mlp, d_mlp, self.d_ssm, self.d_ssm + 2 * self.ngroups * self.d_state, self.nheads],
+            dim=-1
+        )
+
+        # Conv step
+        if causal_conv1d_update is None:
+            conv_state.copy_(torch.roll(conv_state, shifts=-1, dims=-1))  # Update state (B D W)
+            conv_state[:, :, -1] = xBC
+            xBC = torch.sum(conv_state * rearrange(self.conv1d.weight, "d 1 w -> d w"), dim=-1)  # (B D)
+            if self.conv1d.bias is not None:
+                xBC = xBC + self.conv1d.bias
+            xBC = self.act(xBC).to(dtype=dtype)
+        else:
+            xBC = causal_conv1d_update(
+                xBC,
+                conv_state,
+                rearrange(self.conv1d.weight, "d 1 w -> d w"),
+                self.conv1d.bias,
+                self.activation,
+            )
+
+        x, B, C = torch.split(xBC, [self.d_ssm, self.ngroups * self.d_state, self.ngroups * self.d_state], dim=-1)
+        A = -torch.exp(self.A_log.float())  # (nheads,)
+
+        # SSM step
+        if selective_state_update is None:
+            assert self.ngroups == 1, "Only support ngroups=1 for this inference code path"
+            # Discretize A and B
+            dt = F.softplus(dt + self.dt_bias.to(dtype=dt.dtype))  # (batch, nheads)
+            dA = torch.exp(dt * A)  # (batch, nheads)
+            x = rearrange(x, "b (h p) -> b h p", p=self.headdim)
+            dBx = torch.einsum("bh,bn,bhp->bhpn", dt, B, x)
+            ssm_state.copy_(ssm_state * rearrange(dA, "b h -> b h 1 1") + dBx)
+            y = torch.einsum("bhpn,bn->bhp", ssm_state.to(dtype), C)
+            y = y + rearrange(self.D.to(dtype), "h -> h 1") * x
+            y = rearrange(y, "b h p -> b (h p)")
+            if not self.rmsnorm:
+                y = y * self.act(z)  # (B D)
+        else:
+            A = repeat(A, "h -> h p n", p=self.headdim, n=self.d_state).to(dtype=torch.float32)
+            dt = repeat(dt, "b h -> b h p", p=self.headdim)
+            dt_bias = repeat(self.dt_bias, "h -> h p", p=self.headdim)
+            D = repeat(self.D, "h -> h p", p=self.headdim)
+            B = rearrange(B, "b (g n) -> b g n", g=self.ngroups)
+            C = rearrange(C, "b (g n) -> b g n", g=self.ngroups)
+            x_reshaped = rearrange(x, "b (h p) -> b h p", p=self.headdim)
+            if not self.rmsnorm:
+                z = rearrange(z, "b (h p) -> b h p", p=self.headdim)
+            y = selective_state_update(
+                ssm_state, x_reshaped, dt, A, B, C, D, z=z if not self.rmsnorm else None,
+                dt_bias=dt_bias, dt_softplus=True
+            )
+            y = rearrange(y, "b h p -> b (h p)")
+        if self.rmsnorm:
+            y = self.norm(y, z)
+        if d_mlp > 0:
+            y = torch.cat([F.silu(z0) * x0, y], dim=-1)
+        out = self.out_proj(y)
+        return out.unsqueeze(1), conv_state, ssm_state
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        device = self.out_proj.weight.device
+        conv_dtype = self.conv1d.weight.dtype if dtype is None else dtype
+        conv_state = torch.zeros(
+            batch_size, self.d_conv, self.conv1d.weight.shape[0], device=device, dtype=conv_dtype
+        ).transpose(1, 2)
+        ssm_dtype = self.in_proj.weight.dtype if dtype is None else dtype
+        ssm_state = torch.zeros(
+            batch_size, self.nheads, self.headdim, self.d_state, device=device, dtype=ssm_dtype
+        )
+        return conv_state, ssm_state
+
+    def _get_states_from_cache(self, inference_params, batch_size, initialize_states=False):
+        assert self.layer_idx is not None
+        if self.layer_idx not in inference_params.key_value_memory_dict:
+            batch_shape = (batch_size,)
+            conv_state = torch.zeros(
+                batch_size,
+                self.d_conv,
+                self.conv1d.weight.shape[0],
+                device=self.conv1d.weight.device,
+                dtype=self.conv1d.weight.dtype,
+            ).transpose(1, 2)
+            ssm_state = torch.zeros(
+                batch_size,
+                self.nheads,
+                self.headdim,
+                self.d_state,
+                device=self.in_proj.weight.device,
+                dtype=self.in_proj.weight.dtype,
+            )
+            inference_params.key_value_memory_dict[self.layer_idx] = (conv_state, ssm_state)
+        else:
+            conv_state, ssm_state = inference_params.key_value_memory_dict[self.layer_idx]
+            # TODO: What if batch size changes between generation, and we reuse the same states?
+            if initialize_states:
+                conv_state.zero_()
+                ssm_state.zero_()
+        return conv_state, ssm_state
+    
+
+class Mamba2_Fused(nn.Module):
+    def __init__(
+        self,
+        config,
+        layer_idx=None,  # Absorb kwarg for general module
+        reuse_kv=False,
+        conv_init=None,
+        d_ssm=None,  # If not None, we only apply SSM on this many dimensions, the rest uses gated MLP
+        ngroups=1,
+        A_init_range=(1, 16),
+        D_has_hdim=False,
+        rmsnorm=True,
+        norm_before_gate=False,
+        dt_min=0.001,
+        dt_max=0.1,
+        dt_init_floor=1e-4,
+        dt_limit=(0.0, float("inf")),
+        bias=False,
+        conv_bias=True,
+        # Fused kernel and sharding options
+        chunk_size=256,
+        use_mem_eff_path=False, # True,
+        process_group=None,
+        sequence_parallel=True,
+        device=None,
+        dtype=None,
+    ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        
+        self.config = config
+        self.d_model = config.hidden_size
+        self.d_state = config.mamba_d_state
+        self.d_conv = config.mamba_d_conv
+
+        self.conv_init = conv_init
+        self.expand = config.mamba_expand
+        self.process_group = process_group
+        self.sequence_parallel = sequence_parallel
+        self.world_size = 1 if process_group is None else process_group.size()
+        self.local_rank = 0 if process_group is None else process_group.rank()
+        self.d_inner = (self.expand * self.d_model) // self.world_size
+        assert self.d_inner * self.world_size == self.expand * self.d_model
+        self.headdim = config.mamba2_headdim
+        self.d_ssm = self.d_inner if d_ssm is None else d_ssm // self.world_size
+        assert ngroups % self.world_size == 0
+        self.ngroups = ngroups // self.world_size
+        assert self.d_ssm % self.headdim == 0
+        self.nheads = self.d_ssm // self.headdim
+        self.D_has_hdim = D_has_hdim
+        self.rmsnorm = rmsnorm
+        self.norm_before_gate = norm_before_gate
+        self.dt_limit = dt_limit
+        self.activation = "silu"
+        self.chunk_size = chunk_size
+        self.use_mem_eff_path = use_mem_eff_path
+        self.layer_idx = layer_idx
+
+        assert (self.d_model * self.expand / self.headdim) % 8 == 0
+
+        self.fused_multihead_config = config.fused_multihead_config
+        assert self.fused_multihead_config['expand_v'], "Only implemented Hymba for Mamba"
+        
+        self.reuse_kv = reuse_kv
+
+        self.hidden_size = config.hidden_size
+        self.attn_hidden_size = config.hidden_size
+        self.num_attention_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+
+        self.k_hidden_size = int(self.num_key_value_heads/self.num_attention_heads * self.attn_hidden_size)
+        self.v_hidden_size = int(self.num_key_value_heads/self.num_attention_heads * self.attn_hidden_size * self.expand) if self.fused_multihead_config['expand_v'] else int(self.num_key_value_heads/self.num_attention_heads * self.attn_hidden_size)
+    
+        if self.fused_multihead_config['expand_v']:
+            config.v_head_dim = self.d_inner // self.num_attention_heads
+            
+        self.self_attn = config.attn_op(config, layer_idx, attn_only_wo_proj=True, reuse_kv=reuse_kv)
+        
+        if self.reuse_kv:   # Order: [q, z, x, B, C, dt]
+            d_in_proj = self.attn_hidden_size + 2 * self.d_inner + 2 * self.ngroups * self.d_state + self.nheads
+        else:   # Order: [q, k, v, z, x, B, C, dt]
+            d_in_proj = self.attn_hidden_size + self.k_hidden_size + self.v_hidden_size + 2 * self.d_inner + 2 * self.ngroups * self.d_state + self.nheads
+
+        if self.process_group is None:
+            self.in_proj = nn.Linear(self.d_model, d_in_proj, bias=bias, **factory_kwargs)
+        else:
+            self.in_proj = ColumnParallelLinear(self.d_model, d_in_proj * self.world_size, bias=bias,
+                                                process_group=self.process_group, sequence_parallel=self.sequence_parallel,
+                                                **factory_kwargs)
+
+        self.pre_avg_layernorm1 = JambaRMSNorm(self.d_inner, eps=config.rms_norm_eps)
+        self.pre_avg_layernorm2 = JambaRMSNorm(self.d_inner, eps=config.rms_norm_eps)
+            
+        conv_dim = self.d_ssm + 2 * self.ngroups * self.d_state
+        self.conv1d = nn.Conv1d(
+            in_channels=conv_dim,
+            out_channels=conv_dim,
+            bias=conv_bias,
+            kernel_size=self.d_conv,
+            groups=conv_dim,
+            padding=self.d_conv - 1,
+            **factory_kwargs,
+        )
+        if self.conv_init is not None:
+            nn.init.uniform_(self.conv1d.weight, -self.conv_init, self.conv_init)
+
+        self.act = nn.SiLU()
+
+        # Initialize log dt bias
+        dt = torch.exp(
+            torch.rand(self.nheads, **factory_kwargs) * (math.log(dt_max) - math.log(dt_min))
+            + math.log(dt_min)
+        )
+        dt = torch.clamp(dt, min=dt_init_floor)
+        # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
+        inv_dt = dt + torch.log(-torch.expm1(-dt))
+        self.dt_bias = nn.Parameter(inv_dt)
+        # Just to be explicit. Without this we already don't put wd on dt_bias because of the check
+        # name.endswith("bias") in param_grouping.py
+        self.dt_bias._no_weight_decay = True
+
+        assert A_init_range[0] > 0 and A_init_range[1] >= A_init_range[0]
+        A = torch.empty(self.nheads, dtype=torch.float32, device=device).uniform_(*A_init_range)
+        A_log = torch.log(A).to(dtype=dtype)
+        self.A_log = nn.Parameter(A_log)
+        self.A_log._no_weight_decay = True
+
+        # D "skip" parameter
+        self.D = nn.Parameter(torch.ones(self.d_ssm if self.D_has_hdim else self.nheads, device=device))
+        self.D._no_weight_decay = True
+
+        if self.rmsnorm:
+            assert RMSNormGated is not None
+            self.norm = RMSNormGated(self.d_ssm, eps=1e-5, norm_before_gate=self.norm_before_gate,
+                                     group_size=self.d_ssm // ngroups, **factory_kwargs)
+
+        if self.process_group is None:
+            self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=bias, **factory_kwargs)
+        else:
+            self.out_proj = RowParallelLinear(self.d_inner * self.world_size, self.d_model, bias=bias,
+                                              process_group=self.process_group, sequence_parallel=self.sequence_parallel,
+                                              **factory_kwargs)
+
+    def forward(self, hidden_states, attention_mask=None, past_key_value=None, position_ids=None, kv_last_layer=None, use_cache=False, use_swa=False, seqlen=None, seq_idx=None, cu_seqlens=None, inference_params=None):
+        """
+        hidden_states: (batch, seqlen, hidden_dim) if seqlen=None.
+            If seqlen is not None, hidden_states is (batch * seqlen, hidden_dim). This is so that when we
+            split hidden_states during sequence parallel, we split the batch * seqlen dimension
+            (in case batch is small).
+        Returns: same shape as u
+        """
+        # assert past_key_value is None, "Not implemented yet!!!"
+        
+        seqlen_og = seqlen
+        if seqlen is None:
+            batch, seqlen, dim = hidden_states.shape
+        else:
+            batch_seqlen, dim = hidden_states.shape
+            batch = batch_seqlen // seqlen
+
+        conv_state, ssm_state = None, None
+        if inference_params is not None:
+            inference_batch = cu_seqlens.shape[0] - 1 if cu_seqlens is not None else batch
+            conv_state, ssm_state = self._get_states_from_cache(inference_params, inference_batch)
+            if inference_params.seqlen_offset > 0:
+                # The states are updated inplace
+                out, _, _ = self.step(hidden_states, conv_state, ssm_state)
+                return out
+
+        zxbcdt = self.in_proj(hidden_states)  # (B, L, d_in_proj) or (B * L, d_in_proj)
+        
+        if self.reuse_kv:
+            query_states, zxbcdt = zxbcdt.tensor_split((self.attn_hidden_size,), dim=-1)
+            # query_states = query_states.transpose(1,2)
+        else:
+            query_states, key_states, value_states, zxbcdt = zxbcdt.tensor_split((self.attn_hidden_size, self.attn_hidden_size + self.k_hidden_size, self.attn_hidden_size + self.k_hidden_size + self.v_hidden_size), dim=-1)
+
+            # query_states = query_states.transpose(1,2)
+            # key_states = key_states.transpose(1,2)
+            # value_states = value_states.transpose(1,2)
+                        
+        if self.reuse_kv:
+            assert kv_last_layer is not None
+            attn_outputs, attn_key_value = self.self_attn(attention_mask=attention_mask, position_ids=position_ids, query_states=query_states, kv_last_layer=kv_last_layer, use_swa=use_swa, use_cache=use_cache, past_key_value=past_key_value)
+        else:
+            if 'use_linear_attn' in self.fused_multihead_config and self.fused_multihead_config['use_linear_attn'] and self.linear_attn_op == 'gla':
+                attn_outputs, _, attn_key_value = self.self_attn(hidden_states=value_states, position_ids=position_ids, attention_mask=attention_mask, Q=query_states, K=key_states, V=value_states, past_key_value=past_key_value)
+            else:
+                attn_outputs, attn_key_value = self.self_attn(attention_mask=attention_mask, position_ids=position_ids, query_states=query_states, key_states=key_states, value_states=value_states, use_swa=use_swa, use_cache=use_cache, past_key_value=past_key_value)
+
+
+        if seqlen_og is not None:
+            zxbcdt = rearrange(zxbcdt, "(b l) d -> b l d", l=seqlen)
+        # If the model is loaded in fp16, without the .float() here, A might be -inf
+        A = -torch.exp(self.A_log.float())  # (nheads) or (d_inner, d_state)
+        dt_limit_kwargs = {} if self.dt_limit == (0.0, float("inf")) else dict(dt_limit=self.dt_limit)
+        if self.use_mem_eff_path and inference_params is None:
+            out = mamba_split_conv1d_scan_combined(
+                zxbcdt,
+                rearrange(self.conv1d.weight, "d 1 w -> d w"),
+                self.conv1d.bias,
+                self.dt_bias,
+                A,
+                D=rearrange(self.D, "(h p) -> h p", p=self.headdim) if self.D_has_hdim else self.D,
+                chunk_size=self.chunk_size,
+                seq_idx=seq_idx,
+                activation=self.activation,
+                rmsnorm_weight=self.norm.weight if self.rmsnorm else None,
+                rmsnorm_eps=self.norm.eps if self.rmsnorm else 1e-6,
+                outproj_weight=self.out_proj.weight,
+                outproj_bias=self.out_proj.bias,
+                headdim=None if self.D_has_hdim else self.headdim,
+                ngroups=self.ngroups,
+                norm_before_gate=self.norm_before_gate,
+                **dt_limit_kwargs,
+            )
+            if seqlen_og is not None:
+                out = rearrange(out, "b l d -> (b l) d")
+            if self.process_group is not None:
+                reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce
+                out = reduce_fn(out, self.process_group)
+        else:
+            d_mlp = (zxbcdt.shape[-1] - 2 * self.d_ssm - 2 * self.ngroups * self.d_state - self.nheads) // 2
+            
+            z0, x0, z, xBC, dt = torch.split(
+                zxbcdt,
+                [d_mlp, d_mlp, self.d_ssm, self.d_ssm + 2 * self.ngroups * self.d_state, self.nheads],
+                dim=-1
+            )
+            if conv_state is not None:
+                if cu_seqlens is None:
+                    # If we just take xBC[:, :, -self.d_conv :], it will error if seqlen < self.d_conv
+                    # Instead F.pad will pad with zeros if seqlen < self.d_conv, and truncate otherwise.
+                    xBC_t = rearrange(xBC, "b l d -> b d l")
+                    conv_state.copy_(F.pad(xBC_t, (self.d_conv - xBC_t.shape[-1], 0)))  # Update state (B D W)
+                else:
+                    assert causal_conv1d_varlen_states is not None, "varlen inference requires causal_conv1d package"
+                    assert batch == 1, "varlen inference only supports batch dimension 1"
+                    conv_varlen_states = causal_conv1d_varlen_states(
+                        xBC.squeeze(0), cu_seqlens, state_len=conv_state.shape[-1]
+                    )
+                    conv_state.copy_(conv_varlen_states)
+            assert self.activation in ["silu", "swish"]
+            if causal_conv1d_fn is None or self.activation not in ["silu", "swish"]:
+                assert seq_idx is None, "varlen conv1d requires the causal_conv1d package"
+                xBC = self.act(
+                    self.conv1d(xBC.transpose(1, 2)).transpose(1, 2)[:, -(self.dconv - 1):]
+                )  # (B, L, self.d_ssm + 2 * ngroups * d_state)
+            else:
+                xBC = causal_conv1d_fn(
+                    xBC.transpose(1, 2),
+                    rearrange(self.conv1d.weight, "d 1 w -> d w"),
+                    bias=self.conv1d.bias,
+                    activation=self.activation,
+                    # seq_idx=seq_idx,
+                ).transpose(1, 2)
+            x, B, C = torch.split(xBC, [self.d_ssm, self.ngroups * self.d_state, self.ngroups * self.d_state], dim=-1)
+
+            y = mamba_chunk_scan_combined(
+                rearrange(x, "b l (h p) -> b l h p", p=self.headdim),
+                dt,
+                A,
+                rearrange(B, "b l (g n) -> b l g n", g=self.ngroups),
+                rearrange(C, "b l (g n) -> b l g n", g=self.ngroups),
+                chunk_size=self.chunk_size,
+                D=rearrange(self.D, "(h p) -> h p", p=self.headdim) if self.D_has_hdim else self.D,
+                z=rearrange(z, "b l (h p) -> b l h p", p=self.headdim) if not self.rmsnorm else None,
+                dt_bias=self.dt_bias,
+                dt_softplus=True,
+                seq_idx=seq_idx,
+                cu_seqlens=cu_seqlens,
+                **dt_limit_kwargs,
+                return_final_states=ssm_state is not None,
+                return_varlen_states=cu_seqlens is not None and inference_params is not None,
+            )
+            if ssm_state is not None:
+                y, last_state, *rest = y
+                if cu_seqlens is None:
+                    ssm_state.copy_(last_state)
+                else:
+                    varlen_states = rest[0]
+                    ssm_state.copy_(varlen_states)
+            y = rearrange(y, "b l h p -> b l (h p)")
+            if self.rmsnorm:
+                y = self.norm(y, z)
+            if d_mlp > 0:
+                y = torch.cat([F.silu(z0) * x0, y], dim=-1)
+            if seqlen_og is not None:
+                y = rearrange(y, "b l d -> (b l) d")
+
+            scan_outputs = y
+            if 'repeat_v' in self.fused_multihead_config and self.fused_multihead_config['repeat_v']:
+                num_repeat = scan_outputs.shape[-1] // attn_outputs.shape[-1]
+                attn_outputs = attn_outputs.repeat(1, 1, num_repeat)
+
+            hidden_states = (self.pre_avg_layernorm1(attn_outputs) + self.pre_avg_layernorm2(scan_outputs)) / 2
+            out = self.out_proj(hidden_states)
+                
+        return out, attn_key_value, past_key_value
+
+
+    def step(self, hidden_states, conv_state, ssm_state):
+        dtype = hidden_states.dtype
+        assert hidden_states.shape[1] == 1, "Only support decoding with 1 token at a time for now"
+        zxbcdt = self.in_proj(hidden_states.squeeze(1))  # (B 2D)
+        d_mlp = (zxbcdt.shape[-1] - 2 * self.d_ssm - 2 * self.ngroups * self.d_state - self.nheads) // 2
+        z0, x0, z, xBC, dt = torch.split(
+            zxbcdt,
+            [d_mlp, d_mlp, self.d_ssm, self.d_ssm + 2 * self.ngroups * self.d_state, self.nheads],
+            dim=-1
+        )
+
+        # Conv step
+        if causal_conv1d_update is None:
+            conv_state.copy_(torch.roll(conv_state, shifts=-1, dims=-1))  # Update state (B D W)
+            conv_state[:, :, -1] = xBC
+            xBC = torch.sum(conv_state * rearrange(self.conv1d.weight, "d 1 w -> d w"), dim=-1)  # (B D)
+            if self.conv1d.bias is not None:
+                xBC = xBC + self.conv1d.bias
+            xBC = self.act(xBC).to(dtype=dtype)
+        else:
+            xBC = causal_conv1d_update(
+                xBC,
+                conv_state,
+                rearrange(self.conv1d.weight, "d 1 w -> d w"),
+                self.conv1d.bias,
+                self.activation,
+            )
+
+        x, B, C = torch.split(xBC, [self.d_ssm, self.ngroups * self.d_state, self.ngroups * self.d_state], dim=-1)
+        A = -torch.exp(self.A_log.float())  # (nheads,)
+
+        # SSM step
+        if selective_state_update is None:
+            assert self.ngroups == 1, "Only support ngroups=1 for this inference code path"
+            # Discretize A and B
+            dt = F.softplus(dt + self.dt_bias.to(dtype=dt.dtype))  # (batch, nheads)
+            dA = torch.exp(dt * A)  # (batch, nheads)
+            x = rearrange(x, "b (h p) -> b h p", p=self.headdim)
+            dBx = torch.einsum("bh,bn,bhp->bhpn", dt, B, x)
+            ssm_state.copy_(ssm_state * rearrange(dA, "b h -> b h 1 1") + dBx)
+            y = torch.einsum("bhpn,bn->bhp", ssm_state.to(dtype), C)
+            y = y + rearrange(self.D.to(dtype), "h -> h 1") * x
+            y = rearrange(y, "b h p -> b (h p)")
+            if not self.rmsnorm:
+                y = y * self.act(z)  # (B D)
+        else:
+            A = repeat(A, "h -> h p n", p=self.headdim, n=self.d_state).to(dtype=torch.float32)
+            dt = repeat(dt, "b h -> b h p", p=self.headdim)
+            dt_bias = repeat(self.dt_bias, "h -> h p", p=self.headdim)
+            D = repeat(self.D, "h -> h p", p=self.headdim)
+            B = rearrange(B, "b (g n) -> b g n", g=self.ngroups)
+            C = rearrange(C, "b (g n) -> b g n", g=self.ngroups)
+            x_reshaped = rearrange(x, "b (h p) -> b h p", p=self.headdim)
+            if not self.rmsnorm:
+                z = rearrange(z, "b (h p) -> b h p", p=self.headdim)
+            y = selective_state_update(
+                ssm_state, x_reshaped, dt, A, B, C, D, z=z if not self.rmsnorm else None,
+                dt_bias=dt_bias, dt_softplus=True
+            )
+            y = rearrange(y, "b h p -> b (h p)")
+        if self.rmsnorm:
+            y = self.norm(y, z)
+        if d_mlp > 0:
+            y = torch.cat([F.silu(z0) * x0, y], dim=-1)
+        out = self.out_proj(y)
+        
+        print(out)
+        input()
+        return out.unsqueeze(1), conv_state, ssm_state
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        device = self.out_proj.weight.device
+        conv_dtype = self.conv1d.weight.dtype if dtype is None else dtype
+        conv_state = torch.zeros(
+            batch_size, self.d_conv, self.conv1d.weight.shape[0], device=device, dtype=conv_dtype
+        ).transpose(1, 2)
+        ssm_dtype = self.in_proj.weight.dtype if dtype is None else dtype
+        ssm_state = torch.zeros(
+            batch_size, self.nheads, self.headdim, self.d_state, device=device, dtype=ssm_dtype
+        )
+        return conv_state, ssm_state
+
+    def _get_states_from_cache(self, inference_params, batch_size, initialize_states=False):
+        assert self.layer_idx is not None
+        if self.layer_idx not in inference_params.key_value_memory_dict:
+            batch_shape = (batch_size,)
+            conv_state = torch.zeros(
+                batch_size,
+                self.d_conv,
+                self.conv1d.weight.shape[0],
+                device=self.conv1d.weight.device,
+                dtype=self.conv1d.weight.dtype,
+            ).transpose(1, 2)
+            ssm_state = torch.zeros(
+                batch_size,
+                self.nheads,
+                self.headdim,
+                self.d_state,
+                device=self.in_proj.weight.device,
+                dtype=self.in_proj.weight.dtype,
+            )
+            inference_params.key_value_memory_dict[self.layer_idx] = (conv_state, ssm_state)
+        else:
+            conv_state, ssm_state = inference_params.key_value_memory_dict[self.layer_idx]
+            # TODO: What if batch size changes between generation, and we reuse the same states?
+            if initialize_states:
+                conv_state.zero_()
+                ssm_state.zero_()
+        return conv_state, ssm_state
+
+
+class Mamba2_Multihead(nn.Module):
+    def __init__(
+        self,
+        config,
+        conv_init=None,
+        headdim=64,
+        d_ssm=None,  # If not None, we only apply SSM on this many dimensions, the rest uses gated MLP
+        ngroups=1,
+        A_init_range=(1, 16),
+        D_has_hdim=False,
+        rmsnorm=True,
+        norm_before_gate=False,
+        dt_min=0.001,
+        dt_max=0.1,
+        dt_init_floor=1e-4,
+        dt_limit=(0.0, float("inf")),
+        bias=False,
+        conv_bias=True,
+        # Fused kernel and sharding options
+        chunk_size=256,
+        use_mem_eff_path=False, # True,
+        layer_idx=None,  # Absorb kwarg for general module
+        process_group=None,
+        sequence_parallel=True,
+        device=None,
+        dtype=None,
+    ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        
+        self.config = config
+        self.d_model = config.hidden_size
+        self.d_state = config.mamba_d_state
+        self.d_conv = config.mamba_d_conv
+            
+        self.conv_init = conv_init
+        self.expand = config.mamba_expand
+        self.process_group = process_group
+        self.sequence_parallel = sequence_parallel
+        self.world_size = 1 if process_group is None else process_group.size()
+        self.local_rank = 0 if process_group is None else process_group.rank()
+        self.d_inner = (self.expand * self.d_model) // self.world_size
+        assert self.d_inner * self.world_size == self.expand * self.d_model
+        self.headdim = config.mamba2_headdim
+        self.d_ssm = self.d_inner if d_ssm is None else d_ssm // self.world_size
+        assert ngroups % self.world_size == 0
+        self.ngroups = ngroups // self.world_size
+        assert self.d_ssm % self.headdim == 0
+        self.nheads = self.d_ssm // self.headdim
+        self.D_has_hdim = D_has_hdim
+        self.rmsnorm = rmsnorm
+        self.norm_before_gate = norm_before_gate
+        self.dt_limit = dt_limit
+        self.activation = "silu"
+        self.chunk_size = chunk_size
+        self.use_mem_eff_path = use_mem_eff_path
+        self.layer_idx = layer_idx
+
+        assert (self.d_model * self.expand / self.headdim) % 8 == 0
+
+        self.mamba_multihead_config = config.mamba_multihead_config
+        self.share_ratio = self.mamba_multihead_config['share_ratio']
+        
+        self.reuse_ssm = self.mamba_multihead_config['reuse_ssm']
+        self.num_ssm_param = 1 if self.reuse_ssm else self.share_ratio
+        
+        if self.reuse_ssm:
+            if self.mamba_multihead_config['alpha_mode'] == 'learnable':
+                self.alpha = nn.Parameter(torch.ones(self.share_ratio))
+            elif self.mamba_multihead_config['alpha_mode'] == 'manual': 
+                manual_alpha_base = self.mamba_multihead_config['manual_alpha_base']
+                self.alpha = [1 / manual_alpha_base ** k for k in range(self.share_ratio)]
+            else:
+                raise ValueError(f"No such alpha_mode: {self.mamba_multihead_config['alpha_mode']}")
+
+        # Order: [z, x, B, C, dt]
+        d_in_proj = 2 * self.d_inner + 2 * self.ngroups * self.d_state + self.nheads * self.num_ssm_param
+        if self.process_group is None:
+            self.in_proj = nn.Linear(self.d_model, d_in_proj, bias=bias, **factory_kwargs)
+        else:
+            self.in_proj = ColumnParallelLinear(self.d_model, d_in_proj * self.world_size, bias=bias,
+                                                process_group=self.process_group, sequence_parallel=self.sequence_parallel,
+                                                **factory_kwargs)
+
+        conv_dim = self.d_ssm + 2 * self.ngroups * self.d_state
+        self.conv1d = nn.Conv1d(
+            in_channels=conv_dim,
+            out_channels=conv_dim,
+            bias=conv_bias,
+            kernel_size=self.d_conv,
+            groups=conv_dim,
+            padding=self.d_conv - 1,
+            **factory_kwargs,
+        )
+        if self.conv_init is not None:
+            nn.init.uniform_(self.conv1d.weight, -self.conv_init, self.conv_init)
+
+        self.act = nn.SiLU()
+
+        # Initialize log dt bias
+        dt = torch.exp(
+            torch.rand(self.nheads, **factory_kwargs) * (math.log(dt_max) - math.log(dt_min))
+            + math.log(dt_min)
+        )
+        dt = torch.clamp(dt, min=dt_init_floor)
+        # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
+        inv_dt = dt + torch.log(-torch.expm1(-dt))
+        self.dt_bias = nn.ParameterList([nn.Parameter(inv_dt) for _ in range(self.num_ssm_param)])
+        # Just to be explicit. Without this we already don't put wd on dt_bias because of the check
+        # name.endswith("bias") in param_grouping.py
+        self.dt_bias._no_weight_decay = True
+
+        assert A_init_range[0] > 0 and A_init_range[1] >= A_init_range[0]
+        A = torch.empty(self.nheads, dtype=torch.float32, device=device).uniform_(*A_init_range)
+        A_log = torch.log(A).to(dtype=dtype)
+        self.A_log = nn.ParameterList([nn.Parameter(A_log) for _ in range(self.num_ssm_param)])
+        self.A_log._no_weight_decay = True
+
+        # D "skip" parameter
+        self.D = nn.ParameterList([nn.Parameter(torch.ones(self.d_ssm if self.D_has_hdim else self.nheads, device=device)) for _ in range(self.num_ssm_param)])
+        self.D._no_weight_decay = True
+
+        if self.rmsnorm:
+            assert RMSNormGated is not None
+            self.norm = RMSNormGated(self.d_ssm, eps=1e-5, norm_before_gate=self.norm_before_gate,
+                                     group_size=self.d_ssm // ngroups, **factory_kwargs)
+
+        if self.process_group is None:
+            self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=bias, **factory_kwargs)
+        else:
+            self.out_proj = RowParallelLinear(self.d_inner * self.world_size, self.d_model, bias=bias,
+                                              process_group=self.process_group, sequence_parallel=self.sequence_parallel,
+                                              **factory_kwargs)
+
+
+        if self.mamba_multihead_config['merge_op'] == 'norm':
+            self.multihead_layernorm = nn.ModuleList([JambaRMSNorm(self.d_ssm, eps=config.rms_norm_eps) for _ in range(self.share_ratio)])
+        elif self.mamba_multihead_config['merge_op'] == 'scalar_gate':
+            self.multi_head_selection_layer = nn.Linear(self.d_ssm,  self.share_ratio)
+        elif self.mamba_multihead_config['merge_op'] == 'concat':
+            assert self.d_ssm % self.share_ratio == 0
+            self.multihead_layernorm = nn.ModuleList([JambaRMSNorm(self.d_ssm, eps=config.rms_norm_eps) for _ in range(self.share_ratio)])
+            self.reduction_layer = nn.Linear(self.d_ssm,  self.d_ssm//self.share_ratio)
+
+
+    def forward(self, hidden_states, attention_mask=None, past_key_value=None, seqlen=None, seq_idx=None, cu_seqlens=None, inference_params=None):
+        """
+        hidden_states: (batch, seqlen, hidden_dim) if seqlen=None.
+            If seqlen is not None, hidden_states is (batch * seqlen, hidden_dim). This is so that when we
+            split hidden_states during sequence parallel, we split the batch * seqlen dimension
+            (in case batch is small).
+        Returns: same shape as u
+        """
+        assert past_key_value is None, "Not implemented yet!!!"
+
+        seqlen_og = seqlen
+        if seqlen is None:
+            batch, seqlen, dim = hidden_states.shape
+        else:
+            batch_seqlen, dim = hidden_states.shape
+            batch = batch_seqlen // seqlen
+
+        conv_state, ssm_state = None, None
+        if inference_params is not None:
+            inference_batch = cu_seqlens.shape[0] - 1 if cu_seqlens is not None else batch
+            conv_state, ssm_state = self._get_states_from_cache(inference_params, inference_batch)
+            if inference_params.seqlen_offset > 0:
+                # The states are updated inplace
+                out, _, _ = self.step(hidden_states, conv_state, ssm_state)
+                return out
+
+        zxbcdt = self.in_proj(hidden_states)  # (B, L, d_in_proj) or (B * L, d_in_proj)
+        if seqlen_og is not None:
+            zxbcdt = rearrange(zxbcdt, "(b l) d -> b l d", l=seqlen)
+
+        dt_limit_kwargs = {} if self.dt_limit == (0.0, float("inf")) else dict(dt_limit=self.dt_limit)
+        if self.use_mem_eff_path and inference_params is None:
+            # If the model is loaded in fp16, without the .float() here, A might be -inf
+            A = -torch.exp(self.A_log.float())  # (nheads) or (d_inner, d_state)
+
+            out = mamba_split_conv1d_scan_combined(
+                zxbcdt,
+                rearrange(self.conv1d.weight, "d 1 w -> d w"),
+                self.conv1d.bias,
+                self.dt_bias,
+                A,
+                D=rearrange(self.D, "(h p) -> h p", p=self.headdim) if self.D_has_hdim else self.D,
+                chunk_size=self.chunk_size,
+                seq_idx=seq_idx,
+                activation=self.activation,
+                rmsnorm_weight=self.norm.weight if self.rmsnorm else None,
+                rmsnorm_eps=self.norm.eps if self.rmsnorm else 1e-6,
+                outproj_weight=self.out_proj.weight,
+                outproj_bias=self.out_proj.bias,
+                headdim=None if self.D_has_hdim else self.headdim,
+                ngroups=self.ngroups,
+                norm_before_gate=self.norm_before_gate,
+                **dt_limit_kwargs,
+            )
+            if seqlen_og is not None:
+                out = rearrange(out, "b l d -> (b l) d")
+            if self.process_group is not None:
+                reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce
+                out = reduce_fn(out, self.process_group)
+        else:
+            d_mlp = (zxbcdt.shape[-1] - 2 * self.d_ssm - 2 * self.ngroups * self.d_state - self.nheads * self.num_ssm_param) // 2
+            z0, x0, z, xBC, dt = torch.split(
+                zxbcdt,
+                [d_mlp, d_mlp, self.d_ssm, self.d_ssm + 2 * self.ngroups * self.d_state, self.nheads * self.num_ssm_param],
+                dim=-1
+            )
+
+            if conv_state is not None:
+                if cu_seqlens is None:
+                    # If we just take xBC[:, :, -self.d_conv :], it will error if seqlen < self.d_conv
+                    # Instead F.pad will pad with zeros if seqlen < self.d_conv, and truncate otherwise.
+                    xBC_t = rearrange(xBC, "b l d -> b d l")
+                    conv_state.copy_(F.pad(xBC_t, (self.d_conv - xBC_t.shape[-1], 0)))  # Update state (B D W)
+                else:
+                    assert causal_conv1d_varlen_states is not None, "varlen inference requires causal_conv1d package"
+                    assert batch == 1, "varlen inference only supports batch dimension 1"
+                    conv_varlen_states = causal_conv1d_varlen_states(
+                        xBC.squeeze(0), cu_seqlens, state_len=conv_state.shape[-1]
+                    )
+                    conv_state.copy_(conv_varlen_states)
+            assert self.activation in ["silu", "swish"]
+            if causal_conv1d_fn is None or self.activation not in ["silu", "swish"]:
+                assert seq_idx is None, "varlen conv1d requires the causal_conv1d package"
+                xBC = self.act(
+                    self.conv1d(xBC.transpose(1, 2)).transpose(1, 2)[:, -(self.dconv - 1):]
+                )  # (B, L, self.d_ssm + 2 * ngroups * d_state)
+            else:
+                xBC = causal_conv1d_fn(
+                    xBC.transpose(1, 2),
+                    rearrange(self.conv1d.weight, "d 1 w -> d w"),
+                    bias=self.conv1d.bias,
+                    activation=self.activation,
+                    seq_idx=seq_idx,
+                ).transpose(1, 2)
+            x, B, C = torch.split(xBC, [self.d_ssm, self.ngroups * self.d_state, self.ngroups * self.d_state], dim=-1)
+
+            x = rearrange(x, "b l (h p) -> b l h p", p=self.headdim)
+            B = rearrange(B, "b l (g n) -> b l g n", g=self.ngroups)
+            C = rearrange(C, "b l (g n) -> b l g n", g=self.ngroups)
+
+            outputs_list = []
+            dt_list = dt
+            for i in range(self.num_ssm_param):
+                dt = dt_list[..., self.nheads*i:self.nheads*(i+1)]
+                A = -torch.exp(self.A_log[i].float())  # (nheads) or (d_inner, d_state)
+                D = rearrange(self.D[i], "(h p) -> h p", p=self.headdim) if self.D_has_hdim else self.D[i]
+                dt_bias = self.dt_bias[i]
+
+                if self.reuse_ssm:
+                    #### duplicate heads with different decays
+                    if self.mamba_multihead_config['alpha_mode'] == 'learnable':
+                        decay = self.alpha  # [share_ratio]
+                    elif self.mamba_multihead_config['alpha_mode'] == 'manual': 
+                        decay = torch.tensor(self.alpha).to(dt)  # [share_ratio]
+                    
+                    dt = dt.repeat(1, 1, self.share_ratio)  # [bs, seq_len, self.nheads * share_ratio]
+                    decay = decay.view(-1, 1).repeat(1, self.nheads).view(-1)  # [self.nheads * share_ratio]
+                    dt = dt * decay  # [bs, seq_len, nheads * share_ratio]
+
+                    dt_bias = dt_bias.repeat(self.share_ratio) * decay  # [nheads * share_ratio]
+
+                    x = x.repeat(1,1,self.share_ratio,1)  # [bs, seq_len, nheads * share_ratio, head_dim]
+                    D = D.repeat(self.share_ratio,1) if self.D_has_hdim else D.repeat(self.share_ratio)  # [nheads * share_ratio]
+                    A = A.repeat(self.share_ratio)  # [nheads * share_ratio]
+
+                y = mamba_chunk_scan_combined(
+                    x,
+                    dt, 
+                    A,  
+                    B,
+                    C,
+                    chunk_size=self.chunk_size,
+                    D=D,
+                    z=rearrange(z, "b l (h p) -> b l h p", p=self.headdim).repeat(1,1,self.share_ratio,1) if not self.rmsnorm else None,
+                    dt_bias=dt_bias,
+                    dt_softplus=True,
+                    seq_idx=seq_idx,
+                    cu_seqlens=cu_seqlens,
+                    **dt_limit_kwargs,
+                    return_final_states=ssm_state is not None,
+                    return_varlen_states=cu_seqlens is not None and inference_params is not None,
+                )
+                if ssm_state is not None:
+                    y, last_state, *rest = y
+                    if cu_seqlens is None:
+                        ssm_state.copy_(last_state)
+                    else:
+                        varlen_states = rest[0]
+                        ssm_state.copy_(varlen_states)
+                
+                outputs_list.append(y)
+            
+            if len(outputs_list) > 1:
+                y = torch.cat(outputs_list, dim=2)
+            
+            #### merge heads
+            num_repeat = y.shape[2] // self.nheads
+            head_outputs = torch.chunk(y, num_repeat, dim=2)
+            head_outputs = [rearrange(item, "b l h p -> b l (h p)") for item in head_outputs]
+            
+            if self.mamba_multihead_config['merge_op'] == 'norm':
+                y = sum([self.multihead_layernorm[k](item) for k, item in enumerate(head_outputs)])
+
+            elif self.mamba_multihead_config['merge_op'] == 'concat':
+                head_outputs = [self.reduction_layer(self.multihead_layernorm[k](item)) for k, item in enumerate(head_outputs)]
+                y = torch.cat(head_outputs, dim=-1)
+            else:
+                raise ValueError(f"No such merge_op: {self.mamba_multihead_config['merge_op']}")
+
+            if self.rmsnorm:
+                y = self.norm(y, z)
+            if d_mlp > 0:
+                y = torch.cat([F.silu(z0) * x0, y], dim=-1)
+            if seqlen_og is not None:
+                y = rearrange(y, "b l d -> (b l) d")
+            out = self.out_proj(y)
+        return out, past_key_value
+
+    def step(self, hidden_states, conv_state, ssm_state):
+        dtype = hidden_states.dtype
+        assert hidden_states.shape[1] == 1, "Only support decoding with 1 token at a time for now"
+        zxbcdt = self.in_proj(hidden_states.squeeze(1))  # (B 2D)
+        d_mlp = (zxbcdt.shape[-1] - 2 * self.d_ssm - 2 * self.ngroups * self.d_state - self.nheads) // 2
+        z0, x0, z, xBC, dt = torch.split(
+            zxbcdt,
+            [d_mlp, d_mlp, self.d_ssm, self.d_ssm + 2 * self.ngroups * self.d_state, self.nheads],
+            dim=-1
+        )
+
+        # Conv step
+        if causal_conv1d_update is None:
+            conv_state.copy_(torch.roll(conv_state, shifts=-1, dims=-1))  # Update state (B D W)
+            conv_state[:, :, -1] = xBC
+            xBC = torch.sum(conv_state * rearrange(self.conv1d.weight, "d 1 w -> d w"), dim=-1)  # (B D)
+            if self.conv1d.bias is not None:
+                xBC = xBC + self.conv1d.bias
+            xBC = self.act(xBC).to(dtype=dtype)
+        else:
+            xBC = causal_conv1d_update(
+                xBC,
+                conv_state,
+                rearrange(self.conv1d.weight, "d 1 w -> d w"),
+                self.conv1d.bias,
+                self.activation,
+            )
+
+        x, B, C = torch.split(xBC, [self.d_ssm, self.ngroups * self.d_state, self.ngroups * self.d_state], dim=-1)
+        A = -torch.exp(self.A_log.float())  # (nheads,)
+
+        # SSM step
+        if selective_state_update is None:
+            assert self.ngroups == 1, "Only support ngroups=1 for this inference code path"
+            # Discretize A and B
+            dt = F.softplus(dt + self.dt_bias.to(dtype=dt.dtype))  # (batch, nheads)
+            dA = torch.exp(dt * A)  # (batch, nheads)
+            x = rearrange(x, "b (h p) -> b h p", p=self.headdim)
+            dBx = torch.einsum("bh,bn,bhp->bhpn", dt, B, x)
+            ssm_state.copy_(ssm_state * rearrange(dA, "b h -> b h 1 1") + dBx)
+            y = torch.einsum("bhpn,bn->bhp", ssm_state.to(dtype), C)
+            y = y + rearrange(self.D.to(dtype), "h -> h 1") * x
+            y = rearrange(y, "b h p -> b (h p)")
+            if not self.rmsnorm:
+                y = y * self.act(z)  # (B D)
+        else:
+            A = repeat(A, "h -> h p n", p=self.headdim, n=self.d_state).to(dtype=torch.float32)
+            dt = repeat(dt, "b h -> b h p", p=self.headdim)
+            dt_bias = repeat(self.dt_bias, "h -> h p", p=self.headdim)
+            D = repeat(self.D, "h -> h p", p=self.headdim)
+            B = rearrange(B, "b (g n) -> b g n", g=self.ngroups)
+            C = rearrange(C, "b (g n) -> b g n", g=self.ngroups)
+            x_reshaped = rearrange(x, "b (h p) -> b h p", p=self.headdim)
+            if not self.rmsnorm:
+                z = rearrange(z, "b (h p) -> b h p", p=self.headdim)
+            y = selective_state_update(
+                ssm_state, x_reshaped, dt, A, B, C, D, z=z if not self.rmsnorm else None,
+                dt_bias=dt_bias, dt_softplus=True
+            )
+            y = rearrange(y, "b h p -> b (h p)")
+        if self.rmsnorm:
+            y = self.norm(y, z)
+        if d_mlp > 0:
+            y = torch.cat([F.silu(z0) * x0, y], dim=-1)
+        out = self.out_proj(y)
+        return out.unsqueeze(1), conv_state, ssm_state
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        device = self.out_proj.weight.device
+        conv_dtype = self.conv1d.weight.dtype if dtype is None else dtype
+        conv_state = torch.zeros(
+            batch_size, self.d_conv, self.conv1d.weight.shape[0], device=device, dtype=conv_dtype
+        ).transpose(1, 2)
+        ssm_dtype = self.in_proj.weight.dtype if dtype is None else dtype
+        ssm_state = torch.zeros(
+            batch_size, self.nheads, self.headdim, self.d_state, device=device, dtype=ssm_dtype
+        )
+        return conv_state, ssm_state
+
+    def _get_states_from_cache(self, inference_params, batch_size, initialize_states=False):
+        assert self.layer_idx is not None
+        if self.layer_idx not in inference_params.key_value_memory_dict:
+            batch_shape = (batch_size,)
+            conv_state = torch.zeros(
+                batch_size,
+                self.d_conv,
+                self.conv1d.weight.shape[0],
+                device=self.conv1d.weight.device,
+                dtype=self.conv1d.weight.dtype,
+            ).transpose(1, 2)
+            ssm_state = torch.zeros(
+                batch_size,
+                self.nheads,
+                self.headdim,
+                self.d_state,
+                device=self.in_proj.weight.device,
+                dtype=self.in_proj.weight.dtype,
+            )
+            inference_params.key_value_memory_dict[self.layer_idx] = (conv_state, ssm_state)
+        else:
+            conv_state, ssm_state = inference_params.key_value_memory_dict[self.layer_idx]
+            # TODO: What if batch size changes between generation, and we reuse the same states?
+            if initialize_states:
+                conv_state.zero_()
+                ssm_state.zero_()
+        return conv_state, ssm_state
+    
+
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Jamba
+class JambaRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        JambaRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+
+
+
+

model-00001-of-00002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8fc337e437ea93d1d9a6c5d93793512da4b571843edfa616516630418244947b
+size 4995785984

model-00002-of-00002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:46319b16d9db944033225295486a090859d5113e4873eca20c0d836fa38a3f09
+size 491849664

model.safetensors.index.json

@@ -0,0 +1,241 @@
+{
+  "metadata": {
+    "total_size": 5487609216
+  },
+  "weight_map": {
+    "model.embed_tokens.weight": "model-00001-of-00002.safetensors",
+    "model.final_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.0.gla.b_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.gla.k_conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.gla.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.gla.o_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.gla.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.gla.q_conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.gla.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.gla.v_conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.gla.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.moe.experts.0.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.moe.experts.0.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.moe.experts.0.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.1.pre_moe_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.mamba.A_log": "model-00001-of-00002.safetensors",
+    "model.layers.10.mamba.D": "model-00001-of-00002.safetensors",
+    "model.layers.10.mamba.conv1d.bias": "model-00001-of-00002.safetensors",
+    "model.layers.10.mamba.conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.mamba.dt_bias": "model-00001-of-00002.safetensors",
+    "model.layers.10.mamba.in_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.mamba.norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.10.mamba.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.moe.experts.0.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.moe.experts.0.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.moe.experts.0.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.11.pre_moe_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.pre_moe_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.12.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.moe.experts.0.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.moe.experts.0.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.moe.experts.0.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.13.pre_moe_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.mamba.A_log": "model-00001-of-00002.safetensors",
+    "model.layers.14.mamba.D": "model-00001-of-00002.safetensors",
+    "model.layers.14.mamba.conv1d.bias": "model-00001-of-00002.safetensors",
+    "model.layers.14.mamba.conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.mamba.dt_bias": "model-00001-of-00002.safetensors",
+    "model.layers.14.mamba.in_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.mamba.norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.14.mamba.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.moe.experts.0.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.moe.experts.0.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.moe.experts.0.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.15.pre_moe_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.gla.b_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.gla.k_conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.gla.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.gla.o_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.gla.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.gla.q_conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.gla.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.gla.v_conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.gla.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.16.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.moe.experts.0.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.moe.experts.0.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.moe.experts.0.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.17.pre_moe_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.mamba.A_log": "model-00001-of-00002.safetensors",
+    "model.layers.18.mamba.D": "model-00001-of-00002.safetensors",
+    "model.layers.18.mamba.conv1d.bias": "model-00001-of-00002.safetensors",
+    "model.layers.18.mamba.conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.mamba.dt_bias": "model-00001-of-00002.safetensors",
+    "model.layers.18.mamba.in_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.mamba.norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.18.mamba.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.moe.experts.0.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.moe.experts.0.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.moe.experts.0.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.19.pre_moe_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mamba.A_log": "model-00001-of-00002.safetensors",
+    "model.layers.2.mamba.D": "model-00001-of-00002.safetensors",
+    "model.layers.2.mamba.conv1d.bias": "model-00001-of-00002.safetensors",
+    "model.layers.2.mamba.conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mamba.dt_bias": "model-00001-of-00002.safetensors",
+    "model.layers.2.mamba.in_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mamba.norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.2.mamba.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.pre_moe_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.20.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.moe.experts.0.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.moe.experts.0.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.moe.experts.0.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.21.pre_moe_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.mamba.A_log": "model-00001-of-00002.safetensors",
+    "model.layers.22.mamba.D": "model-00001-of-00002.safetensors",
+    "model.layers.22.mamba.conv1d.bias": "model-00001-of-00002.safetensors",
+    "model.layers.22.mamba.conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.mamba.dt_bias": "model-00001-of-00002.safetensors",
+    "model.layers.22.mamba.in_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.mamba.norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.22.mamba.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.23.moe.experts.0.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.23.moe.experts.0.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.23.moe.experts.0.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.23.pre_moe_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.24.gla.b_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.24.gla.k_conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.24.gla.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.24.gla.o_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.24.gla.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.24.gla.q_conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.24.gla.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.24.gla.v_conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.24.gla.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.24.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.25.moe.experts.0.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.25.moe.experts.0.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.25.moe.experts.0.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.25.pre_moe_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.26.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.26.mamba.A_log": "model-00001-of-00002.safetensors",
+    "model.layers.26.mamba.D": "model-00001-of-00002.safetensors",
+    "model.layers.26.mamba.conv1d.bias": "model-00001-of-00002.safetensors",
+    "model.layers.26.mamba.conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.26.mamba.dt_bias": "model-00001-of-00002.safetensors",
+    "model.layers.26.mamba.in_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.26.mamba.norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.26.mamba.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.27.moe.experts.0.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.27.moe.experts.0.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.27.moe.experts.0.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.27.pre_moe_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.28.gla.b_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.28.gla.k_conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.28.gla.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.28.gla.o_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.28.gla.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.28.gla.q_conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.28.gla.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.28.gla.v_conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.28.gla.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.28.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.29.moe.experts.0.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.29.moe.experts.0.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.29.moe.experts.0.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.29.pre_moe_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.moe.experts.0.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.moe.experts.0.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.moe.experts.0.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.3.pre_moe_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.30.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.30.mamba.A_log": "model-00001-of-00002.safetensors",
+    "model.layers.30.mamba.D": "model-00001-of-00002.safetensors",
+    "model.layers.30.mamba.conv1d.bias": "model-00001-of-00002.safetensors",
+    "model.layers.30.mamba.conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.30.mamba.dt_bias": "model-00001-of-00002.safetensors",
+    "model.layers.30.mamba.in_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.30.mamba.norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.30.mamba.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.31.moe.experts.0.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.31.moe.experts.0.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.31.moe.experts.0.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.31.pre_moe_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.32.gla.b_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.32.gla.k_conv1d.weight": "model-00002-of-00002.safetensors",
+    "model.layers.32.gla.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.32.gla.o_norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.32.gla.o_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.32.gla.q_conv1d.weight": "model-00002-of-00002.safetensors",
+    "model.layers.32.gla.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.32.gla.v_conv1d.weight": "model-00002-of-00002.safetensors",
+    "model.layers.32.gla.v_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.32.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.33.moe.experts.0.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.33.moe.experts.0.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.33.moe.experts.0.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.33.pre_moe_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.34.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.34.mamba.A_log": "model-00002-of-00002.safetensors",
+    "model.layers.34.mamba.D": "model-00002-of-00002.safetensors",
+    "model.layers.34.mamba.conv1d.bias": "model-00002-of-00002.safetensors",
+    "model.layers.34.mamba.conv1d.weight": "model-00002-of-00002.safetensors",
+    "model.layers.34.mamba.dt_bias": "model-00002-of-00002.safetensors",
+    "model.layers.34.mamba.in_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.34.mamba.norm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.34.mamba.out_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.35.moe.experts.0.down_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.35.moe.experts.0.gate_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.35.moe.experts.0.up_proj.weight": "model-00002-of-00002.safetensors",
+    "model.layers.35.pre_moe_layernorm.weight": "model-00002-of-00002.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.pre_moe_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.4.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.moe.experts.0.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.moe.experts.0.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.moe.experts.0.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.5.pre_moe_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.mamba.A_log": "model-00001-of-00002.safetensors",
+    "model.layers.6.mamba.D": "model-00001-of-00002.safetensors",
+    "model.layers.6.mamba.conv1d.bias": "model-00001-of-00002.safetensors",
+    "model.layers.6.mamba.conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.mamba.dt_bias": "model-00001-of-00002.safetensors",
+    "model.layers.6.mamba.in_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.mamba.norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.6.mamba.out_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.moe.experts.0.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.moe.experts.0.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.moe.experts.0.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.7.pre_moe_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.gla.b_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.gla.k_conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.gla.k_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.gla.o_norm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.gla.o_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.gla.q_conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.gla.q_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.gla.v_conv1d.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.gla.v_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.8.input_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.moe.experts.0.down_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.moe.experts.0.gate_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.moe.experts.0.up_proj.weight": "model-00001-of-00002.safetensors",
+    "model.layers.9.pre_moe_layernorm.weight": "model-00001-of-00002.safetensors",
+    "model.memory_tokens": "model-00001-of-00002.safetensors"
+  }
+}