Upload 2 files

configuration_motif.py

@@ -0,0 +1,300 @@
+import torch
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class MotifConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`MotifModel`]. It is used to instantiate a
+    Motif model according to the specified arguments, defining the model architecture.
+    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 151936):
+            Vocabulary size of the Motif model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`MotifModel`]
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 22016):
+            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 32):
+            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 `32`.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 32768):
+            The maximum sequence length that this model might ever be used with.
+        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`.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether the model's input and output word embeddings should be tied.
+        rope_theta (`float`, *optional*, defaults to 1000000.0):
+            The base period of the RoPE embeddings.
+        rope_scaling (`Dict`, *optional*):
+            Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
+            and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
+            accordingly.
+            Expected contents:
+                `rope_type` (`str`):
+                    The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
+                    'llama3'], with 'default' being the original RoPE implementation.
+                `factor` (`float`, *optional*):
+                    Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
+                    most scaling types, a `factor` of x will enable the model to handle sequences of length x *
+                    original maximum pre-trained length.
+                `original_max_position_embeddings` (`int`, *optional*):
+                    Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
+                    pretraining.
+                `attention_factor` (`float`, *optional*):
+                    Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
+                    computation. If unspecified, it defaults to value recommended by the implementation, using the
+                    `factor` field to infer the suggested value.
+                `beta_fast` (`float`, *optional*):
+                    Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 32.
+                `beta_slow` (`float`, *optional*):
+                    Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 1.
+                `short_factor` (`List[float]`, *optional*):
+                    Only used with 'longrope'. The scaling factor to be applied to short contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `long_factor` (`List[float]`, *optional*):
+                    Only used with 'longrope'. The scaling factor to be applied to long contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `low_freq_factor` (`float`, *optional*):
+                    Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
+                `high_freq_factor` (`float`, *optional*):
+                    Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
+        use_sliding_window (`bool`, *optional*, defaults to `False`):
+            Whether to use sliding window attention.
+        sliding_window (`int`, *optional*, defaults to 4096):
+            Sliding window attention (SWA) window size. If not specified, will default to `4096`.
+        max_window_layers (`int`, *optional*, defaults to 28):
+            The number of layers that use SWA (Sliding Window Attention). The bottom layers use SWA while the top use full attention.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+    ```python
+    >>> from transformers import MotifModel, MotifConfig
+    >>> # Initializing a Motif style configuration
+    >>> configuration = MotifConfig()
+    >>> # Initializing a model from the Motif-102B style configuration
+    >>> model = MotifModel(configuration)
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "Motif"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    base_model_tp_plan = {
+        # Attention
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        # Dense MLP
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+        # MoE experts (fused gate+up)
+        "layers.*.moe.experts.gate_up_proj": "packed_colwise",
+        "layers.*.moe.experts.down_proj": "rowwise",
+        # Shared experts
+        "layers.*.moe.shared_experts.gate_proj": "colwise",
+        "layers.*.moe.shared_experts.up_proj": "colwise",
+        "layers.*.moe.shared_experts.down_proj": "rowwise",
+    }
+
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        vocab_size=151936,
+        hidden_size=4096,
+        intermediate_size=22016,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=32,
+        hidden_act="silu",
+        max_position_embeddings=32768,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=1000000.0,
+        rope_scaling=None,
+        use_sliding_window=False,
+        sliding_window=4096,
+        max_window_layers=28,
+        sliding_window_pattern="interleave",
+        sliding_window_period=2,
+        attention_dropout=0.0,
+        # Differential Attention parameters
+        head_dim=None,
+        num_noise_heads=0,
+        k_ratio=1,
+        # MoE parameters
+        num_experts=0,
+        experts_top_k=2,
+        num_shared_experts=0,
+        interleave_moe_layer_step=0,
+        moe_intermediate_size=None,
+        score_func="softmax",
+        route_norm=False,
+        route_scale=1.0,
+        load_balance_coeff=None,
+        score_before_experts=False,
+        _debug_force_load_balance=False,
+        output_router_logits=False,
+        router_aux_loss_coef=0.0,
+        # MHC (Manifold-constrained Hyper-Connections) parameters
+        mhc_enabled=False,
+        mhc_expansion_rate=4,
+        mhc_identity_init=False,
+        mhc_sinkhorn_iters=20,
+        # DiffAttention V2 / Attention class
+        diff_v2=False,
+        attention_cls="basic",
+        # GDLA (Grouped Differential Latent Attention) parameters
+        q_lora_rank=0,
+        kv_lora_rank=0,
+        qk_rope_head_dim=None,
+        v_head_dim=None,
+        original_seq_len=32768,
+        rope_factor=1.0,
+        mscale=1.0,
+        swa_rope_theta=None,
+        # Attention output gating
+        headwise_attn_output_gate=False,
+        elementwise_attn_output_gate=False,
+        # MoE: first N layers always dense (no MoE), regardless of interleave schedule
+        n_dense_first_layers=0,
+        # MTP (Multi-Token Prediction) speculative decoding
+        num_nextn_predict_layers=0,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_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.use_sliding_window = use_sliding_window
+        self.sliding_window = sliding_window if use_sliding_window else None
+        self.max_window_layers = max_window_layers
+        self.sliding_window_pattern = sliding_window_pattern
+        self.sliding_window_period = sliding_window_period
+
+        # 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.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_dropout = attention_dropout
+
+        # Differential Attention configuration
+        self.head_dim = head_dim
+        self.num_noise_heads = num_noise_heads
+        self.k_ratio = k_ratio
+
+        # MoE configuration
+        self.num_experts = num_experts
+        self.experts_top_k = experts_top_k
+        self.num_shared_experts = num_shared_experts
+        self.interleave_moe_layer_step = interleave_moe_layer_step
+        self.moe_intermediate_size = moe_intermediate_size if moe_intermediate_size is not None else intermediate_size
+        self.score_func = score_func
+        self.route_norm = route_norm
+        self.route_scale = route_scale
+        self.load_balance_coeff = load_balance_coeff
+        self.score_before_experts = score_before_experts
+        self._debug_force_load_balance = _debug_force_load_balance
+        self.output_router_logits = output_router_logits
+        self.router_aux_loss_coef = router_aux_loss_coef
+
+        # MHC configuration
+        self.mhc_enabled = mhc_enabled
+        self.mhc_expansion_rate = mhc_expansion_rate
+        self.mhc_identity_init = mhc_identity_init
+        self.mhc_sinkhorn_iters = mhc_sinkhorn_iters
+
+        # DiffAttention V2 / Attention class
+        self.diff_v2 = diff_v2
+        self.attention_cls = attention_cls
+
+        # GDLA parameters
+        self.q_lora_rank = q_lora_rank
+        self.kv_lora_rank = kv_lora_rank
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.v_head_dim = v_head_dim
+        self.original_seq_len = original_seq_len
+        self.rope_factor = rope_factor
+        self.mscale = mscale
+        self.swa_rope_theta = swa_rope_theta
+
+        # Attention output gating
+        self.headwise_attn_output_gate = headwise_attn_output_gate
+        self.elementwise_attn_output_gate = elementwise_attn_output_gate
+
+        # MoE dense-first layers
+        self.n_dense_first_layers = n_dense_first_layers
+
+        # MTP speculative decoding
+        self.num_nextn_predict_layers = num_nextn_predict_layers
+
+        # Validate the correctness of rotary position embeddings parameters
+        # BC: if there is a 'type' field, move it to 'rope_type'.
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+        # Motif applies the YaRN mscale on the attention softmax scale
+        # (DeepSeek-style, full-attention layers only), never on the cos/sin table.
+        # Default attention_factor=1.0 so transformers' YaRN rope init does not ALSO
+        # scale cos/sin (which would double-apply mscale); apply_yarn_scaling=False
+        # makes that intent explicit for the modeling code (MotifRotaryEmbedding).
+        if self.rope_scaling is not None and self.rope_scaling.get("rope_type") == "yarn":
+            self.rope_scaling.setdefault("attention_factor", 1.0)
+            self.rope_scaling.setdefault("apply_yarn_scaling", False)
+        if callable(getattr(type(self), "validate_rope", None)):
+            self.validate_rope()
+
+        # The per-expert PolyNorm activation (GroupedPolyNorm) is only correct via
+        # the EAGER experts loop: the grouped_mm / batched_mm interfaces call
+        # _apply_gate once over all expert-sorted tokens with no per-expert index,
+        # so they cannot apply per-expert coefficients. Force eager dispatch for
+        # MoE Motif models unless the caller explicitly overrides it. (This also
+        # covers ROCm, where torch._grouped_mm has no runtime kernel.)
+        if self.num_experts > 0 and "experts_implementation" not in kwargs:
+            kwargs["experts_implementation"] = "eager"
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        logger.info(f" kwargs : {kwargs}")

