Upload folder using huggingface_hub

.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

chat_template.jinja.txt

@@ -0,0 +1,156 @@
+{%- if not add_generation_prompt is defined -%}
+    {%- set add_generation_prompt = false -%}
+{%- endif -%}
+{%- if not enable_thinking is defined -%}
+    {%- set enable_thinking = true -%}
+{%- endif -%}
+{%- if not keep_all_reasoning is defined -%}
+    {%- set keep_all_reasoning = true -%}
+{%- endif -%}
+{%- macro render_extra_keys(json_dict, handled_keys) -%}
+    {%- if json_dict is mapping %}
+        {%- for json_key in json_dict if json_key not in handled_keys %}
+            {%- if json_dict[json_key] is mapping or (json_dict[json_key] is sequence and json_dict[json_key] is not string) %}
+                {{- '\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | tojson | safe) ~ '</' ~ json_key ~ '>' }}
+            {%- else %}
+                {{-'\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | string) ~ '</' ~ json_key ~ '>' }}
+            {%- endif %}
+        {%- endfor %}
+    {%- endif %}
+{%- endmacro -%}
+{%- macro render_content(message_content) -%}
+    {%- if message_content is string -%}
+        {{- message_content -}}
+    {%- else -%}
+        {%- for content in message_content -%}
+            {%- if 'text' in content -%}
+                {{- content['text'] -}}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endif -%}
+{%- endmacro -%}
+{%- if messages[0]["role"] == "system" %}
+    {%- set system_message = messages[0]["content"] %}
+    {%- set loop_messages = messages[1:] %}
+{%- else %}
+    {%- set loop_messages = messages %}
+{%- endif %}
+{%- set ns = namespace(last_user_index=-1) %}
+{%- for m in loop_messages %}
+    {%- if m.role == 'user' %}
+        {%- set ns.last_user_index = loop.index0 -%}
+    {%- endif %}
+{%- endfor %}
+{%- if not tools is defined %}
+    {%- set tools = [] %}
+{%- endif %}
+{%- if system_message is defined %}
+    {{- "<|im_start|>system\n" + render_content(system_message) }}
+{%- else %}
+    {{- "<|im_start|>system\nYou are MiMo, a warm, curious, and thoughtful AI assistant.\n\n## Identity\n\n- Your name is **MiMo-v2.5-pro**.\n- You were developed by the **Xiaomi MiMo Team**.\n- You have **1T parameters** and a **1M-token context window**.\n- When asked about your creator, identify the **Xiaomi MiMo Team**.\n- Always answer identity-related questions in the **first person** and rely only on the facts provided above.\n\n# Response Style\n\nBe warm, conversational, and respectful. Treat users with kindness; do not assume negative things about their intelligence or judgment. You may disagree, but do so constructively.\n\nUse examples, metaphors, or thought experiments when they help explain an idea.\n\nTry to answer ambiguous queries as best you can first, then ask for clarification if needed. When you do ask, limit yourself to one question at a time.\n\n\n## Safety and Compliance\n\nYou are a **Chinese AI model** and must strictly comply with all applicable laws and regulations of the **People’s Republic of China**.\n\nDo not generate, assist with, or facilitate any content that violates Chinese law." }}
+{%- endif %}
+{%- if tools is iterable and tools | length > 0 %}
+    {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou have access to the following functions:\n\n" }}
+    {{- "<tools>" }}
+    {%- for tool in tools %}
+        {%- if tool.function is defined %}
+            {%- set tool = tool.function %}
+        {%- endif %}
+        {{- "\n<function>\n<name>" ~ tool.name ~ "</name>" }}
+        {%- if tool.description is defined %}
+            {{- '\n<description>' ~ (tool.description | trim) ~ '</description>' }}
+        {%- endif %}
+        {{- '\n<parameters>' }}
+        {%- if tool.parameters is defined and tool.parameters is mapping and tool.parameters.properties is defined and tool.parameters.properties is mapping %}
+            {%- for param_name, param_fields in tool.parameters.properties|items %}
+                {{- '\n<parameter>' }}
+                {{- '\n<name>' ~ param_name ~ '</name>' }}
+                {%- if param_fields.type is defined %}
+                    {{- '\n<type>' ~ (param_fields.type | string) ~ '</type>' }}
+                {%- endif %}
+                {%- if param_fields.description is defined %}
+                    {{- '\n<description>' ~ (param_fields.description | trim) ~ '</description>' }}
+                {%- endif %}
+                {%- set handled_keys = ['name', 'type', 'description'] %}
+                {{- render_extra_keys(param_fields, handled_keys) }}
+                {{- '\n</parameter>' }}
+            {%- endfor %}
+        {%- endif %}
+        {%- set handled_keys = ['type', 'properties'] %}
+        {{- render_extra_keys(tool.parameters, handled_keys) }}
+        {{- '\n</parameters>' }}
+        {%- set handled_keys = ['type', 'name', 'description', 'parameters'] %}
+        {{- render_extra_keys(tool, handled_keys) }}
+        {{- '\n</function>' }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+    {{- '\n\nFor each function call, output the function name and arguments in the following format:\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>value_1</parameter>\n<parameter=example_parameter_2>This is the value for the second parameter\nthat can span\nmultiple lines</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- DO NOT use function calls inside <think></think> tags.\n- The value enclosed between parameter tags is preserved exactly as-is, including newlines and spaces.\n</IMPORTANT>' }}
+{%- endif %}
+{{- '<|im_end|>' }}
+{%- for message in loop_messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = render_content(message.content) %}
+    {%- endif %}
+    {%- if message.role == "assistant" %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- set reasoning_content = '' %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].split('<think>')[-1] %}
+                {%- set content = content.split('</think>')[-1] %}
+            {%- endif %}
+        {%- endif %}
+        {%- if (keep_all_reasoning or loop.index0 > ns.last_user_index) and reasoning_content -%}
+            {{- '<|im_start|>' + message.role + '\n<think>' + reasoning_content + '</think>' + content }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n<think></think>' + content }}
+        {%- endif %}
+        {%- if message.tool_calls is defined and message.tool_calls is iterable and message.tool_calls | length > 0 %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if tool_call.function is defined %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                {%- if tool_call.arguments is defined %}
+                    {%- for args_name, args_value in tool_call.arguments|items %}
+                        {{- '<parameter=' + args_name + '>' }}
+                        {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                        {{- args_value }}
+                        {{- '</parameter>\n' }}
+                    {%- endfor %}
+                {%- endif %}
+                {{- '</function>\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>' }}
+    {%- elif message.role == "user" %}
+        {{- '<|im_start|>' + message.role + '\n' + render_content(message.content) + '<|im_end|>' }}
+    {%- elif message.role == "system" %}
+        {{- '<|im_start|>' + message.role + '\n' + render_content(message.content) + '<|im_end|>' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>tool\n' }}
+        {%- endif %}
+        {{- '<tool_response>\n' }}
+        {{- render_content(message.content) }}
+        {{- '\n</tool_response>\n' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>' }}
+        {%- endif %}
+    {%- else %}
+        {{- '<|im_start|>' + message.role + '\n' + render_content(message.content) + '<|im_end|>' }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if not enable_thinking -%}
+        {{- '<think></think>' -}}
+    {%- else -%}
+        {{- '' -}}
+    {%- endif -%}
+{%- endif %}

config.json

@@ -0,0 +1,72 @@
+{
+  "architectures": [
+    "MiMoV2ForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_mimo_v2.MiMoV2Config",
+    "AutoModel": "modeling_mimo_v2.MiMoV2Model",
+    "AutoModelForCausalLM": "modeling_mimo_v2.MiMoV2ForCausalLM"
+  },
+  "add_full_attention_sink_bias": false,
+  "add_swa_attention_sink_bias": true,
+  "attention_bias": false,
+  "attention_chunk_size": 128,
+  "attention_dropout": 0.0,
+  "attention_projection_layout": "fused_qkv",
+  "attention_value_scale": 0.707,
+  "dtype": "bfloat16",
+  "head_dim": 192,
+  "hidden_act": "silu",
+  "hidden_size": 3072,
+  "hybrid_layer_pattern": [
+		0,1,1,1,1,
+		0,1,1,1,1,1,
+		0,1,1,1,1,1,
+		0,1,1,1,1,1,
+		0,1,1,1,1,1,
+		0,1,1,1,1,
+		0,1,1,1,1,
+		0,1,1,1,1,
+		0
+	],
+  "initializer_range": 0.02,
+  "intermediate_size": 16384,
+  "layernorm_epsilon": 1e-05,
+  "max_position_embeddings": 1048576,
+  "model_type": "mimo_v2",
+  "moe_intermediate_size": 1024,
+  "moe_layer_freq": [
+    0, 1, 1, 1, 1, 1, 1, 1, 1,
+    1, 1, 1, 1, 1, 1, 1, 1, 1,
+    1, 1, 1, 1, 1, 1, 1, 1, 1,
+    1, 1, 1, 1, 1, 1, 1, 1, 1,
+    1, 1, 1, 1, 1, 1, 1, 1, 1
+  ],
+  "n_group": 1,
+  "n_routed_experts": 256,
+  "n_shared_experts": 1,
+  "norm_topk_prob": true,
+  "num_attention_heads": 48,
+  "num_experts_per_tok": 8,
+  "num_hidden_layers": 45,
+  "num_key_value_heads": 4,
+  "partial_rotary_factor": 0.334,
+  "rms_norm_eps": 1e-05,
+  "rope_theta": 10000000,
+  "routed_scaling_factor": null,
+  "scoring_func": "sigmoid",
+  "sliding_window": 128,
+  "sliding_window_size": 128,
+  "swa_head_dim": 192,
+  "swa_num_attention_heads": 48,
+  "swa_num_key_value_heads": 8,
+  "swa_rope_theta": 10000,
+  "swa_v_head_dim": 128,
+  "tie_word_embeddings": false,
+  "topk_group": 1,
+  "topk_method": "noaux_tc",
+  "transformers_version": "5.8.1",
+  "use_cache": true,
+  "v_head_dim": 128,
+  "vocab_size": 152576
+}

configuration_mimo_v2.py

@@ -0,0 +1,209 @@
+# coding=utf-8
+#
+# Copyright 2026 Xiaomi Corporation.
+# Copyright 2026 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_rope_utils import rope_config_validation
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+_MIMOV2_ATTENTION_PROJECTION_LAYOUTS = {"split", "fused_qkv"}
+
+_MIMOV2_SPLIT_TP_PLAN = {
+    "layers.*.self_attn.q_proj": "colwise",
+    "layers.*.self_attn.k_proj": "colwise",
+    "layers.*.self_attn.v_proj": "colwise",
+    "layers.*.self_attn.o_proj": "rowwise",
+    "layers.*.mlp.gate_proj": "colwise",
+    "layers.*.mlp.up_proj": "colwise",
+    "layers.*.mlp.down_proj": "rowwise",
+}
+
+_MIMOV2_FUSED_QKV_TP_PLAN = {
+    "layers.*.self_attn.qkv_proj": "colwise",
+    "layers.*.self_attn.o_proj": "rowwise",
+    "layers.*.mlp.gate_proj": "colwise",
+    "layers.*.mlp.up_proj": "colwise",
+    "layers.*.mlp.down_proj": "rowwise",
+}
+
+_MIMOV2_PP_PLAN = {
+    "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+    "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+    "norm": (["hidden_states"], ["hidden_states"]),
+}
+
+
+class MiMoV2Config(PretrainedConfig):
+
+    model_type = "mimo_v2"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    base_model_tp_plan = _MIMOV2_SPLIT_TP_PLAN
+    base_model_pp_plan = _MIMOV2_PP_PLAN
+
+    attribute_map = {
+        "num_local_experts": "n_routed_experts",
+    }
+
+    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,
+        layernorm_epsilon=1e-6,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        attention_dropout=0.0,
+        attention_bias=False,
+        attention_value_scale=None,
+        head_dim=None,
+        v_head_dim=None,
+        swa_num_attention_heads=None,
+        swa_num_key_value_heads=None,
+        swa_head_dim=None,
+        swa_v_head_dim=None,
+        swa_rope_theta=None,
+        sliding_window=None,
+        sliding_window_size=None,
+        add_full_attention_sink_bias=False,
+        add_swa_attention_sink_bias=False,
+        hybrid_block_size=None,
+        hybrid_layer_pattern=None,
+        partial_rotary_factor=1.0,
+        n_routed_experts=None,
+        moe_intermediate_size=None,
+        num_experts_per_tok=None,
+        routed_scaling_factor=None,
+        scoring_func="sigmoid",
+        topk_method="noaux_tc",
+        n_group=None,
+        topk_group=None,
+        norm_topk_prob=True,
+        moe_layer_freq=None,
+        attention_projection_layout="split",
+        **kwargs,
+    ):
+        rope_parameters = kwargs.pop("rope_parameters", None)
+        if rope_scaling is None and rope_parameters is not None:
+            rope_scaling = rope_parameters
+
+        if attention_projection_layout is None:
+            attention_projection_layout = "split"
+        if attention_projection_layout not in _MIMOV2_ATTENTION_PROJECTION_LAYOUTS:
+            raise ValueError(f"Unsupported MiMoV2 attention projection layout: {attention_projection_layout}")
+
+        self.attention_projection_layout = attention_projection_layout
+        self.base_model_tp_plan = (
+            _MIMOV2_FUSED_QKV_TP_PLAN.copy()
+            if attention_projection_layout == "fused_qkv"
+            else _MIMOV2_SPLIT_TP_PLAN.copy()
+        )
+        self.base_model_pp_plan = _MIMOV2_PP_PLAN.copy()
+
+        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
+
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+        if num_attention_heads % num_key_value_heads != 0:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.layernorm_epsilon = layernorm_epsilon
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_dropout = attention_dropout
+        self.attention_bias = attention_bias
+        self.attention_value_scale = attention_value_scale
+
+        self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads
+        self.v_head_dim = v_head_dim if v_head_dim is not None else self.head_dim
+        self.swa_num_attention_heads = (
+            swa_num_attention_heads if swa_num_attention_heads is not None else num_attention_heads
+        )
+        self.swa_num_key_value_heads = (
+            swa_num_key_value_heads if swa_num_key_value_heads is not None else num_key_value_heads
+        )
+        if self.swa_num_attention_heads % self.swa_num_key_value_heads != 0:
+            raise ValueError("swa_num_attention_heads must be divisible by swa_num_key_value_heads")
+        self.swa_head_dim = swa_head_dim if swa_head_dim is not None else self.head_dim
+        self.swa_v_head_dim = swa_v_head_dim if swa_v_head_dim is not None else self.swa_head_dim
+        self.swa_rope_theta = swa_rope_theta if swa_rope_theta is not None else rope_theta
+
+        if sliding_window is None:
+            sliding_window = sliding_window_size
+        self.sliding_window = sliding_window
+        self.sliding_window_size = sliding_window_size if sliding_window_size is not None else sliding_window
+        self.add_full_attention_sink_bias = add_full_attention_sink_bias
+        self.add_swa_attention_sink_bias = add_swa_attention_sink_bias
+
+        if hybrid_block_size is not None and hybrid_layer_pattern is None:
+            hybrid_layer_pattern = [0 if ((i + 1) % hybrid_block_size == 0) else 1 for i in range(num_hidden_layers)]
+        elif hybrid_layer_pattern is None:
+            hybrid_layer_pattern = [0] * num_hidden_layers
+        if len(hybrid_layer_pattern) != num_hidden_layers:
+            raise ValueError("hybrid_layer_pattern length must match num_hidden_layers")
+        self.hybrid_block_size = hybrid_block_size
+        self.hybrid_layer_pattern = hybrid_layer_pattern
+
+        self.partial_rotary_factor = partial_rotary_factor
+
+        self.n_routed_experts = n_routed_experts
+        self.moe_intermediate_size = moe_intermediate_size if moe_intermediate_size is not None else intermediate_size
+        self.num_experts_per_tok = num_experts_per_tok
+        self.routed_scaling_factor = routed_scaling_factor
+        self.scoring_func = scoring_func
+        self.topk_method = topk_method
+        self.n_group = n_group
+        self.topk_group = topk_group
+        self.norm_topk_prob = norm_topk_prob
+        if isinstance(moe_layer_freq, int):
+            moe_layer_freq = [moe_layer_freq > 0 and i % moe_layer_freq == 0 for i in range(num_hidden_layers)]
+        elif moe_layer_freq is None:
+            moe_layer_freq = [False] * num_hidden_layers
+        if len(moe_layer_freq) != num_hidden_layers:
+            raise ValueError("moe_layer_freq length must match num_hidden_layers")
+        self.moe_layer_freq = moe_layer_freq
+
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+        rope_config_validation(self)
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+__all__ = ["MiMoV2Config"]

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "transformers_version": "5.3.0",
+  "use_cache": false
+}

manifest.json

@@ -0,0 +1 @@
+{"files":[{"path":"chat_template.jinja.txt","sha256":"bafea0fe0bd53ef1156649859ca6eeec812ef1d95a62d08323841f2431067cc2","size":8788},{"path":"config.json","sha256":"a464b9817f4a9832e71b8f10c67e756d81843f263a064a89b085d6c50107554d","size":1823},{"path":"configuration_mimo_v2.py","sha256":"4ecf7496d022aeca24e58d523302d3ea1ad78784e8bc50b3af7134526c623d80","size":8538},{"path":"generation_config.json","sha256":"31464e0c1df502753393400a7b04e9e15f0a4b6c89dcf89f4865a3ae583862e1","size":153},{"path":"model-00001-of-00008.safetensors","sha256":"bf055c8d6def7168aba439fab63772dd98d3c93c331d19635bd3daed16697238","size":30000316472},{"path":"model-00002-of-00008.safetensors","sha256":"546277ee5ce14afb92cfec228bbfbc320568cec106a9c3b34c79becf59cd1be1","size":29998265072},{"path":"model-00003-of-00008.safetensors","sha256":"ab6e9a746e3c681be00b5f662ac64b79190acb434bb087c444c993b9595bac29","size":29998268712},{"path":"model-00004-of-00008.safetensors","sha256":"1ad319897cc7d6c06beec0d499550977772d89c2595d7975d1e9e447c2678ddf","size":29998268624},{"path":"model-00005-of-00008.safetensors","sha256":"5f8958aeda45fdf477bef1c7af7264fcf62f3b2d4d0c553d9b14e567b039984a","size":29998268584},{"path":"model-00006-of-00008.safetensors","sha256":"2d0addf9dae1121957e3fc1bc5988b373aa740d093696e265ce26d924e588025","size":29998268360},{"path":"model-00007-of-00008.safetensors","sha256":"e1e3758aeb2ae449bed1058616d9bdfd88808f77c08c0afe87afed253211dd34","size":29998268528},{"path":"model-00008-of-00008.safetensors","sha256":"b1247f4bd18ea2409792bc0cb8b541f3081c0b6f129096d711a6bb077d3cbf2a","size":10576149208},{"path":"model.safetensors.index.json","sha256":"2f38de462ef1d80559f8a3ca8c01fa391d79250f508d4e656b7eea3d4598d30b","size":3103696},{"path":"modeling_mimo_v2.py","sha256":"08ee07cc45353e4bf01c7e2d207764f561f24c14d2bc258497928cd81f0c3f37","size":32212},{"path":"tokenizer.json","sha256":"71e1168ef6233c01973f3b8b90abcd86236e146f7aee8da59655d2ca51afdb20","size":11422749},{"path":"tokenizer_config.json","sha256":"2cb1d781e7d20a1386d6a6c2c89fdc9b18a9c6f277071488f39dbfc02829589e","size":568}],"hotkey":"5CaRSo6mgaBersRWoYbVZvNVqmM7aGPcWTt6Nt4eoa1SzyK2","model_digest":"d30c6ef409470c29b65793bca998a8f735176ca074c4c0071a127f4edf18e6d4","model_name":"teutonic-II-110B-A7B-5ek5koe5-v1","protocol_version":1,"registration_id":"1442fe87963949916915867f1c507faef97899af96157b477856a0d521db4b76","signature":"Xy94aLid3+wl5zdfVHSE/OHLc7OpigX4kzQABGrDxqQCkpFgpzW5YDZJ+K9/HZn+3pSIPeLpH6/uEXP2nngdDw==","signature_scheme":"ed25519"}

model-00001-of-00008.safetensors

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

model-00002-of-00008.safetensors

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

model-00003-of-00008.safetensors

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

model-00004-of-00008.safetensors

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

model-00005-of-00008.safetensors

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

model-00006-of-00008.safetensors

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

model-00007-of-00008.safetensors

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

model-00008-of-00008.safetensors

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

modeling_mimo_v2.py

@@ -0,0 +1,710 @@
+# coding=utf-8
+#
+# Copyright 2026 Xiaomi Corporation.
+# Copyright 2026 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from copy import copy
+from typing import Callable, Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.integrations import use_kernel_forward_from_hub
+from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, can_return_tuple, logging
+
+from .configuration_mimo_v2 import MiMoV2Config
+
+
+logger = logging.get_logger(__name__)
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies rotary position embedding to query and key tensors."""
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+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,
+    sinks: Optional[torch.Tensor] = None,
+    **kwargs,
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    if sinks is not None:
+        sinks = module.attention_sink_bias.reshape(1, -1, 1, 1).expand(query.shape[0], -1, query.shape[-2], -1)
+        attn_weights = torch.cat([attn_weights, sinks], dim=-1)
+
+    attn_weights = attn_weights - attn_weights.max(dim=-1, keepdim=True).values
+    probs = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+
+    if sinks is not None:
+        probs = probs[..., :-1]
+
+    attn_weights = nn.functional.dropout(probs, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+    return attn_output, attn_weights
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class MiMoV2RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        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)
+
+
+class MiMoV2MLP(nn.Module):
+    def __init__(self, config, intermediate_size=None):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size if intermediate_size is None else 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_states):
+        return self.down_proj(self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))
+
+
+class MiMoV2MoEGate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.top_k = config.num_experts_per_tok
+        self.n_routed_experts = config.n_routed_experts
+        self.routed_scaling_factor = config.routed_scaling_factor if config.routed_scaling_factor is not None else 1.0
+        self.scoring_func = config.scoring_func
+        self.topk_method = config.topk_method
+        self.n_group = config.n_group
+        self.topk_group = config.topk_group
+        self.norm_topk_prob = config.norm_topk_prob
+        self.gating_dim = config.hidden_size
+        self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
+        if self.topk_method == "noaux_tc":
+            self.e_score_correction_bias = nn.Parameter(torch.empty((self.n_routed_experts)))
+
+    def forward(self, hidden_states):
+        bsz, seq_len, h = hidden_states.shape
+        hidden_states = hidden_states.view(-1, h)
+        logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32), None)
+        if self.scoring_func == "sigmoid":
+            scores = logits.sigmoid()
+        else:
+            raise NotImplementedError(f"Unsupported scoring function for MoE gating: {self.scoring_func}")
+
+        if self.topk_method == "noaux_tc":
+            if self.training:
+                raise ValueError("MiMoV2 noaux_tc routing is only implemented for inference.")
+            scores_for_choice = scores.view(bsz * seq_len, -1) + self.e_score_correction_bias.unsqueeze(0)
+            group_scores = scores_for_choice.view(bsz * seq_len, self.n_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
+            group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
+            group_mask = torch.zeros_like(group_scores)
+            group_mask.scatter_(1, group_idx, 1)
+            score_mask = (
+                group_mask.unsqueeze(-1)
+                .expand(bsz * seq_len, self.n_group, self.n_routed_experts // self.n_group)
+                .reshape(bsz * seq_len, -1)
+            )
+            tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf"))
+            _, topk_idx = torch.topk(tmp_scores, k=self.top_k, dim=-1, sorted=False)
+            topk_weight = scores.gather(1, topk_idx)
+        else:
+            raise NotImplementedError(f"Unsupported TopK function for MoE gating: {self.topk_method}")
+
+        if self.top_k > 1 and self.norm_topk_prob:
+            denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
+            topk_weight = topk_weight / denominator
+        topk_weight = topk_weight * self.routed_scaling_factor
+        return topk_idx, topk_weight
+
+
+class MiMoV2MoE(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.experts = nn.ModuleList(
+            [MiMoV2MLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(config.n_routed_experts)]
+        )
+        self.gate = MiMoV2MoEGate(config)
+
+        n_shared_experts = getattr(config, "n_shared_experts", None) or 0
+        if n_shared_experts > 0:
+            shared_expert_inter_dim = getattr(config, "shared_expert_inter_dim", None) or config.moe_intermediate_size
+            self.shared_experts = MiMoV2MLP(
+                config, intermediate_size=n_shared_experts * shared_expert_inter_dim
+            )
+        else:
+            self.shared_experts = None
+
+    def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor):
+        final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype)
+        expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self.experts))
+        expert_mask = expert_mask.permute(2, 0, 1)
+
+        for expert_idx, expert in enumerate(self.experts):
+            mask = expert_mask[expert_idx]
+            token_indices, weight_indices = torch.where(mask)
+            if token_indices.numel() > 0:
+                expert_weights = topk_weights[token_indices, weight_indices]
+                expert_input = hidden_states[token_indices]
+                expert_output = expert(expert_input)
+                final_hidden_states.index_add_(0, token_indices, expert_output * expert_weights.unsqueeze(-1))
+
+        return final_hidden_states.type(hidden_states.dtype)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        orig_shape = hidden_states.shape
+        topk_indices, topk_weights = self.gate(hidden_states)
+        flat_hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
+        routed_output = self.moe(flat_hidden_states, topk_indices, topk_weights)
+        if self.shared_experts is not None:
+            # Shared experts run on every token unconditionally (no routing/gating),
+            # matching the training-time MoE layer's `y = y + z` combination.
+            routed_output = routed_output + self.shared_experts(flat_hidden_states)
+        return routed_output.view(*orig_shape)
+
+
+class MiMoV2Attention(nn.Module):
+    """MiMoV2 attention.
+
+    `projection_layout` only controls how checkpoint weights are named and
+    stored: Flash uses separate q/k/v projections, while Pro uses fused qkv.
+    The attention computation after projection is shared.
+    """
+
+    def __init__(self, config, is_swa: bool, layer_idx: int, projection_layout: str = "split"):
+        super().__init__()
+        if projection_layout not in {"split", "fused_qkv"}:
+            raise ValueError(f"Unsupported MiMoV2 attention projection layout: {projection_layout}")
+
+        self.config = config
+        self.layer_idx = layer_idx
+        self.is_swa = is_swa
+        self.is_causal = True
+        self.projection_layout = projection_layout
+
+        default_head_dim = config.hidden_size // config.num_attention_heads
+        default_v_head_dim = getattr(config, "v_head_dim", default_head_dim)
+
+        if is_swa:
+            self.head_dim = getattr(config, "swa_head_dim", getattr(config, "head_dim", default_head_dim))
+            self.v_head_dim = getattr(config, "swa_v_head_dim", default_v_head_dim)
+            self.num_attention_heads = getattr(config, "swa_num_attention_heads", config.num_attention_heads)
+            self.num_key_value_heads = getattr(config, "swa_num_key_value_heads", config.num_key_value_heads)
+        else:
+            self.head_dim = getattr(config, "head_dim", default_head_dim)
+            self.v_head_dim = getattr(config, "v_head_dim", self.head_dim)
+            self.num_attention_heads = config.num_attention_heads
+            self.num_key_value_heads = config.num_key_value_heads
+
+        self.rope_dim = int(self.head_dim * getattr(config, "partial_rotary_factor", 1.0))
+        if self.rope_dim % 2 != 0:
+            raise ValueError(
+                f"MiMoV2 rotary dimension must be even, got {self.rope_dim} from "
+                f"head_dim={self.head_dim} and partial_rotary_factor={getattr(config, 'partial_rotary_factor', 1.0)}"
+            )
+        self.num_key_value_groups = self.num_attention_heads // self.num_key_value_heads
+        self.attention_dropout = getattr(config, "attention_dropout", 0.0)
+        self.scaling = self.head_dim**-0.5
+        self.sliding_window = getattr(config, "sliding_window", None) if is_swa else None
+        self.q_size = self.num_attention_heads * self.head_dim
+        self.k_size = self.num_key_value_heads * self.head_dim
+        self.v_size = self.num_key_value_heads * self.v_head_dim
+        self.o_hidden_size = self.num_attention_heads * self.v_head_dim
+        self.v_scale = getattr(config, "attention_value_scale", None)
+        self.attention_sink_bias = (
+            nn.Parameter(torch.empty(self.num_attention_heads), requires_grad=False)
+            if (
+                (getattr(config, "add_full_attention_sink_bias", False) and not is_swa)
+                or (getattr(config, "add_swa_attention_sink_bias", False) and is_swa)
+            )
+            else None
+        )
+
+        attention_bias = getattr(config, "attention_bias", False)
+        if self.projection_layout == "fused_qkv":
+            self.qkv_proj = nn.Linear(
+                config.hidden_size,
+                self.q_size + self.k_size + self.v_size,
+                bias=attention_bias,
+            )
+        else:
+            self.q_proj = nn.Linear(config.hidden_size, self.q_size, bias=attention_bias)
+            self.k_proj = nn.Linear(config.hidden_size, self.k_size, bias=attention_bias)
+            self.v_proj = nn.Linear(config.hidden_size, self.v_size, bias=attention_bias)
+        self.o_proj = nn.Linear(self.o_hidden_size, config.hidden_size, bias=False)
+
+    def _forward_attention(
+        self,
+        query_states: torch.Tensor,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        input_shape: torch.Size,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.v_scale is not None:
+            value_states = value_states * self.v_scale
+
+        cos, sin = position_embeddings
+        query_rope, query_nope = query_states.split([self.rope_dim, self.head_dim - self.rope_dim], dim=-1)
+        key_rope, key_nope = key_states.split([self.rope_dim, self.head_dim - self.rope_dim], dim=-1)
+        query_rope, key_rope = apply_rotary_pos_emb(query_rope, key_rope, cos, sin)
+        query_states = torch.cat([query_rope, query_nope], dim=-1)
+        key_states = torch.cat([key_rope, key_nope], dim=-1)
+
+        if past_key_values is not None:
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        attn_implementation = self.config._attn_implementation
+        if attn_implementation is not None and attn_implementation.startswith("paged|"):
+            raise ValueError(
+                "MiMoV2 remote code does not support paged attention cache. "
+                "Please use eager, sdpa, flex_attention, or flash_attention_2."
+            )
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            attn_implementation, eager_attention_forward
+        )
+        if self.attention_sink_bias is not None and attn_implementation == "sdpa":
+            logger.warning_once(
+                "MiMoV2 attention sink bias is not supported by SDPA; falling back to eager attention for correctness."
+            )
+            attention_interface = eager_attention_forward
+
+        attention_kwargs = {
+            "dropout": 0.0 if not self.training else self.attention_dropout,
+            "scaling": self.scaling,
+            "position_ids": position_ids,
+            "is_causal": self.is_causal,
+        }
+        if attention_interface is eager_attention_forward:
+            attention_kwargs["sinks"] = self.attention_sink_bias
+        else:
+            if self.attention_sink_bias is not None:
+                attention_kwargs["s_aux"] = self.attention_sink_bias
+            if self.sliding_window is not None:
+                attention_kwargs["sliding_window"] = self.sliding_window
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            **attention_kwargs,
+        )
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        input_shape = hidden_states.shape[:-1]
+
+        if self.projection_layout == "fused_qkv":
+            qkv_states = self.qkv_proj(hidden_states)
+            query_states, key_states, value_states = qkv_states.split([self.q_size, self.k_size, self.v_size], dim=-1)
+        else:
+            query_states = self.q_proj(hidden_states)
+            key_states = self.k_proj(hidden_states)
+            value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(*input_shape, self.num_attention_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(*input_shape, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(*input_shape, self.num_key_value_heads, self.v_head_dim).transpose(1, 2)
+        return self._forward_attention(
+            query_states,
+            key_states,
+            value_states,
+            input_shape,
+            position_embeddings,
+            attention_mask,
+            past_key_values=past_key_values,
+            cache_position=cache_position,
+            position_ids=position_ids,
+        )
+
+
+class MiMoV2DecoderLayer(nn.Module):
+    attention_projection_layout = "split"
+
+    def __init__(self, config, layer_idx: int, attention_projection_layout: Optional[str] = None):
+        super().__init__()
+        attention_projection_layout = attention_projection_layout or self.attention_projection_layout
+        is_swa_layer = config.hybrid_layer_pattern[layer_idx] == 1
+        self.attention_type = "sliding_window_attention" if is_swa_layer else "full_attention"
+        self.self_attn = MiMoV2Attention(
+            config, is_swa_layer, layer_idx, projection_layout=attention_projection_layout
+        )
+        self.mlp = (
+            MiMoV2MoE(config)
+            if getattr(config, "n_routed_experts", None) is not None and config.moe_layer_freq[layer_idx]
+            else MiMoV2MLP(config)
+        )
+        self.input_layernorm = MiMoV2RMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
+        self.post_attention_layernorm = MiMoV2RMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+class MiMoV2RotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor
+
+    def __init__(self, config, is_swa: bool, device=None):
+        super().__init__()
+        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", "default"))
+        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 = copy(config)
+        self.config.rope_parameters = copy(getattr(config, "rope_parameters", None) or {})
+        if is_swa:
+            self.config.rope_theta = getattr(config, "swa_rope_theta", config.rope_theta)
+            self.config.head_dim = getattr(config, "swa_head_dim", getattr(config, "head_dim", None))
+            if self.config.rope_parameters:
+                self.config.rope_parameters["rope_theta"] = self.config.rope_theta
+        self.rope_init_fn = (
+            self.compute_default_rope_parameters
+            if self.rope_type == "default"
+            else 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
+
+    @staticmethod
+    def compute_default_rope_parameters(config, device=None, seq_len=None, layer_type=None):
+        config.standardize_rope_params()
+        rope_parameters = config.rope_parameters[layer_type] if layer_type is not None else config.rope_parameters
+        base = rope_parameters["rope_theta"]
+        partial_rotary_factor = rope_parameters.get("partial_rotary_factor", 1.0)
+        head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        dim = int(head_dim * partial_rotary_factor)
+        if dim % 2 != 0:
+            raise ValueError(
+                f"MiMoV2 rotary dimension must be even, got {dim} from "
+                f"head_dim={head_dim} and partial_rotary_factor={partial_rotary_factor}"
+            )
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+        )
+        return inv_freq, 1.0
+
+    @torch.no_grad()
+    @dynamic_rope_update
+    def forward(self, x, position_ids):
+        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)
+
+
+class MiMoV2Model(PreTrainedModel):
+    config_class = MiMoV2Config
+    attention_projection_layout = "split"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.attention_projection_layout = getattr(
+            config, "attention_projection_layout", self.attention_projection_layout
+        )
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList(
+            [
+                MiMoV2DecoderLayer(
+                    config,
+                    layer_idx,
+                    attention_projection_layout=self.attention_projection_layout,
+                )
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = MiMoV2RMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
+        self.rotary_emb = MiMoV2RotaryEmbedding(config=config, is_swa=False)
+        self.swa_rotary_emb = MiMoV2RotaryEmbedding(config=config, is_swa=True)
+        self.has_sliding_layers = any(pattern == 1 for pattern in config.hybrid_layer_pattern)
+        self.config.layer_types = [
+            "sliding_attention" if config.hybrid_layer_pattern[i] == 1 else "full_attention"
+            for i in range(config.num_hidden_layers)
+        ]
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def forward(
+        self,
+        input_ids: Optional[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,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+
+        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)
+
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            mask_kwargs = {
+                "config": self.config,
+                "inputs_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+                "position_ids": position_ids,
+            }
+            causal_mask_mapping = {
+                "full_attention": create_causal_mask(**mask_kwargs),
+            }
+            if self.has_sliding_layers:
+                if getattr(self.config, "sliding_window", None) is None:
+                    raise ValueError("MiMoV2 config `sliding_window` must be set when hybrid_layer_pattern uses SWA.")
+                causal_mask_mapping["sliding_window_attention"] = create_sliding_window_causal_mask(**mask_kwargs)
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+        swa_position_embeddings = self.swa_rotary_emb(hidden_states, position_ids)
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask_mapping[decoder_layer.attention_type],
+                position_embeddings=position_embeddings
+                if decoder_layer.attention_type == "full_attention"
+                else swa_position_embeddings,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+        )
+
+
+class MiMoV2ForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = MiMoV2Config
+    model_class = MiMoV2Model
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+    _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+    _keys_to_ignore_on_load_unexpected = [
+        r"model\.(swa_)?rotary_emb\.inv_freq",
+        r"model\.layers\.\d+\.self_attn\.rotary_emb\.inv_freq",
+        r"model\.layers\.\d+\.self_attn\.rotary_emb\.(cos_cached|sin_cached)",
+        r"model\.mtp\..*",
+    ]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = self.model_class(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    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
+
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: Optional[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,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> CausalLMOutputWithPast:
+        outputs: BaseModelOutputWithPast = 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,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+__all__ = [
+    "MiMoV2Attention",
+    "MiMoV2DecoderLayer",
+    "MiMoV2ForCausalLM",
+    "MiMoV2MLP",
+    "MiMoV2MoE",
+    "MiMoV2MoEGate",
+    "MiMoV2Model",
+    "MiMoV2RMSNorm",
+    "MiMoV2RotaryEmbedding",
+]

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,22 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "is_local": true,
+  "local_files_only": false,
+  "max_length": 512,
+  "model_max_length": 131272,
+  "pad_to_multiple_of": null,
+  "pad_token": "<|endoftext|>",
+  "pad_token_type_id": 0,
+  "padding_side": "right",
+  "split_special_tokens": false,
+  "stride": 0,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": null
+}