Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,7 @@
+---
+language: en
+tags:
+- mlx
+pipeline_tag: text-generation
+library_name: mlx
+---

chat_template.jinja

@@ -0,0 +1,157 @@
+{%- macro visible_text(content) -%}
+    {%- if content is string -%}
+        {{- content -}}
+    {%- elif content is iterable and content is not mapping -%}
+        {%- for item in content -%}
+            {%- if item is mapping and item.type == 'text' -%}
+                {{- item.text -}}
+            {%- elif item is string -%}
+                {{- item -}}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- else -%}
+        {{- content -}}
+    {%- endif -%}
+{%- endmacro -%}
+
+{{- '<|beginoftext|>' -}}
+
+{%- set ns = namespace(has_system=false, last_user_index=-1, last_assistant_index=-1) -%}
+{%- if messages | length > 0 and messages[0].role == 'system' -%}
+    {%- set ns.has_system = true -%}
+{%- endif -%}
+{%- for m in messages -%}
+    {%- if m.role == 'user' -%}
+        {%- set ns.last_user_index = loop.index0 -%}
+    {%- elif m.role == 'assistant' -%}
+        {%- set ns.last_assistant_index = loop.index0 -%}
+    {%- endif -%}
+{%- endfor -%}
+
+{#- ── System / Tools block ── -#}
+{%- if tools is iterable and tools | length > 0 -%}
+    {{- '<|startofturn|><|system|>' -}}
+    {{- '# Tools\n\nYou may call one or more functions to assist with the user query.\n\n' -}}
+    {{- 'You are provided with function signatures within <tools></tools> XML tags:\n\n<tools>' -}}
+    {%- for tool in tools -%}
+        {%- if tool.function is defined -%}
+            {%- set tool = tool.function -%}
+        {%- endif -%}
+        {{- '\n' ~ (tool | tojson) -}}
+    {%- endfor -%}
+    {{- '\n</tools>' -}}
+    {{- '\n\nFor each function call, output in JSON within <tool_call> tags:\n' -}}
+    {%- for tool in tools -%}
+        {%- if tool.function is defined -%}
+            {%- set tool = tool.function -%}
+        {%- endif -%}
+        {%- set _props = tool.parameters.properties if (tool.parameters is defined and tool.parameters.properties is defined) else {} -%}
+        {{- '\n<tool_call>{"name": "' ~ tool.name ~ '", "arguments": {' -}}
+        {%- set _keys = _props | list -%}
+        {%- for k in _keys -%}
+            {{- '"' ~ k ~ '": <' ~ k ~ '>' -}}
+            {%- if not loop.last -%}{{- ', ' -}}{%- endif -%}
+        {%- endfor -%}
+        {{- '}}</tool_call>' -}}
+    {%- endfor -%}
+    {%- if ns.has_system -%}
+        {{- '\n\n' ~ visible_text(messages[0].content) -}}
+    {%- endif -%}
+    {{- '<|endofturn|>' -}}
+{%- elif ns.has_system -%}
+    {{- '<|startofturn|><|system|>' ~ visible_text(messages[0].content) ~ '<|endofturn|>' -}}
+{%- endif -%}
+
+{#- ── Conversation turns ── -#}
+{%- for m in messages -%}
+
+    {#- [Fix 1] continue 대신 if/elif 체인으로 첫 system 스킵 -#}
+    {%- if loop.index0 == 0 and m.role == 'system' -%}
+        {#- already rendered above, skip -#}
+
+    {#- [Fix 2] 중간에 나오는 system 메시지도 처리 -#}
+    {%- elif m.role == 'system' -%}
+        {{- '<|startofturn|><|system|>' ~ visible_text(m.content) ~ '<|endofturn|>' -}}
+
+    {%- elif m.role == 'user' -%}
+        {{- '<|startofturn|><|user|>' -}}
+        {%- if m.references is defined and m.references -%}
+            {{- '<|reference|>' ~ m.references ~ '\n' -}}
+        {%- endif -%}
+        {{- visible_text(m.content) -}}
+        {{- '<|endofturn|>' -}}
+
+    {%- elif m.role == 'assistant' -%}
+        {{- '<|startofturn|><|assistant|>' -}}
+
+        {#- [순서 정책] think → plan → content → tool_calls -#}
+
+        {#- Reasoning block -#}
+        {%- set _content = visible_text(m.content) -%}
+        {%- set _reasoning = '' -%}
+        {%- if m.reasoning_content is string -%}
+            {%- set _reasoning = m.reasoning_content -%}
+        {%- elif '</think>' in _content -%}
+            {%- set _reasoning = _content.split('</think>')[0].split('<think>')[-1].strip() -%}
+            {%- set _content = _content.split('</think>', 1)[-1].lstrip('\n') -%}
+        {%- endif -%}
+        {%- set _has_tools = (tools is defined and tools is iterable and tools | length > 0) -%}
+        {%- set _emit_think = _reasoning and (_has_tools or loop.index0 == ns.last_assistant_index) -%}
+        {%- if _emit_think -%}
+            {{- '<think>' ~ _reasoning.strip() ~ '</think>' -}}
+        {%- endif -%}
+
+        {#- Text content -#}
+        {%- if _content.strip() -%}
+            {{- _content.strip() -}}
+        {%- endif -%}
+
+        {#- Tool calls — IDs are rendered for intermediate assistant turns
+             (context for call↔response correlation) but omitted for the last
+             assistant turn (prediction target — model should not learn to
+             generate IDs). After multi-turn data expansion, intermediate turns
+             become GRAY context and the last turn is GREEN. -#}
+        {%- if m.tool_calls is defined and m.tool_calls -%}
+            {%- set _is_last_assistant = (loop.index0 == ns.last_assistant_index) and not add_generation_prompt -%}
+            {%- for tc in m.tool_calls -%}
+                {%- set _tc_id = tc.id if (tc.id is defined and not _is_last_assistant) else none -%}
+                {%- if tc.function is defined -%}
+                    {%- set tc = tc.function -%}
+                {%- endif -%}
+                {%- set _id_suffix = ', "id": ' ~ (_tc_id | tojson) if _tc_id is not none else '' -%}
+                {%- if tc.arguments is not defined or not tc.arguments or tc.arguments == "" -%}
+                    {{- '\n<tool_call>' ~ '{"name": "' ~ tc.name ~ '", "arguments": ' ~ null ~ _id_suffix ~ '}' ~ '</tool_call>' -}}
+                {%- elif tc.arguments is string -%}
+                    {{- '\n<tool_call>' ~ '{"name": "' ~ tc.name ~ '", "arguments": ' ~ tc.arguments ~ _id_suffix ~ '}' ~ '</tool_call>' -}}
+                {%- else -%}
+                    {{- '\n<tool_call>' ~ '{"name": "' ~ tc.name ~ '", "arguments": ' ~ (tc.arguments | tojson) ~ _id_suffix ~ '}' ~ '</tool_call>' -}}
+                {%- endif -%}
+            {%- endfor -%}
+        {%- endif -%}
+
+        {{- '<|endofturn|>' -}}
+
+    {%- elif m.role == 'tool' -%}
+        {#- [Fix 3] loop.previtem/nextitem으로 인덱스 오버플로 제거 -#}
+        {%- if loop.first or loop.previtem.role != 'tool' -%}
+            {{- '<|startofturn|><|tool|>' -}}
+        {%- endif -%}
+        {{- '<tool_response>' ~ ({"tool_call_id": m.tool_call_id, "content": m.content} | tojson) ~ '</tool_response>' -}}
+        {%- if loop.last or loop.nextitem.role != 'tool' -%}
+            {{- '<|endofturn|>' -}}
+        {%- endif -%}
+
+    {%- endif -%}
+{%- endfor -%}
+
+{#- ── Generation prompt ── -#}
+{%- if add_generation_prompt -%}
+    {{- '<|startofturn|><|assistant|>' -}}
+    {%- if enable_thinking is defined and not enable_thinking -%}
+        {{- '<think></think>' -}}
+    {%- else -%}
+        {{- '<think>' -}}
+    {%- endif -%}
+{%- else -%}
+    {{- '<|endoftext|>' -}}
+{%- endif -%}

configuration_motif.py

@@ -0,0 +1,294 @@
+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"]
+        if callable(getattr(type(self), "validate_rope", None)):
+            self.validate_rope()
+
+        # On ROCm, torch._grouped_mm is not supported at runtime even though
+        # transformers auto-selects the grouped_mm expert backend for torch>=2.9.
+        # Force eager (for-loop) dispatch so MoE models work on ROCm.
+        if (
+            self.num_experts > 0
+            and hasattr(torch.version, "hip")
+            and torch.version.hip is not None
+            and "experts_implementation" not in kwargs
+        ):
+            kwargs["experts_implementation"] = "eager"
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        logger.info(f" kwargs : {kwargs}")

generation_config.json

@@ -0,0 +1,13 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "pad_token_id": 0,
+  "eos_token_id": [0, 3, 6],
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "transformers_version": "5.7.0",
+  "use_cache": true,
+  "do_sample": true,
+  "temperature": 1.0,
+  "top_p": 0.95
+}

model-00001-of-00018.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:27351e5ce2a8f87c3b58f376ab877f72d033f01381116a991cee59291c6a8ab7
+size 5095036671

model-00002-of-00018.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7a3ccc916df9a9c3a95c7322756336b8a14305bc4fa4bcb6c63a99920a556722
+size 5254098018

model-00003-of-00018.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9dec39afaea6b4ec8874f92300de38b6101cb6bcaf245fbb0096b074e6b94439
+size 5254098194

model-00004-of-00018.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b505e32190537c81e02b50fe56d08292c72e852addecc382d1836122cc4c416
+size 5254098351

model-00005-of-00018.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:09d665bb8295e70b61c61d2616bb78921d617cadcd3ceedc9b05009ac5b96a2e
+size 5254098443

model-00006-of-00018.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4789eb3bd8c42eb2f6bdfb69936ec29d661c5a463f848088d905cc5c9aacb04a
+size 5254098445

model-00007-of-00018.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:46e365c83c0b30788b326024f3b3c7db4711d958b9455dcd73b2927ea3b4aa0e
+size 5254098267

model-00008-of-00018.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:60f88393402825e5a37a9c385e19c1e620a1982acd1831791b60b304d6fea193
+size 5254098409

model-00009-of-00018.safetensors

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

model-00010-of-00018.safetensors

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

model-00011-of-00018.safetensors

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

model-00012-of-00018.safetensors

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

model-00013-of-00018.safetensors

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

model-00014-of-00018.safetensors

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

model-00015-of-00018.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7826cb0aab92144dd617037416bb2cdf59873d8eaa1b3749795a0366373b93f9
+size 5254098363

model-00016-of-00018.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:324c92c5c35ca5031a9b659d422cd0c19fc6ce6e50ff11edc9471fbe901d0a9c
+size 5254098397

model-00017-of-00018.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:78bb5a4c8078ea9d1b2c31c72af2027c42a4678c5ee4f925f1273d6ae9aab7a0
+size 5254098285

model-00018-of-00018.safetensors

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

modeling_motif.py

@@ -0,0 +1,1943 @@
+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.integrations import use_experts_implementation
+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):
+        super(PolyNormTorch, self).__init__()
+        self.weight = torch.nn.Parameter(torch.ones(3) / 3)
+        self.bias = torch.nn.Parameter(torch.zeros(1))
+        self.eps = eps
+
+    def _norm(self, x):
+        return x / torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        return (
+            self.weight[0] * self._norm(x**3)
+            + self.weight[1] * self._norm(x**2)
+            + self.weight[2] * self._norm(x)
+            + self.bias
+        )
+
+
+PolyNorm = PolyNormKernel if PolyNormKernel is not None else PolyNormTorch
+
+
+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):
+        super().__init__()
+        self.num_experts = num_experts
+        self.eps = eps
+        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, expert_idx: int) -> torch.Tensor:
+        w = self.weight[expert_idx]  # [3]
+        b = self.bias[expert_idx]  # [1]
+        return w[0] * self._norm(x**3) + w[1] * self._norm(x**2) + w[2] * self._norm(x) + b
+
+
+CUSTOM_ACT2CLS = {"poly_norm": PolyNorm}
+ACT2CLS = {**_ACT2CLS, **CUSTOM_ACT2CLS}
+ACT2FN = ClassInstantier(ACT2CLS)
+
+
+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,
+    ):
+        super().__init__()
+        self.expansion_rate = expansion_rate
+        self.num_dim = num_dim
+        self.sinkhorn_iters = sinkhorn_iters
+
+        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 = 2 * 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}"
+
+
+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
+
+        else:
+            self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+            inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        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)
+
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        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/sin shape: (B, S, D) -> unsqueeze to (B, 1, S, D) for broadcasting with (B, NH, S, D)
+    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)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, hidden_state):
+        hidden_state = self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state)
+        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 MotifAttention(nn.Module):
+    """
+    Grouped Differential Attention module.
+
+    Implements Grouped Differential Attention (https://arxiv.org/pdf/2510.06949)
+    with support for eager, Flash Attention, and SDPA backends via the attention
+    function registry.
+    """
+
+    def __init__(self, config: MotifConfig, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        if layer_idx is None:
+            logger.warning_once(
+                f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
+                "lead to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class."
+            )
+
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads if config.head_dim is None else config.head_dim
+        self.num_key_value_heads = config.num_key_value_heads
+        self.is_causal = True
+        self.attention_dropout = config.attention_dropout
+        self.scaling = 1.0 / math.sqrt(self.head_dim)
+
+        # Grouped Differential Transformer
+        self.num_noise_heads = config.num_noise_heads
+        self.grouped_ratio = (self.num_heads - self.num_noise_heads) // self.num_noise_heads
+        self.q_heads = (self.grouped_ratio + 1) * self.num_noise_heads
+        self.n_signal_heads = self.grouped_ratio * self.num_noise_heads  # = q_heads - noise_heads
+        self.expanded = getattr(config, "expanded", False)
+
+        self.diff_v2 = getattr(config, "diff_v2", False)
+        self.elementwise_attn_output_gate = getattr(config, "elementwise_attn_output_gate", False)
+        self.headwise_attn_output_gate = getattr(config, "headwise_attn_output_gate", False)
+
+        if (self.head_dim * self.num_heads) != self.hidden_size and not self.expanded:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+                f" and `num_heads`: {self.num_heads})."
+            )
+
+        # q_proj size depends on gating mode
+        if self.elementwise_attn_output_gate:
+            self.q_proj = nn.Linear(
+                self.hidden_size, (self.num_heads + self.n_signal_heads * 2) * self.head_dim, bias=False
+            )
+        elif self.headwise_attn_output_gate:
+            self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim + self.n_signal_heads, bias=False)
+        else:
+            self.q_proj = nn.Linear(self.hidden_size, self.q_heads * self.head_dim, bias=False)
+
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+
+        if self.diff_v2:
+            # V2: single V projection, smaller O projection, lambda from projection
+            self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+            self.o_proj = nn.Linear(self.n_signal_heads * self.head_dim, self.hidden_size, bias=False)
+            self.lambda_proj = nn.Linear(self.hidden_size, self.n_signal_heads, bias=False)
+        else:
+            # V1: split V projection, subln, learnable lambda scalars
+            self.k_ratio = config.k_ratio
+            k_noise_heads = self.num_key_value_heads // (self.k_ratio + 1)
+            self.kv_repeat = self.num_noise_heads // k_noise_heads
+            self.v_proj = nn.Linear(self.hidden_size, 2 * k_noise_heads * self.head_dim, bias=False)
+            self.o_proj = nn.Linear(
+                2 * self.grouped_ratio * self.num_noise_heads * self.head_dim, self.hidden_size, bias=False
+            )
+            self.lambda_proj = None
+            for name in ["lambda_q1", "lambda_k1", "lambda_q2", "lambda_k2"]:
+                setattr(self, name, nn.Parameter(torch.zeros(self.head_dim, dtype=torch.float32)))
+                getattr(self, name).data.normal_(mean=0.0, std=0.1)
+            RMSNorm = kernelRMSNorm if kernelRMSNorm is not None else MotifRMSNorm
+            self.subln = RMSNorm(2 * self.head_dim, eps=1e-5)
+            self.lambda_init = 0.8 - 0.6 * math.exp(-0.3 * (layer_idx - 1))
+
+        # Sliding window config for this layer (interleave pattern matching Titan)
+        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)
+            # +1 to match torchtitan convention: torchtitan passes window_size=(W, 0)
+            # directly to flash_attn (right=0 since attention is causal), but HF
+            # converts sliding_window=X to window_size=(X-1, X-1). Adding 1
+            # compensates the left side so X+1 → (X, X). The right side differs
+            # (W vs 0) but is irrelevant because causal masking already prevents
+            # attending to future tokens.
+            effective_window = config.sliding_window + 1
+            if pattern == "all":
+                self.sliding_window = effective_window
+            elif pattern == "interleave" and layer_idx % period != 0:
+                self.sliding_window = effective_window
+            else:
+                self.sliding_window = None
+        else:
+            self.sliding_window = None
+
+    def _reshape_heads(self, tensor, grouped_ratio, num_groups):
+        """2-way head split tensor reshape"""
+        tensor = einops.rearrange(
+            tensor,
+            "... (num_groups group_size) D -> ... num_groups group_size D",
+            num_groups=num_groups,
+            group_size=grouped_ratio + 1,
+        )
+        tensor1 = tensor[..., :grouped_ratio, :]
+        tensor2 = tensor[..., grouped_ratio:, :]
+        return tensor1.contiguous(), tensor2.contiguous()
+
+    def _restore_shape(self, tensor, batch_size, seq_len):
+        """restore tensor"""
+        return tensor.reshape(batch_size, seq_len, -1, self.head_dim)
+
+    def _compute_attention_via_interface(
+        self,
+        attention_interface: Callable,
+        query_states,
+        key_states,
+        value_states,
+        attention_mask,
+        dropout_rate,
+        **kwargs,
+    ):
+        """Compute attention via the registered attention interface."""
+        # Transpose to [batch, num_heads, seq_len, head_dim] for attention interface
+        q = query_states.transpose(1, 2)
+        k = key_states.transpose(1, 2)
+        v = value_states.transpose(1, 2)
+
+        attn_output, _ = attention_interface(
+            self,
+            q,
+            k,
+            v,
+            attention_mask,
+            dropout=dropout_rate,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            is_causal=self.is_causal,
+            **kwargs,
+        )
+
+        # attn_output from interface is [batch, seq_len, num_heads, head_dim] (already transposed)
+        # But some interfaces return [batch, num_heads, seq_len, head_dim] - handle both
+        if attn_output.dim() == 4 and attn_output.shape[1] != query_states.shape[1]:
+            attn_output = attn_output.transpose(1, 2)
+        return attn_output
+
+    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()
+
+        # Lambda for V2 (input-dependent)
+        lambda_full = None
+        if self.diff_v2:
+            lambda_full = self.lambda_proj(hidden_states)  # (bsz, q_len, n_signal_heads)
+
+        # Project Q, K, V
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        # Extract gate score from Q if gating is enabled
+        gate_score = None
+        if self.headwise_attn_output_gate:
+            q_flat = query_states[..., : self.num_heads * self.head_dim]
+            gate_flat = query_states[..., self.num_heads * self.head_dim :]  # (bsz, q_len, n_signal_heads)
+            query_states = q_flat.view(bsz, q_len, self.num_heads, self.head_dim)
+            gate_score = gate_flat.unsqueeze(-1)  # (bsz, q_len, n_signal_heads, 1)
+        elif self.elementwise_attn_output_gate:
+            query_states = query_states.view(bsz, q_len, -1, self.head_dim)
+            q_main = query_states[:, :, : self.num_heads, :]
+            gate_heads = query_states[:, :, self.num_heads :, :]  # (bsz, q_len, n_signal_heads*2, head_dim)
+            gate_score = gate_heads.view(bsz, q_len, self.n_signal_heads, self.head_dim * 2)
+            query_states = q_main
+        else:
+            query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim)
+
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim)
+        value_states = value_states.view(bsz, q_len, -1, self.head_dim)
+
+        # Transpose to [batch, num_heads, seq_len, head_dim] for RoPE
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        # Apply RoPE
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids=position_ids)
+
+        # Handle KV cache
+        if past_key_value is not None:
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        kv_seq_len = key_states.shape[-2]
+        dropout_rate = 0.0 if not self.training else self.attention_dropout
+
+        # Cast dtype if needed (PEFT float32 workaround)
+        input_dtype = query_states.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.q_proj.weight.dtype
+            logger.warning_once(
+                f"The input hidden states seems to be silently casted in float32, this might be related to"
+                f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+                f" {target_dtype}."
+            )
+            query_states = query_states.to(target_dtype)
+            key_states = key_states.to(target_dtype)
+            value_states = value_states.to(target_dtype)
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        if self.diff_v2:
+            attn_output = self._forward_v2(
+                attention_interface,
+                query_states,
+                key_states,
+                value_states,
+                lambda_full,
+                gate_score,
+                attention_mask,
+                dropout_rate,
+                bsz,
+                q_len,
+                **kwargs,
+            )
+        else:
+            attn_output = self._forward_v1(
+                attention_interface,
+                query_states,
+                key_states,
+                value_states,
+                gate_score,
+                attention_mask,
+                dropout_rate,
+                bsz,
+                q_len,
+                kv_seq_len,
+                **kwargs,
+            )
+
+        attn_output = attn_output.reshape(bsz, q_len, -1)
+        attn_output = self.o_proj(attn_output)
+        return attn_output, None, past_key_value
+
+    def _forward_v1(
+        self,
+        attention_interface,
+        query_states,
+        key_states,
+        value_states,
+        gate_score,
+        attention_mask,
+        dropout_rate,
+        bsz,
+        q_len,
+        kv_seq_len,
+        **kwargs,
+    ):
+        """Differential Attention V1: two flash-attn calls with learnable scalar lambda."""
+        # forward() transposes to [B, H, S, D] for RoPE; undo to [B, S, H, D]
+        # so _reshape_heads (splits H) and _compute_attention_via_interface work correctly.
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+        num_groups = self.q_heads // (self.grouped_ratio + 1)
+        q1, q2 = self._reshape_heads(query_states, self.grouped_ratio, num_groups)
+
+        num_kv_groups = self.num_key_value_heads // (self.k_ratio + 1)
+        k1, k2 = self._reshape_heads(key_states, self.k_ratio, num_kv_groups)
+        v1, v2 = self._reshape_heads(value_states, 1, num_kv_groups)
+
+        q1, q2 = self._restore_shape(q1, bsz, q_len), self._restore_shape(q2, bsz, q_len)
+        k1, k2 = self._restore_shape(k1, bsz, kv_seq_len), self._restore_shape(k2, bsz, kv_seq_len)
+        v1, v2 = self._restore_shape(v1, bsz, kv_seq_len), self._restore_shape(v2, bsz, kv_seq_len)
+
+        q_f = torch.cat([q1, q2], dim=2)
+
+        k1 = repeat_kv(k1, 2, self.kv_repeat)
+        k2 = repeat_kv(k2, 2, self.kv_repeat)
+        v1 = repeat_kv(v1, 2, self.kv_repeat)
+        v2 = repeat_kv(v2, 2, self.kv_repeat)
+
+        if self.k_ratio == 1:
+            k_f = torch.cat([repeat_kv(k1, 2, self.grouped_ratio), k2], dim=2)
+        else:
+            k_f = torch.cat([k1, k2], dim=2)
+        v1_f = torch.cat([repeat_kv(v1, 2, self.grouped_ratio), v1], dim=2)
+        v2_f = torch.cat([repeat_kv(v2, 2, self.grouped_ratio), v2], dim=2)
+
+        attn_1 = self._compute_attention_via_interface(
+            attention_interface, q_f, k_f, v1_f, attention_mask, dropout_rate, **kwargs
+        )
+        attn_2 = self._compute_attention_via_interface(
+            attention_interface, q_f, k_f, v2_f, attention_mask, dropout_rate, **kwargs
+        )
+
+        merged_attn = torch.cat([attn_1, attn_2], dim=-1)
+        attn_o = merged_attn[..., :-num_kv_groups, :]
+        attn_n_group = merged_attn[..., -num_kv_groups:, :]
+        attn_n = repeat_kv(attn_n_group, 2, self.grouped_ratio)
+
+        lambda_q1 = self.lambda_q1.unsqueeze(0).expand([bsz, self.lambda_q1.shape[0]])
+        lambda_q2 = self.lambda_q2.unsqueeze(0).expand([bsz, self.lambda_q2.shape[0]])
+        lambda_1 = torch.exp(torch.sum(lambda_q1 * self.lambda_k1, dim=-1).float()).type_as(attn_o)
+        lambda_2 = torch.exp(torch.sum(lambda_q2 * self.lambda_k2, dim=-1).float()).type_as(attn_n)
+        lambda_full = lambda_1 - lambda_2 + self.lambda_init
+
+        attn_output = attn_o - lambda_full.view([bsz, 1, 1, 1]) * attn_n
+        attn_output = self.subln(attn_output)
+        attn_output = attn_output * (1 - self.lambda_init)
+
+        if gate_score is not None:
+            attn_output = attn_output * torch.sigmoid(gate_score)
+
+        expected = (bsz, q_len, self.grouped_ratio * self.num_noise_heads, self.head_dim * 2)
+        if attn_output.size() != expected:
+            raise ValueError(f"`attn_output` should be of size {expected}, but is {attn_output.size()}")
+        return attn_output
+
+    def _forward_v2(
+        self,
+        attention_interface,
+        query_states,
+        key_states,
+        value_states,
+        lambda_full,
+        gate_score,
+        attention_mask,
+        dropout_rate,
+        bsz,
+        q_len,
+        **kwargs,
+    ):
+        """Differential Attention V2: single attention call with input-dependent lambda."""
+        # forward() transposes to [B, H, S, D] for RoPE; undo to [B, S, H, D]
+        # so _compute_attention_via_interface and einops rearrange work correctly.
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+        attn_output = self._compute_attention_via_interface(
+            attention_interface,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout_rate,
+            **kwargs,
+        )
+        # attn_output: (bsz, q_len, n_heads, head_dim)
+
+        # Split heads into signal and noise groups
+        num_groups = self.num_noise_heads  # n_heads // (grouped_ratio + 1)
+        attn_reshaped = einops.rearrange(
+            attn_output,
+            "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.head_dim)
+        attn2_group = attn_reshaped[:, :, :, self.grouped_ratio :, :].reshape(bsz, q_len, num_groups, self.head_dim)
+        attn2 = repeat_kv(attn2_group, 2, self.grouped_ratio)  # (bsz, q_len, n_signal_heads, head_dim)
+
+        # Differential: signal - sigmoid(lambda) * noise
+        attn_output = attn1 - torch.sigmoid(lambda_full).unsqueeze(-1) * attn2
+
+        if gate_score is not None:
+            attn_output = attn_output * torch.sigmoid(gate_score)
+
+        return attn_output
+
+
+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)
+
+        # Softmax scaling with optional mscale correction (DeepSeek-style)
+        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 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."
+            )
+
+        # Sliding window (same interleave pattern as basic Attention)
+        self.sliding_window = None
+        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
+            elif pattern == "interleave" and (layer_idx + 1) % period != 0:
+                self.sliding_window = effective_window
+
+    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()
+
+        # Q path: down-project -> norm -> up-project
+        q_latent = self.q_norm(self.wq_a(hidden_states))  # (bsz, q_len, q_lora_rank)
+        q = self.wq_b(q_latent).view(bsz, q_len, self.num_heads, self.head_dim)
+
+        # 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)
+
+        # Apply RoPE only to the rope parts (qk_rope_head_dim dims)
+        # position_embeddings were computed with qk_rope_head_dim by MotifModel
+        cos, sin = position_embeddings  # (bsz, q_len, qk_rope_head_dim)
+        q_pe = apply_rotary_pos_emb_single(q_pe, cos, sin)
+        # k_pe: (bsz, q_len, qk_rope_head_dim) -> add head dim for apply_rotary_pos_emb_single
+        k_pe = apply_rotary_pos_emb_single(k_pe.unsqueeze(2), cos, sin)  # (bsz, q_len, 1, qk_rope_head_dim)
+
+        # Reconstruct full Q with nope + rope
+        q_total = torch.cat([q_nope, q_pe], dim=-1)  # (bsz, q_len, n_heads, head_dim)
+
+        # 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)
+        # q_t = q_total
+        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)
+
+        # Differential combination: signal - sigmoid(lambda) * noise
+        attn_output = attn1 - torch.sigmoid(lambda_full).unsqueeze(-1) * 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 shape (bs*slen, num_experts)
+        scores = self.gate(x)
+
+        # 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}")
+
+        # top scores shape (bs*slen, experts_top_k)
+        # NOTE: The expert_bias is only used for routing. The gating value
+        #       top_scores is still derived from the original scores.
+
+        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)
+
+
+@use_experts_implementation
+class MotifExperts(nn.Module):
+    """Collection of expert weights stored as fused 3D tensors."""
+
+    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))
+        # Per-expert poly norm (grouped) or shared activation
+        if config.hidden_act == "poly_norm":
+            self.act_fn = GroupedPolyNorm(self.num_experts)
+        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)
+        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 * top_k_weights[token_idx, top_k_pos, None]
+            final_hidden_states.index_add_(0, token_idx, current_hidden.to(final_hidden_states.dtype))
+
+        return final_hidden_states
+
+    def _apply_gate(self, gate_up_output: torch.Tensor, expert_idx: int = 0) -> torch.Tensor:
+        gate, up = gate_up_output.chunk(2, dim=-1)
+        # .chunk() returns non-contiguous views; kernel act_fn requires contiguous input
+        gate = gate.contiguous()
+        if isinstance(self.act_fn, GroupedPolyNorm):
+            return self.act_fn.forward_single(gate, expert_idx) * up
+        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
+
+        # Auxiliary-loss-free load balancing (https://arxiv.org/abs/2408.15664)
+        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)
+
+        # Track expert usage for load balancing
+
+        # Dispatch to experts
+        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)
+
+        # Shared experts
+        if self.shared_experts is not None:
+            final_hidden_states = final_hidden_states + self.shared_experts(x)
+
+        # Return router logits for auxiliary loss computation
+        router_logits = self.router.gate(x.view(-1, dim)) if hasattr(self.router, "gate") else None
+
+        return final_hidden_states.reshape(bs, slen, dim), 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) @ 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 = {
+        "basic": MotifAttention,
+        "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)
+        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
+            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),
+            )
+            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),
+            )
+
+    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,
+        )
+
+        # Expand attention output: (B, S, D) -> (B, S, E, D)
+        attn_expanded = MHCLayer.apply_h_post(attn_out, h_post_attn)
+        # Apply H_res to residual stream
+        x_res_attn = MHCLayer.apply_h_res(x, h_res_attn)
+        h = x_res_attn + attn_expanded
+
+        # === 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)
+
+        # Expand FFN output: (B, S, D) -> (B, S, E, D)
+        ffn_expanded = MHCLayer.apply_h_post(ffn_out, h_post_ffn)
+        h_res_ffn_out = MHCLayer.apply_h_res(h, h_res_ffn)
+        out = h_res_ffn_out + ffn_expanded
+
+        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"))
+        elif isinstance(module, MotifRotaryEmbedding):
+            effective_head_dim = module.inv_freq.shape[0] * 2
+            inv_freq = 1.0 / (
+                self.config.rope_theta
+                ** (torch.arange(0, effective_head_dim, 2, dtype=torch.int64).float() / effective_head_dim)
+            )
+            module.register_buffer("inv_freq", inv_freq, persistent=False)
+
+
+@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)
+        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
+
+        # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in
+        # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail
+        # to infer the attention mask.
+        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
+        ):
+            # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
+            # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
+            # Details: https://github.com/pytorch/pytorch/issues/110213
+            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 we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also
+                # the check is needed to verify is current checkpoint was trained with sliding window or not
+                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,
+        )

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5dc7fe0e557b116bd25eb0185aeaa7672a19368def3e10630a83d48f08447df5
+size 17548406

tokenizer_config.json

@@ -0,0 +1,29 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|beginoftext|>",
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": [
+    "<|system|>",
+    "<|user|>",
+    "<|assistant|>",
+    "<|startofturn|>",
+    "<|endofturn|>",
+    "<|tool|>",
+    "<|reference|>",
+    "<|plan|>",
+    "<|endofplan|>",
+    "<tool_response>",
+    "</tool_response>",
+    "<latent>",
+    "<latent_pad>",
+    "</latent>"
+  ],
+  "is_local": true,
+  "local_files_only": false,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|endoftext|>",
+  "padding_side": "left",
+  "split_special_tokens": false,
+  "tokenizer_class": "TokenizersBackend",
+  "truncation_side": "left"
+}