modeling_motif.py

@@ -0,0 +1,1690 @@
+import copy
+import math
+from typing import Callable, Literal, Optional, Tuple
+
+import einops
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import CrossEntropyLoss
+from transformers.activations import ACT2CLS as _ACT2CLS
+from transformers.activations import ClassInstantier
+from transformers.cache_utils import Cache, DynamicCache, StaticCache
+from transformers.generation import GenerationMixin
+from transformers.modeling_attn_mask_utils import AttentionMaskConverter
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.utils import auto_docstring, can_return_tuple, logging
+
+from .configuration_motif import MotifConfig
+
+logger = logging.get_logger(__name__)
+
+if hasattr(torch.version, "hip") and torch.version.hip is not None:
+    activation = None
+    logger.warning_once("Using HIP")
+    logger.warning_once("Due to the HIP, we do not utilize the kernel ops for precision.")
+    logger.warning_once("Using torch ops")
+    kernelRMSNorm = None
+    PolyNormKernel = None
+else:
+    logger.warning_once("Using CUDA")
+    try:
+        import kernels
+
+        activation = kernels.get_kernel("Motif-Technologies/activation")
+        kernelRMSNorm = activation.layers.RMSNorm
+        PolyNormKernel = activation.layers.PolyNorm
+    except Exception as e:
+        activation = None
+        kernelRMSNorm = None
+        PolyNormKernel = None
+        logger.warning_once(f"Failed to import kernel ops: {e}")
+        logger.warning_once("Using torch ops")
+
+
+class PolyNormTorch(torch.nn.Module):
+    """
+    A trainable activation function introduced in https://arxiv.org/html/2411.03884v1.
+    The code is copied from https://github.com/BryceZhuo/PolyCom?tab=readme-ov-file/README.md,
+    with the change `* torch.rsqrt` => `/ torch.sqrt`.
+    """
+
+    def __init__(self, eps=1e-6, sigmoid_weight: bool = True):
+        super(PolyNormTorch, self).__init__()
+        self.weight = torch.nn.Parameter(torch.ones(3) / 3)
+        self.bias = torch.nn.Parameter(torch.zeros(1))
+        self.eps = eps
+        self.sigmoid_weight = sigmoid_weight
+
+    def _norm(self, x):
+        return x / torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def _coeffs(self):
+        w = self.weight.float()
+        return torch.sigmoid(w) if self.sigmoid_weight else w
+
+    def _poly(self, x):
+        w = self._coeffs()
+        return w[0] * self._norm(x**3) + w[1] * self._norm(x**2) + w[2] * self._norm(x) + self.bias.float()
+
+    def forward(self, x):
+        orig_dtype = x.dtype
+        return self._poly(x.float()).to(orig_dtype)
+
+    def forward_mul(self, x, mul):
+        orig_dtype = x.dtype
+        return (self._poly(x.float()) * mul.float()).to(orig_dtype)
+
+
+PolyNorm = PolyNormKernel if PolyNormKernel is not None else PolyNormTorch
+CUSTOM_ACT2CLS = {"poly_norm": PolyNorm}
+ACT2CLS = {**_ACT2CLS, **CUSTOM_ACT2CLS}
+ACT2FN = ClassInstantier(ACT2CLS)
+
+
+class GroupedPolyNorm(nn.Module):
+    """Per-expert PolyNorm: weight [num_experts, 3], bias [num_experts, 1].
+
+    Mirrors titan's GroupedExpertsPolyNorm — each expert has independent
+    polynomial normalization coefficients.
+    """
+
+    def __init__(
+        self,
+        num_experts: int,
+        eps: float = 1e-6,
+        sigmoid_weight: bool = True,
+        bias_clamp: Optional[float] = None,
+        output_scale: float = 1.0,
+        hidden_clamp: Optional[float] = None,
+    ):
+        super().__init__()
+        self.num_experts = num_experts
+        self.eps = eps
+        self.sigmoid_weight = sigmoid_weight
+        self.bias_clamp = bias_clamp
+        self.output_scale = output_scale
+        self.hidden_clamp = hidden_clamp
+        self.weight = nn.Parameter(torch.ones(num_experts, 3) / 3)
+        self.bias = nn.Parameter(torch.zeros(num_experts, 1))
+
+    def _norm(self, x: torch.Tensor) -> torch.Tensor:
+        return x / torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward_single(self, x: torch.Tensor, mul: torch.Tensor, expert_idx: int) -> torch.Tensor:
+        orig_dtype = x.dtype
+        w = self.weight[expert_idx].float()
+        if self.sigmoid_weight:
+            w = torch.sigmoid(w)
+        b = self.bias[expert_idx].float()
+        if self.bias_clamp is not None:
+            b = b.clamp(-self.bias_clamp, self.bias_clamp)
+        xf = x.float()
+        mf = mul.float()
+        if self.hidden_clamp is not None:
+            xf = xf.clamp(-self.hidden_clamp, self.hidden_clamp)
+            mf = mf.clamp(-self.hidden_clamp, self.hidden_clamp)
+        poly = w[0] * self._norm(xf**3) + w[1] * self._norm(xf**2) + w[2] * self._norm(xf) + b
+        result = poly * mf
+        if self.hidden_clamp is not None:
+            result = result.clamp(-self.hidden_clamp, self.hidden_clamp)
+        result = result * self.output_scale
+        return result.to(orig_dtype)
+
+
+class MotifRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        MotifRMSNorm 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)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class MHCLayer(nn.Module):
+    """Manifold-constrained Hyper-Connections (MHC) layer.
+
+    Written inline in the HF model (no llm_training.layers.mhc dependency).
+    apply_h_res uses a pure-PyTorch einsum instead of the Triton kernel.
+
+    Reference: https://arxiv.org/abs/2512.24880
+    """
+
+    def __init__(
+        self,
+        expansion_rate: int,
+        num_dim: int,
+        identity_init: bool = False,
+        sinkhorn_iters: int = 20,
+        h_post_coeff: float = 2.0,
+    ):
+        super().__init__()
+        self.expansion_rate = expansion_rate
+        self.num_dim = num_dim
+        self.sinkhorn_iters = sinkhorn_iters
+        self.h_post_coeff = float(h_post_coeff)
+
+        E, D = expansion_rate, num_dim
+        self.proj_pre = nn.Linear(E * D, E, bias=False)
+        self.proj_post = nn.Linear(E * D, E, bias=False)
+        self.proj_res = nn.Linear(E * D, E * E, bias=False)
+
+        RMSNorm = kernelRMSNorm if kernelRMSNorm is not None else MotifRMSNorm
+        self.rms_norm = RMSNorm(E * D, eps=1e-6)
+
+        self.bias_pre = nn.Parameter(torch.empty(E))
+        self.bias_post = nn.Parameter(torch.empty(E))
+        self.bias_res = nn.Parameter(torch.empty(E, E))
+        self.alpha_pre = nn.Parameter(torch.empty(1))
+        self.alpha_post = nn.Parameter(torch.empty(1))
+        self.alpha_res = nn.Parameter(torch.empty(1))
+
+        self._init_weights(identity_init)
+
+    def _init_weights(self, identity_init: bool) -> None:
+        if hasattr(self.rms_norm, "reset_parameters"):
+            self.rms_norm.reset_parameters()
+        if identity_init:
+            nn.init.zeros_(self.alpha_pre)
+            nn.init.zeros_(self.alpha_post)
+            nn.init.zeros_(self.alpha_res)
+            nn.init.xavier_uniform_(self.proj_pre.weight)
+            nn.init.xavier_uniform_(self.proj_post.weight)
+            nn.init.xavier_uniform_(self.proj_res.weight)
+            uniform_weight = 1.0 / self.expansion_rate
+            bias_pre_value = math.log(uniform_weight / (1 - uniform_weight)) if 0 < uniform_weight < 1 else 0.0
+            nn.init.constant_(self.bias_pre, bias_pre_value)
+            nn.init.zeros_(self.bias_post)
+            nn.init.constant_(self.bias_res, -10.0)
+            self.bias_res.data.fill_diagonal_(0.0)
+        else:
+            nn.init.normal_(self.alpha_pre, mean=0.0, std=0.1)
+            nn.init.normal_(self.alpha_post, mean=0.0, std=0.1)
+            nn.init.normal_(self.alpha_res, mean=0.0, std=0.1)
+            nn.init.xavier_uniform_(self.proj_pre.weight)
+            nn.init.xavier_uniform_(self.proj_post.weight)
+            nn.init.xavier_uniform_(self.proj_res.weight)
+            nn.init.zeros_(self.bias_pre)
+            nn.init.zeros_(self.bias_post)
+            nn.init.normal_(self.bias_res, mean=0.0, std=0.1)
+
+    def _sinkhorn_knopp_batch(self, matrix: torch.Tensor) -> torch.Tensor:
+        orig_dtype = matrix.dtype
+        # Run Sinkhorn-Knopp in float32 (bf16/fp16 exp() is numerically unstable)
+        m = matrix.float().clamp(-20.0, 20.0).exp()
+        for _ in range(self.sinkhorn_iters):
+            m = m / m.sum(dim=-1, keepdim=True).clamp(min=1e-8)
+            m = m / m.sum(dim=-2, keepdim=True).clamp(min=1e-8)
+        return m.to(orig_dtype)
+
+    def forward(self, x: torch.Tensor):
+        batch_size, seq_len, expansion_rate, dim = x.shape
+        x_reshaped = x.reshape(batch_size, seq_len, expansion_rate * dim)
+        x_norm = self.rms_norm(x_reshaped)
+
+        # Cast projection outputs to float32 (paper §4.3.1)
+        proj_pre_out = self.proj_pre(x_norm).float()
+        proj_post_out = self.proj_post(x_norm).float()
+        proj_res_out = self.proj_res(x_norm).float().reshape(batch_size, seq_len, expansion_rate, expansion_rate)
+
+        h_pre = torch.sigmoid((self.alpha_pre * proj_pre_out + self.bias_pre).clamp(-10.0, 10.0))
+        h_post = self.h_post_coeff * torch.sigmoid((self.alpha_post * proj_post_out + self.bias_post).clamp(-10.0, 10.0))
+        h_res = self._sinkhorn_knopp_batch(self.alpha_res * proj_res_out + self.bias_res)
+
+        return h_pre, h_post, h_res
+
+    @classmethod
+    def apply_h_res(cls, x: torch.Tensor, h_res: torch.Tensor) -> torch.Tensor:
+        """h_res: (B, S, E, E), x: (B, S, E, D) -> (B, S, E, D)."""
+        return torch.einsum("bsij,bsjd->bsid", h_res, x.float()).to(x.dtype)
+
+    @classmethod
+    def apply_h_pre(cls, x: torch.Tensor, h_pre: torch.Tensor) -> torch.Tensor:
+        """Weighted sum over expansion dim: (B, S, E, D) -> (B, S, D)."""
+        return (x * h_pre.unsqueeze(-1)).sum(dim=2).to(x.dtype)
+
+    @classmethod
+    def apply_h_post(cls, x: torch.Tensor, h_post: torch.Tensor) -> torch.Tensor:
+        """Expand: (B, S, D) -> (B, S, E, D)."""
+        return (h_post.unsqueeze(-1) * x.unsqueeze(2)).to(x.dtype)
+
+    def extra_repr(self) -> str:
+        return f"expansion_rate={self.expansion_rate}, sinkhorn_iters={self.sinkhorn_iters}"
+
+
+def _compute_yarn_inv_freq(
+    dim: int,
+    end: int,
+    theta: float,
+    original_seq_len: int,
+    rope_factor: float,
+    beta_fast: float = 32.0,
+    beta_slow: float = 1.0,
+    device=None,
+) -> torch.Tensor:
+    """YaRN frequency interpolation, matching llm-training's
+    ``precompute_freqs_cis_yarn`` (model.py) exactly. Returns ``inv_freq`` of
+    shape ``[dim // 2]`` (the per-position freqs table is built from it in
+    ``MotifRotaryEmbedding.forward``). Only the frequencies are interpolated —
+    the cos/sin table is never scaled by mscale (Motif folds mscale into the
+    attention softmax scale of full-attention layers instead)."""
+
+    def find_correction_dim(num_rotations, d, base, max_seq_len):
+        return d * math.log(max_seq_len / (num_rotations * 2 * math.pi)) / (2 * math.log(base))
+
+    def find_correction_range(low_rot, high_rot, d, base, max_seq_len):
+        low = math.floor(find_correction_dim(low_rot, d, base, max_seq_len))
+        high = math.ceil(find_correction_dim(high_rot, d, base, max_seq_len))
+        return max(low, 0), min(high, d - 1)
+
+    def linear_ramp_factor(lo, hi, d):
+        if lo == hi:
+            hi += 0.001
+        ramp = (torch.arange(d, dtype=torch.float32, device=device) - lo) / (hi - lo)
+        return torch.clamp(ramp, 0, 1)
+
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim))
+    if end > original_seq_len:
+        low, high = find_correction_range(beta_fast, beta_slow, dim, theta, original_seq_len)
+        smooth = 1 - linear_ramp_factor(low, high, dim // 2)
+        freqs = freqs / rope_factor * (1 - smooth) + freqs * smooth
+    return freqs
+
+
+class MotifRotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor
+
+    def __init__(self, config: MotifConfig, device=None, rope_head_dim: Optional[int] = None):
+        super().__init__()
+        # BC: "rope_type" was originally "type"
+        if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
+            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        # Use rope_head_dim if provided (e.g. for GDLA which only applies RoPE to qk_rope_head_dim dims)
+        effective_head_dim = (
+            rope_head_dim
+            if rope_head_dim is not None
+            else (config.head_dim if config.head_dim is not None else config.hidden_size // config.num_attention_heads)
+        )
+        if self.rope_type == "default":
+            self.rope_init_fn = None
+            inv_freq = 1.0 / (
+                config.rope_theta
+                ** (torch.arange(0, effective_head_dim, 2, dtype=torch.int64).float().to(device) / effective_head_dim)
+            )
+            self.attention_scaling = 1.0
+
+        elif self.rope_type == "yarn":
+            self.rope_init_fn = None
+            rope_scaling = config.rope_scaling if isinstance(config.rope_scaling, dict) else {}
+            factor = float(rope_scaling.get("factor", getattr(config, "rope_factor", 1.0)))
+            original_seq_len = int(
+                rope_scaling.get(
+                    "original_max_position_embeddings",
+                    getattr(config, "original_seq_len", config.max_position_embeddings),
+                )
+            )
+            beta_fast = float(rope_scaling.get("beta_fast", 32))
+            beta_slow = float(rope_scaling.get("beta_slow", 1))
+            theta = float(rope_scaling.get("rope_theta", config.rope_theta))
+            inv_freq = _compute_yarn_inv_freq(
+                effective_head_dim,
+                config.max_position_embeddings,
+                theta,
+                original_seq_len,
+                factor,
+                beta_fast,
+                beta_slow,
+                device,
+            )
+            self.attention_scaling = 1.0
+
+        else:
+            self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+            inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+            rope_scaling = config.rope_scaling if isinstance(config.rope_scaling, dict) else {}
+            if not rope_scaling.get("apply_yarn_scaling", False):
+                self.attention_scaling = 1.0
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    def _dynamic_frequency_update(self, position_ids, device):
+        seq_len = torch.max(position_ids) + 1
+        if seq_len > self.max_seq_len_cached:
+            if self.rope_init_fn is not None:
+                inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, seq_len=seq_len)
+                self.register_buffer("inv_freq", inv_freq, persistent=False)
+            self.max_seq_len_cached = seq_len
+
+        if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len:
+            self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
+            self.max_seq_len_cached = self.original_max_seq_len
+
+    @torch.no_grad()
+    def forward(self, x, position_ids):
+        if "dynamic" in self.rope_type:
+            self._dynamic_frequency_update(position_ids, device=x.device)
+
+        if (not getattr(self, "_inv_ready", False)) or self.inv_freq.is_meta \
+                or float(self.inv_freq.detach().abs().sum()) == 0.0:
+            _rhd = self.inv_freq.shape[0] * 2
+            if self.rope_type == "yarn":
+                _rs = self.config.rope_scaling if isinstance(self.config.rope_scaling, dict) else {}
+                _iv = _compute_yarn_inv_freq(
+                    _rhd, self.config.max_position_embeddings,
+                    float(_rs.get("rope_theta", self.config.rope_theta)),
+                    int(_rs.get("original_max_position_embeddings",
+                                getattr(self.config, "original_seq_len", self.config.max_position_embeddings))),
+                    float(_rs.get("factor", getattr(self.config, "rope_factor", 1.0))),
+                    float(_rs.get("beta_fast", 32)), float(_rs.get("beta_slow", 1)), x.device)
+            else:  # plain RoPE (SWA layers: rope_theta=swa_rope_theta, no YaRN)
+                _iv = 1.0 / (self.config.rope_theta ** (
+                    torch.arange(0, _rhd, 2, dtype=torch.int64).float().to(x.device) / _rhd))
+            self.register_buffer("inv_freq", _iv.to(device=x.device, dtype=torch.float32), persistent=False)
+            self.original_inv_freq = self.inv_freq
+            self._inv_ready = True
+        inv_freq_expanded = self.inv_freq[None, :, None].to(device=x.device, dtype=torch.float32).expand(position_ids.shape[0], -1, 1)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def rotate_half(x):
+    """
+    Rotates half of the dimensions of the input tensor using torch.roll and in-place negation.
+
+    Args:
+    x (torch.Tensor): The input tensor.
+
+    Returns:
+    torch.Tensor: A tensor where the latter half of the dimensions are negated
+                  and moved before the first half.
+    """
+    half_size = x.shape[-1] // 2
+    rotated_tensor = torch.roll(x, shifts=-half_size, dims=-1)
+    rotated_tensor[..., :half_size] *= -1
+
+    return rotated_tensor
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """
+    Applies rotary position embeddings to the input tensors.
+    Args:
+        q (torch.Tensor): Query tensor of shape (B, NH, S, D_KV).
+        k (torch.Tensor): Key tensor of shape (B, NH, S, D_KV).
+        cos (torch.Tensor): Cosine values for rotary embedding, shape (B, S, D) from MotifRotaryEmbedding.
+        sin (torch.Tensor): Sine values for rotary embedding, shape (B, S, D) from MotifRotaryEmbedding.
+        position_ids: Unused, kept for API compatibility.
+        unsqueeze_dim (int, optional): Dimension along which `cos` and `sin` are unsqueezed.
+            Defaults to 1 (head dimension).
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Transformed query and key tensors.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def apply_rotary_pos_emb_single(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> torch.Tensor:
+    """Apply RoPE to a single tensor in (B, S, NH, D) format.
+
+    Used by GDLA to apply positional encoding only to the rope portion of Q/K.
+    cos/sin shape: (B, S, D) — broadcast over NH dimension.
+    """
+    cos = cos.unsqueeze(2)  # (B, S, 1, D)
+    sin = sin.unsqueeze(2)  # (B, S, 1, D)
+    return x * cos + rotate_half(x) * sin
+
+
+class MotifMLP(nn.Module):
+    def __init__(self, config, intermediate_size: int | None = None):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = intermediate_size if intermediate_size is not None else config.intermediate_size
+
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        if config.hidden_act == "poly_norm":
+            self.act_fn = PolyNormTorch(sigmoid_weight=getattr(config, "polynorm_sigmoid_weight", True))
+        else:
+            self.act_fn = ACT2FN[config.hidden_act]
+        self.hidden_clamp = getattr(config, "hidden_clamp", None)
+        self.polynorm_output_scale = float(getattr(config, "polynorm_output_scale", 1.0))
+
+    def forward(self, hidden_state):
+        gate = self.gate_proj(hidden_state)
+        up = self.up_proj(hidden_state)
+        if isinstance(self.act_fn, PolyNormTorch):
+            if self.hidden_clamp is not None:
+                gate = gate.clamp(-self.hidden_clamp, self.hidden_clamp)
+                up = up.clamp(-self.hidden_clamp, self.hidden_clamp)
+            hidden_state = self.act_fn.forward_mul(gate, up)
+            if self.polynorm_output_scale != 1.0:
+                hidden_state = hidden_state * self.polynorm_output_scale
+        else:
+            hidden_state = self.act_fn(gate) * up
+        return self.down_proj(hidden_state)
+
+
+def repeat_kv(hidden_states: torch.Tensor, dim: int, n_rep: int) -> torch.Tensor:
+    return torch.repeat_interleave(hidden_states, dim=dim, repeats=n_rep)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs,
+):
+    """Eager attention forward compatible with ALL_ATTENTION_FUNCTIONS interface.
+    Expects query/key/value in [batch, num_heads, seq_len, head_dim] format.
+    """
+    attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+class MotifGDLAttention(nn.Module):
+    """Grouped Differential Latent Attention (GDLA) for HF Transformers.
+
+    Ports GDLAttention from torchtitan. Uses low-rank Q and KV projections
+    (MLA-style) with RoPE applied only to the qk_rope_head_dim dimensions.
+    Only diff_v2=True is fully supported (the motif3 configuration).
+    """
+
+    def __init__(self, config: MotifConfig, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim if config.head_dim is not None else self.hidden_size // self.num_heads
+        self.is_causal = True
+        self.attention_dropout = config.attention_dropout
+
+        # Head split
+        self.num_noise_heads = config.num_noise_heads
+        self.grouped_ratio = (self.num_heads - self.num_noise_heads) // self.num_noise_heads
+        self.n_signal_heads = self.grouped_ratio * self.num_noise_heads
+
+        # GDLA dimensions
+        self.q_lora_rank = config.q_lora_rank
+        self.kv_lora_rank = config.kv_lora_rank
+        self.qk_rope_head_dim = config.qk_rope_head_dim if config.qk_rope_head_dim is not None else self.head_dim // 2
+        self.qk_nope_head_dim = self.head_dim - self.qk_rope_head_dim
+        self.v_head_dim = config.v_head_dim if config.v_head_dim is not None else self.head_dim
+        self.diff_v2 = getattr(config, "diff_v2", True)
+        self.sliding_window = None
+        is_swa_layer = False
+        if config.use_sliding_window and getattr(config, "sliding_window", None) is not None:
+            pattern = getattr(config, "sliding_window_pattern", "interleave")
+            period = getattr(config, "sliding_window_period", 2)
+            effective_window = config.sliding_window + 1
+            if pattern == "all":
+                self.sliding_window = effective_window
+                is_swa_layer = True
+            elif pattern == "interleave" and layer_idx % period != 0:
+                self.sliding_window = effective_window
+                is_swa_layer = True
+        self.is_swa_layer = is_swa_layer
+        self.scaling = self.head_dim**-0.5
+        original_seq_len = getattr(config, "original_seq_len", 32768)
+        rope_factor = getattr(config, "rope_factor", 1.0)
+        mscale = getattr(config, "mscale", 1.0)
+        if (not is_swa_layer) and config.max_position_embeddings > original_seq_len:
+            mscale_val = 0.1 * mscale * math.log(rope_factor) + 1.0
+            self.scaling = self.scaling * mscale_val * mscale_val
+
+        # Output gating
+        self.elementwise_attn_output_gate = getattr(config, "elementwise_attn_output_gate", False)
+        self.headwise_attn_output_gate = getattr(config, "headwise_attn_output_gate", False)
+
+        # Required by transformers SDPA interface for GQA repeat_kv dispatch
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+
+        RMSNorm = kernelRMSNorm if kernelRMSNorm is not None else MotifRMSNorm
+
+        # Query LoRA
+        self.wq_a = nn.Linear(self.hidden_size, self.q_lora_rank, bias=False)
+        self.q_norm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
+
+        if self.elementwise_attn_output_gate:
+            # Separate gate projection; for V2: n_signal_heads gates; for V1: n_signal_heads*2
+            gate_extra = self.n_signal_heads if self.diff_v2 else self.n_signal_heads * 2
+            self.wq_b = nn.Linear(self.q_lora_rank, self.num_heads * self.head_dim, bias=False)
+            self.wq_b_gate = nn.Linear(self.q_lora_rank, gate_extra * self.v_head_dim, bias=False)
+        else:
+            self.wq_b = nn.Linear(self.q_lora_rank, self.num_heads * self.head_dim, bias=False)
+            self.wq_b_gate = None
+
+        # KV LoRA: output kv_lora_rank + qk_rope_head_dim
+        self.wkv_a = nn.Linear(self.hidden_size, self.kv_lora_rank + self.qk_rope_head_dim, bias=False)
+
+        if self.diff_v2:
+            self.kv_norm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
+            self.wkv_b = nn.Linear(
+                self.kv_lora_rank,
+                self.num_key_value_heads * (self.qk_nope_head_dim + self.v_head_dim),
+                bias=False,
+            )
+            self.lambda_proj = nn.Linear(self.hidden_size, self.n_signal_heads, bias=False)
+            self.wo = nn.Linear(self.n_signal_heads * self.v_head_dim, self.hidden_size, bias=False)
+        else:
+            raise NotImplementedError(
+                "GDLA V1 (diff_v2=False) is not yet implemented in HF. Use attention_cls='gdla' only with diff_v2=True."
+            )
+
+        self.swa_rotary_emb = None
+        swa_rope_theta = getattr(config, "swa_rope_theta", None)
+        if self.is_swa_layer and swa_rope_theta is not None:
+            swa_config = copy.copy(config)
+            swa_config.rope_scaling = None
+            swa_config.rope_theta = swa_rope_theta
+            self.swa_rotary_emb = MotifRotaryEmbedding(swa_config, rope_head_dim=self.qk_rope_head_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        bsz, q_len, _ = hidden_states.size()
+        _wqa = F.linear(hidden_states.float(), self.wq_a.weight.float())
+        q_latent = self.q_norm(_wqa)  # fp32
+        q = F.linear(q_latent, self.wq_b.weight.float()).view(
+            bsz, q_len, self.num_heads, self.head_dim).to(hidden_states.dtype)
+        q_latent = q_latent.to(hidden_states.dtype)  # bf16 for the gate below
+
+        # Gate score (elementwise)
+        gate_score = None
+        if self.elementwise_attn_output_gate and self.wq_b_gate is not None:
+            gate_score = self.wq_b_gate(q_latent).view(bsz, q_len, -1, self.v_head_dim)
+
+        # Split Q into nope (no-positional) and rope parts
+        q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+
+        # KV path: project to kv_lora_rank + qk_rope_head_dim, then split
+        kv_raw = self.wkv_a(hidden_states)  # (bsz, q_len, kv_lora_rank + qk_rope_head_dim)
+        kv_latent, k_pe = torch.split(kv_raw, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+
+        if self.swa_rotary_emb is not None:
+            cos, sin = self.swa_rotary_emb(hidden_states, position_ids)
+        else:
+            cos, sin = position_embeddings
+        _rope_dtype = q_pe.dtype
+        cos_f, sin_f = cos.float(), sin.float()
+        q_pe = apply_rotary_pos_emb_single(q_pe.float(), cos_f, sin_f).to(_rope_dtype)
+        k_pe = apply_rotary_pos_emb_single(k_pe.unsqueeze(2).float(), cos_f, sin_f).to(_rope_dtype)
+
+        # Reconstruct full Q with nope + rope
+        q_total = torch.cat([q_nope, q_pe], dim=-1)
+
+        # KV projection: norm -> project -> split k_nope and v
+        kv_latent = kv_latent.contiguous()
+        kv_proj = self.wkv_b(self.kv_norm(kv_latent))
+        kv_proj = kv_proj.view(bsz, q_len, self.num_key_value_heads, self.qk_nope_head_dim + self.v_head_dim)
+        k_nope, v = torch.split(kv_proj, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+
+        # Assemble full K: k_nope + k_pe (broadcast shared rope over all kv heads)
+        k_full = torch.cat([k_nope, k_pe.expand(-1, -1, self.num_key_value_heads, -1)], dim=-1)
+
+        # Lambda (input-dependent for V2)
+        lambda_full = self.lambda_proj(hidden_states)  # (bsz, q_len, n_signal_heads)
+
+        # Transpose to (B, H, S, D) before cache (DynamicCache expects this format)
+        k_full = k_full.transpose(1, 2)  # (bsz, n_kv_heads, q_len, head_dim)
+        v = v.transpose(1, 2)  # (bsz, n_kv_heads, q_len, v_head_dim)
+
+        # KV cache
+        if past_key_value is not None:
+            cache_kwargs = {"cache_position": cache_position}
+            k_full, v = past_key_value.update(k_full, v, self.layer_idx, cache_kwargs)
+
+        dropout_rate = 0.0 if not self.training else self.attention_dropout
+
+        # Cast dtype if needed
+        input_dtype = q_total.dtype
+        if input_dtype == torch.float32:
+            if torch.is_autocast_enabled():
+                target_dtype = torch.get_autocast_gpu_dtype()
+            elif hasattr(self.config, "_pre_quantization_dtype"):
+                target_dtype = self.config._pre_quantization_dtype
+            else:
+                target_dtype = self.wq_b.weight.dtype
+            q_total = q_total.to(target_dtype)
+            k_full = k_full.to(target_dtype)
+            v = v.to(target_dtype)
+
+        # Attention: pad v to head_dim if v_head_dim != head_dim (flash_attn compatibility)
+        need_v_pad = self.v_head_dim != self.head_dim
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        # k_full and v are already (B, H, S, D); transpose q for interface
+        q_t = q_total.transpose(1, 2)  # (bsz, n_heads, q_len, head_dim)
+        k_t = k_full  # (bsz, n_kv_heads, kv_seq_len, head_dim)
+        v_t = v  # (bsz, n_kv_heads, kv_seq_len, v_head_dim)
+
+        if need_v_pad:
+            v_t = F.pad(v_t, [0, self.head_dim - self.v_head_dim])
+
+        # Truncate mask to actual kv length (sliding window cache may return fewer tokens than target_length)
+        if attention_mask is not None:
+            attention_mask = attention_mask[:, -k_t.shape[-2] :]
+
+        attn_out, _ = attention_interface(
+            self,
+            q_t,
+            k_t,
+            v_t,
+            attention_mask,
+            dropout=dropout_rate,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            is_causal=self.is_causal,
+            **kwargs,
+        )
+
+        # Normalize to (B, S, H, D)
+        if attn_out.shape[1] == self.num_heads:  # (B, H, S, D) from SDPA
+            attn_out = attn_out.transpose(1, 2)
+        # Now (B, S, H, D)
+
+        if need_v_pad:
+            attn_out = attn_out[..., : self.v_head_dim].contiguous()
+
+        # Split heads into signal and noise groups
+        num_groups = self.num_noise_heads  # n_heads // (grouped_ratio + 1)
+        attn_reshaped = einops.rearrange(
+            attn_out,
+            "b s (g gs) d -> b s g gs d",
+            g=num_groups,
+            gs=self.grouped_ratio + 1,
+        )
+        attn1 = attn_reshaped[:, :, :, : self.grouped_ratio, :].reshape(bsz, q_len, -1, self.v_head_dim)
+        attn2_group = attn_reshaped[:, :, :, self.grouped_ratio :, :].reshape(bsz, q_len, num_groups, self.v_head_dim)
+        attn2 = repeat_kv(attn2_group, 2, self.grouped_ratio)  # (bsz, q_len, n_signal_heads, v_head_dim)
+
+        lambda_scale = torch.sigmoid(lambda_full.float()).to(attn1.dtype).unsqueeze(-1)
+        attn_output = attn1 - lambda_scale * attn2
+
+        if gate_score is not None:
+            attn_output = attn_output * torch.sigmoid(gate_score)
+
+        # Output projection
+        attn_output = attn_output.reshape(bsz, q_len, -1)
+        attn_output = self.wo(attn_output)
+
+        return attn_output, None, past_key_value
+
+
+class TokenChoiceTopKRouter(nn.Module):
+    """This class implements token-choice routing. In token-choice top-K routing, each token is
+        routed to top K experts based on the router scores.
+
+    Args:
+        dim (int): Dimension of input tokens.
+        num_experts (int): Number of experts in each moe layer.
+        experts_top_k (int): Number of experts each token will be routed to in token-choice routing.
+        score_func (Literal["softmax", "sigmoid"]): Whether to use sigmoid or softmax for router scores.
+        route_norm (bool): Whether to normalize the routing scores when using sigmoid.
+        route_scale (float): Scaling factor applied to the routing scores.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_experts: int,
+        experts_top_k: int,
+        score_func: Literal["softmax", "sigmoid"],
+        route_norm: bool,
+        route_scale: float,
+        _debug_force_load_balance: bool = False,
+    ):
+        super().__init__()
+        self.gate = nn.Linear(hidden_size, num_experts, bias=False)
+        self.num_experts = num_experts
+        self.experts_top_k = experts_top_k
+        self.score_func = score_func
+        self.route_norm = route_norm
+        self.route_scale = route_scale
+        self._debug_force_load_balance = _debug_force_load_balance
+
+    def _debug_force_load_balance_routing(self, scores: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Balanced round-robin expert assignment.
+        Returns (selected_experts_indices [N, K] LongTensor, top_scores [N, K] FloatTensor).
+        """
+        n_tokens = scores.size(0)
+        # Round-robin indices with exact balance
+        selected_experts_indices = (
+            torch.arange(n_tokens * self.experts_top_k, device=scores.device, dtype=torch.int64).reshape(
+                n_tokens, self.experts_top_k
+            )
+            % self.num_experts
+        )
+        top_scores = scores.gather(dim=1, index=selected_experts_indices)  # [N,K]
+        return selected_experts_indices, top_scores
+
+    def forward(
+        self, x: torch.Tensor, expert_bias: torch.Tensor | None = None
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            x (torch.Tensor): Input tensor with shape ``(bs*slen, dim)``.
+            expert_bias (torch.Tensor | None, optional): Optional bias tensor for experts with shape ``(num_experts,)``.
+                Used for load balancing. Defaults to None.
+
+        Returns:
+            tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+                - top_scores (torch.Tensor):
+                    Routing scores for selected experts with shape ``(bs*slen, experts_top_k)``.
+                - selected_experts_indices (torch.Tensor):
+                    Expert indices selected for each token with shape ``(bs*slen, experts_top_k)``.
+                - num_tokens_per_expert (torch.Tensor):
+                    Number of tokens assigned to each expert with shape ``(num_experts,)``.
+        """
+        scores = F.linear(x.to(torch.float32), self.gate.weight.to(torch.float32))
+
+        # By default, sigmoid or softmax is performed in float32 to avoid loss explosion
+        if self.score_func == "sigmoid":
+            scores = torch.sigmoid(scores.to(torch.float32))
+        elif self.score_func == "softmax":
+            scores = F.softmax(scores.to(torch.float32), dim=1)
+        else:
+            raise NotImplementedError(f"Unknown score function {self.score_func}")
+
+
+        if expert_bias is not None:
+            _, selected_experts_indices = torch.topk(scores + expert_bias, k=self.experts_top_k, dim=1)
+            top_scores = scores.gather(dim=1, index=selected_experts_indices)
+        else:
+            top_scores, selected_experts_indices = torch.topk(scores, k=self.experts_top_k, dim=1)
+
+        # debug override: balanced round-robin routing
+        if self._debug_force_load_balance:
+            (
+                selected_experts_indices,
+                top_scores,
+            ) = self._debug_force_load_balance_routing(scores)
+
+        if self.route_norm:
+            denominator = top_scores.sum(dim=-1, keepdim=True) + 1e-20
+            top_scores = top_scores / denominator
+        top_scores = top_scores * self.route_scale
+
+        num_tokens_per_expert = torch.bincount(selected_experts_indices.view(-1).long(), minlength=self.num_experts).to(
+            dtype=torch.float32
+        )
+
+        return top_scores, selected_experts_indices, num_tokens_per_expert
+
+    def init_weights(self, init_std: float):
+        nn.init.trunc_normal_(self.gate.weight, mean=0.0, std=init_std)
+
+
+class MotifExperts(nn.Module):
+    """Collection of expert weights stored as fused 3D tensors.
+
+    NOTE: the @use_experts_implementation decorator was intentionally removed.
+    That decorator dispatches forward() through config._experts_implementation,
+    whose default (grouped_mm) calls _apply_gate once over all expert-sorted
+    tokens with no per-expert index — incompatible with per-expert PolyNorm. We
+    use the explicit eager per-expert loop below directly, with no dispatch."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.num_experts = config.num_experts
+        self.hidden_size = config.hidden_size
+        moe_intermediate = getattr(config, "moe_intermediate_size", config.intermediate_size)
+        self.intermediate_dim = moe_intermediate
+
+        # Fused gate+up: [num_experts, 2*intermediate, hidden_size]
+        self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_size))
+        # Down projection: [num_experts, hidden_size, intermediate]
+        self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_size, self.intermediate_dim))
+        if config.hidden_act == "poly_norm":
+            self.act_fn = GroupedPolyNorm(
+                self.num_experts,
+                sigmoid_weight=getattr(config, "polynorm_sigmoid_weight", True),
+                bias_clamp=getattr(config, "polynorm_bias_clamp", None),
+                output_scale=float(getattr(config, "polynorm_output_scale", 1.0)),
+                hidden_clamp=getattr(config, "hidden_clamp", None),
+            )
+        else:
+            self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        top_k_index: torch.Tensor,
+        top_k_weights: torch.Tensor,
+    ) -> torch.Tensor:
+        """Eager expert dispatch (loops over experts).
+
+        Args:
+            hidden_states: [total_tokens, hidden_size]
+            top_k_index: [total_tokens, top_k] expert indices
+            top_k_weights: [total_tokens, top_k] routing weights
+
+        Returns:
+            [total_tokens, hidden_size]
+        """
+        final_hidden_states = torch.zeros_like(hidden_states, dtype=torch.float32)
+        expert_mask = F.one_hot(top_k_index, num_classes=self.num_experts).permute(2, 1, 0)
+
+        for expert_idx in range(self.num_experts):
+            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
+            if token_idx.shape[0] == 0:
+                continue
+
+            current_state = hidden_states[token_idx]
+            gate_up = current_state @ self.gate_up_proj[expert_idx].T  # [T, 2*I]
+            current_hidden = self._apply_gate(gate_up, expert_idx) @ self.down_proj[expert_idx].T  # [T, H]
+
+            current_hidden = current_hidden.float() * top_k_weights[token_idx, top_k_pos, None].float()
+            final_hidden_states.index_add_(0, token_idx, current_hidden)
+
+        return final_hidden_states  # fp32; MoE.forward downcasts after adding shared
+
+    def _apply_gate(self, gate_up_output: torch.Tensor, expert_idx: int = 0) -> torch.Tensor:
+        gate, up = gate_up_output.chunk(2, dim=-1)
+        gate = gate.contiguous()
+        if isinstance(self.act_fn, GroupedPolyNorm):
+            return self.act_fn.forward_single(gate, up, expert_idx)
+        return self.act_fn(gate) * up
+
+
+class MoE(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        self.num_experts = config.num_experts
+        self.experts_top_k = config.experts_top_k
+
+        self.experts = MotifExperts(config)
+
+        self.router = TokenChoiceTopKRouter(
+            hidden_size=config.hidden_size,
+            num_experts=config.num_experts,
+            experts_top_k=config.experts_top_k,
+            score_func=config.score_func,
+            route_norm=config.route_norm,
+            route_scale=config.route_scale,
+            _debug_force_load_balance=config._debug_force_load_balance,
+        )
+
+        moe_intermediate = getattr(config, "moe_intermediate_size", config.intermediate_size)
+        self.shared_experts = (
+            MotifMLP(config, intermediate_size=moe_intermediate) if config.num_shared_experts > 0 else None
+        )
+        self.score_before_experts = config.score_before_experts
+        self.load_balance_coeff = config.load_balance_coeff
+        if self.load_balance_coeff is not None:
+            assert self.load_balance_coeff > 0.0
+            self.expert_bias = nn.Parameter(torch.zeros(config.num_experts, dtype=torch.float32), requires_grad=False)
+        else:
+            self.expert_bias = None
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            x (torch.Tensor): Input tensor with shape ``(bs, slen, dim)``.
+
+        Returns:
+            tuple: (output tensor (bs, slen, dim), router_logits (bs*slen, num_experts))
+        """
+        bs, slen, dim = x.shape
+        x = x.view(-1, dim)
+
+        # Route tokens
+        top_scores, selected_experts_indices, num_tokens_per_expert = self.router(x, self.expert_bias)
+        if self.score_before_experts:
+            final_hidden_states = self._score_before_forward(x, selected_experts_indices, top_scores)
+        else:
+            final_hidden_states = self.experts(x, selected_experts_indices, top_scores)
+
+        if self.shared_experts is not None:
+            final_hidden_states = final_hidden_states + self.shared_experts(x).float()
+
+        router_logits = self.router.gate(x.view(-1, dim)) if hasattr(self.router, "gate") else None
+
+        return final_hidden_states.reshape(bs, slen, dim).to(x.dtype), router_logits
+
+    def _score_before_forward(
+        self,
+        hidden_states: torch.Tensor,
+        top_k_index: torch.Tensor,
+        top_k_weights: torch.Tensor,
+    ) -> torch.Tensor:
+        """Custom dispatch for score_before_experts mode.
+
+        Pre-weights inputs by routing scores before expert computation,
+        rather than weighting expert outputs (the standard approach).
+        """
+        final_hidden_states = torch.zeros_like(hidden_states)
+        expert_mask = F.one_hot(top_k_index, num_classes=self.num_experts).permute(2, 1, 0)
+
+        for expert_idx in range(self.num_experts):
+            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
+            if token_idx.shape[0] == 0:
+                continue
+
+            current_state = hidden_states[token_idx]
+            weights = top_k_weights[token_idx, top_k_pos, None]
+            # Pre-weight input
+            current_state = (current_state.to(torch.float32) * weights).to(hidden_states.dtype)
+
+            gate_up = current_state @ self.experts.gate_up_proj[expert_idx].T  # [T, 2*I]
+            current_hidden = self.experts._apply_gate(gate_up, expert_idx) @ self.experts.down_proj[expert_idx].T  # [T, H]
+
+            final_hidden_states.index_add_(0, token_idx, current_hidden.to(final_hidden_states.dtype))
+
+        return final_hidden_states
+
+    def init_weights(self, init_std: float, buffer_device: torch.device):
+        nn.init.trunc_normal_(self.experts.gate_up_proj, mean=0.0, std=0.02)
+        nn.init.trunc_normal_(self.experts.down_proj, mean=0.0, std=init_std)
+        self.router.init_weights(init_std)
+        if self.shared_experts is not None:
+            nn.init.trunc_normal_(self.shared_experts.gate_proj.weight, mean=0.0, std=0.02)
+            nn.init.trunc_normal_(self.shared_experts.up_proj.weight, mean=0.0, std=init_std)
+            nn.init.trunc_normal_(self.shared_experts.down_proj.weight, mean=0.0, std=init_std)
+        nn.init.zeros_(self.expert_bias)
+
+
+class MotifDecoderLayer(GradientCheckpointingLayer):
+    _ATTN_CLS = {
+        "gdla": MotifGDLAttention,
+    }
+
+    def __init__(self, config: MotifConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        attention_cls_name = getattr(config, "attention_cls", "basic")
+        attn_cls = self._ATTN_CLS.get(attention_cls_name)
+        if attn_cls is None:
+            raise ValueError(f"Unknown attention_cls={attention_cls_name!r}, expected one of {list(self._ATTN_CLS)}")
+        self.self_attn = attn_cls(config, layer_idx)
+
+        # n_dense_first_layers: first N layers always dense (no MoE)
+        n_dense_first = getattr(config, "n_dense_first_layers", 0)
+        self.moe_enabled = (
+            layer_idx >= n_dense_first and (layer_idx + 1) % config.interleave_moe_layer_step == 0
+            if config.interleave_moe_layer_step != 0
+            else False
+        )
+
+        if self.moe_enabled:
+            self.moe = MoE(config)
+            self.moe.layer_idx = layer_idx
+        else:
+            self.mlp = MotifMLP(config)
+
+        RMSNorm = kernelRMSNorm if kernelRMSNorm is not None else MotifRMSNorm
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # MHC (Manifold-constrained Hyper-Connections) layers
+        self.mhc_enabled = getattr(config, "mhc_enabled", False)
+        if self.mhc_enabled:
+            mhc_expansion_rate = config.mhc_expansion_rate
+            # h_post = (1 + mhc_h_post_alpha_end) * sigmoid(...); motif3 alpha_end=0 -> 1.0.
+            mhc_h_post_coeff = 1.0 + float(getattr(config, "mhc_h_post_alpha_end", 0.0))
+            self.mhc_attn = MHCLayer(
+                expansion_rate=mhc_expansion_rate,
+                num_dim=config.hidden_size,
+                identity_init=getattr(config, "mhc_identity_init", False),
+                sinkhorn_iters=getattr(config, "mhc_sinkhorn_iters", 20),
+                h_post_coeff=mhc_h_post_coeff,
+            )
+            self.mhc_ffn = MHCLayer(
+                expansion_rate=mhc_expansion_rate,
+                num_dim=config.hidden_size,
+                identity_init=getattr(config, "mhc_identity_init", False),
+                sinkhorn_iters=getattr(config, "mhc_sinkhorn_iters", 20),
+                h_post_coeff=mhc_h_post_coeff,
+            )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
+                `(batch, sequence_length)` where padding elements are indicated by 0.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
+            cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+                Indices depicting the position of the input sequence tokens in the sequence.
+            position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
+                Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
+                with `head_dim` being the embedding dimension of each attention head.
+            kwargs (`dict`, *optional*):
+                Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
+                into the model
+        """
+
+        if self.mhc_enabled:
+            return self._forward_with_mhc(
+                hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+            )
+
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        hidden_states, self_attn_weights, present_key_value = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+
+        router_logits = None
+        if self.moe_enabled:
+            hidden_states, router_logits = self.moe(hidden_states)
+        else:
+            hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        outputs += (router_logits,)
+
+        return outputs
+
+    def _forward_with_mhc(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> Tuple:
+        """MHC residual path. hidden_states: (batch, seq_len, expansion_rate, dim)."""
+        x = hidden_states
+
+        # === Attention sublayer with MHC ===
+        h_pre_attn, h_post_attn, h_res_attn = self.mhc_attn(x)
+
+        # Reduce for attention input: (B, S, E, D) -> (B, S, D)
+        x_reduced = MHCLayer.apply_h_pre(x, h_pre_attn)
+
+        attn_in = self.input_layernorm(x_reduced)
+        attn_out, self_attn_weights, present_key_value = self.self_attn(
+            hidden_states=attn_in,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+        )
+
+        _res_attn = torch.einsum("bsij,bsjd->bsid", h_res_attn, x.float())
+        _post_attn = h_post_attn.unsqueeze(-1) * attn_out.float().unsqueeze(2)
+        h = (_res_attn + _post_attn).to(x.dtype)
+
+        # === FFN sublayer with MHC ===
+        h_pre_ffn, h_post_ffn, h_res_ffn = self.mhc_ffn(h)
+
+        # Reduce for FFN input: (B, S, E, D) -> (B, S, D)
+        h_reduced = MHCLayer.apply_h_pre(h, h_pre_ffn)
+        n_out = self.post_attention_layernorm(h_reduced)
+
+        router_logits = None
+        if self.moe_enabled:
+            ffn_out, router_logits = self.moe(n_out)
+        else:
+            ffn_out = self.mlp(n_out)
+
+        # MHC FFN combine: out = H_res @ h + H_post * ffn_out, fp32 + single downcast.
+        _res_ffn = torch.einsum("bsij,bsjd->bsid", h_res_ffn, h.float())
+        _post_ffn = h_post_ffn.unsqueeze(-1) * ffn_out.float().unsqueeze(2)
+        out = (_res_ffn + _post_ffn).to(h.dtype)
+
+        outputs = (out,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        outputs += (router_logits,)
+
+        return outputs
+
+
+@auto_docstring
+class MotifPreTrainedModel(PreTrainedModel):
+    config_class = MotifConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MotifDecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _supports_attention_backend = True
+    _supports_cache_class = True
+    _supports_quantized_cache = True
+    _supports_static_cache = True
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data = torch.where(abs(module.weight.data) > 3 * std, 0, module.weight.data)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data = torch.where(abs(module.weight.data) > 3 * std, 0, module.weight.data)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, MoE):
+            module.init_weights(std, buffer_device=torch.device("cpu"))
+
+
+@auto_docstring
+class MotifModel(MotifPreTrainedModel):
+    """
+    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MotifDecoderLayer`]
+
+    Args:
+        config: MotifConfig
+    """
+
+    def __init__(self, config: MotifConfig):
+        super().__init__(config)
+        # Only flash_attention_2 supports GDLA's GQA + per-layer sliding window (see error).
+        if getattr(config, "_attn_implementation", None) != "flash_attention_2":
+            raise ValueError(
+                "Motif requires attn_implementation='flash_attention_2' (got "
+                f"{getattr(config, '_attn_implementation', None)!r}). The GDLA attention uses GQA "
+                "and per-layer sliding-window attention, which the eager backend (no GQA KV repeat) "
+                "and the sdpa backend (sliding-window mask mismatch past `sliding_window` tokens "
+                "during decode) do not handle correctly yet. Please load with "
+                "attn_implementation='flash_attention_2'."
+            )
+        self.padding_idx = getattr(config, "pad_token_id", None)
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [MotifDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self._attn_implementation = config._attn_implementation
+        RMSNorm = kernelRMSNorm if kernelRMSNorm is not None else MotifRMSNorm
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        # GDLA applies RoPE only to qk_rope_head_dim dimensions
+        rope_head_dim = (
+            getattr(config, "qk_rope_head_dim", None) if getattr(config, "attention_cls", "basic") == "gdla" else None
+        )
+        self.rotary_emb = MotifRotaryEmbedding(config=config, rope_head_dim=rope_head_dim)
+
+        self.mhc_enabled = getattr(config, "mhc_enabled", False)
+        self.mhc_expansion_rate = getattr(config, "mhc_expansion_rate", 4)
+
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_router_logits: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+    ) -> MoeModelOutputWithPast:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        output_router_logits = (
+            output_router_logits
+            if output_router_logits is not None
+            else getattr(self.config, "output_router_logits", False)
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        causal_mask = self._update_causal_mask(
+            attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
+        )
+
+        hidden_states = inputs_embeds
+
+        # Create position embeddings BEFORE MHC expansion (uses (B, S, D) for dtype/device)
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # Expand to (B, S, E, D) for MHC
+        if self.mhc_enabled:
+            hidden_states = hidden_states.unsqueeze(2).expand(-1, -1, self.mhc_expansion_rate, -1).contiguous()
+
+        # Decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        all_router_logits = () if output_router_logits else None
+        next_decoder_cache = None
+
+        for decoder_layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_values,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+            )
+
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+            # Router logits are always the last element
+            if output_router_logits:
+                all_router_logits += (layer_outputs[-1],)
+
+        # Reduce from (B, S, E, D) back to (B, S, D) for MHC
+        if self.mhc_enabled:
+            hidden_states = hidden_states.mean(dim=2)
+
+        hidden_states = self.norm(hidden_states)
+
+        # Add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = next_decoder_cache if use_cache else None
+
+        return MoeModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+            router_logits=all_router_logits,
+        )
+
+    def _update_causal_mask(
+        self,
+        attention_mask: torch.Tensor,
+        input_tensor: torch.Tensor,
+        cache_position: torch.Tensor,
+        past_key_values: Cache,
+        output_attentions: bool,
+    ):
+        if self.config._attn_implementation == "flash_attention_2":
+            if attention_mask is not None and 0.0 in attention_mask:
+                return attention_mask
+            return None
+
+        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+        using_static_cache = isinstance(past_key_values, StaticCache)
+
+        # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
+        if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions:
+            if AttentionMaskConverter._ignore_causal_mask_sdpa(
+                attention_mask,
+                inputs_embeds=input_tensor,
+                past_key_values_length=past_seen_tokens,
+                sliding_window=self.config.sliding_window,
+                is_training=self.training,
+            ):
+                return None
+
+        dtype, device = input_tensor.dtype, input_tensor.device
+        min_dtype = torch.finfo(dtype).min
+        sequence_length = input_tensor.shape[1]
+        # StaticCache
+        if using_static_cache:
+            target_length = past_key_values.get_max_cache_shape()
+        # DynamicCache or no cache
+        else:
+            target_length = (
+                attention_mask.shape[-1]
+                if isinstance(attention_mask, torch.Tensor)
+                else past_seen_tokens + sequence_length
+            )
+
+        # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).
+        causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(
+            attention_mask,
+            sequence_length=sequence_length,
+            target_length=target_length,
+            dtype=dtype,
+            device=device,
+            cache_position=cache_position,
+            batch_size=input_tensor.shape[0],
+            config=self.config,
+            past_key_values=past_key_values,
+        )
+
+        if (
+            self.config._attn_implementation == "sdpa"
+            and attention_mask is not None
+            and attention_mask.device.type == "cuda"
+            and not output_attentions
+        ):
+            causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
+
+        return causal_mask
+
+    @staticmethod
+    def _prepare_4d_causal_attention_mask_with_cache_position(
+        attention_mask: torch.Tensor,
+        sequence_length: int,
+        target_length: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        cache_position: torch.Tensor,
+        batch_size: int,
+        config: MotifConfig,
+        past_key_values: Cache,
+    ):
+        """
+        Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
+        `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.
+
+        Args:
+            attention_mask (`torch.Tensor`):
+                A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`.
+            sequence_length (`int`):
+                The sequence length being processed.
+            target_length (`int`):
+                The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet.
+            dtype (`torch.dtype`):
+                The dtype to use for the 4D attention mask.
+            device (`torch.device`):
+                The device to plcae the 4D attention mask on.
+            cache_position (`torch.Tensor`):
+                Indices depicting the position of the input sequence tokens in the sequence.
+            batch_size (`torch.Tensor`):
+                Batch size.
+            config (`MotifConfig`):
+                The model's configuration class
+            past_key_values (`Cache`):
+                The cache class that is being used currently to generate
+        """
+        if attention_mask is not None and attention_mask.dim() == 4:
+            # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
+            causal_mask = attention_mask
+        else:
+            min_dtype = torch.finfo(dtype).min
+            causal_mask = torch.full(
+                (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device
+            )
+            diagonal_attend_mask = torch.arange(target_length, device=cache_position.device) > cache_position.reshape(
+                -1, 1
+            )
+            if config.sliding_window is not None:
+                if sequence_length > target_length:
+                    sliding_attend_mask = torch.arange(target_length, device=device) <= (
+                        cache_position.reshape(-1, 1) - config.sliding_window
+                    )
+                    diagonal_attend_mask.bitwise_or_(sliding_attend_mask)
+            causal_mask *= diagonal_attend_mask
+            causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
+            if attention_mask is not None:
+                causal_mask = causal_mask.clone()  # copy to contiguous memory for in-place edit
+                if attention_mask.shape[-1] > target_length:
+                    attention_mask = attention_mask[:, :target_length]
+                mask_length = attention_mask.shape[-1]
+                padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
+                padding_mask = padding_mask == 0
+                causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
+                    padding_mask, min_dtype
+                )
+        return causal_mask
+
+
+class MotifForCausalLM(MotifPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+    _tp_plan = {"lm_head": "colwise_gather_output"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = MotifModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+        if config.tie_word_embeddings:
+            self.tie_weights()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_router_logits: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: int = 0,
+        **kwargs,
+    ) -> MoeCausalLMOutputWithPast:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        output_router_logits = (
+            output_router_logits
+            if output_router_logits is not None
+            else getattr(self.config, "output_router_logits", False)
+        )
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            output_router_logits=output_router_logits,
+            cache_position=cache_position,
+        )
+
+        hidden_states = outputs[0]
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        logits = self.lm_head(hidden_states[:, -logits_to_keep:, :])
+        logits = logits.float()
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        return MoeCausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            router_logits=outputs.router_logits,
+        )