Add files using upload-large-folder tool

README.md

@@ -4,7 +4,7 @@ license: mit
 ---
 
 <div align="center">
-Upconverted to BFloat16 by <a href='https://x.com/qubitium'>@Qubitum</a> at ModelCloud
+Upconverted to BFloat16 by <a href='https://x.com/qubitium'>@Qubitum</a> at ModelCloud.
 
 <svg width="60%" height="auto" viewBox="0 0 144 48" fill="none" xmlns="http://www.w3.org/2000/svg">
 <path d="M26.6782 7.96523C26.6782 7.02436 25.913 6.26087 24.9739 6.26087C24.0348 6.26087 23.2695 7.0261 23.2695 7.96523V36.2139C23.2695 38.4 21.4904 40.1791 19.3043 40.1791C17.1183 40.1791 15.3391 38.4 15.3391 36.2139V18.0904C15.3391 17.1496 14.5739 16.3861 13.6348 16.3861C12.6956 16.3861 11.9304 17.1513 11.9304 18.0904V25.7722C11.9304 27.9583 10.1513 29.7374 7.96518 29.7374C5.7791 29.7374 4 27.9583 4 25.7722V22.9878C4 22.3635 4.50609 21.8574 5.13043 21.8574C5.75478 21.8574 6.26087 22.3635 6.26087 22.9878V25.7722C6.26087 26.713 7.02605 27.4765 7.96518 27.4765C8.90431 27.4765 9.66954 26.7113 9.66954 25.7722V18.0904C9.66954 15.9044 11.4487 14.1252 13.6348 14.1252C15.8209 14.1252 17.6 15.9044 17.6 18.0904V36.2139C17.6 37.1548 18.3652 37.9183 19.3043 37.9183C20.2435 37.9183 21.0087 37.153 21.0087 36.2139V25.1322V7.96523C21.0087 5.77914 22.7878 4 24.9739 4C27.16 4 28.9391 5.77914 28.9391 7.96523V31.3565C28.9391 31.9809 28.433 32.487 27.8087 32.487C27.1843 32.487 26.6782 31.9809 26.6782 31.3565V7.96523ZM47.6539 14.1252C45.4678 14.1252 43.6887 15.9044 43.6887 18.0904V33.2296C43.6887 34.1704 42.9235 34.9339 41.9843 34.9339C41.0452 34.9339 40.28 34.1687 40.28 33.2296V7.96523C40.28 5.77914 38.5008 4 36.3148 4C34.1287 4 32.3496 5.77914 32.3496 7.96523V40.0348C32.3496 40.9756 31.5843 41.7391 30.6452 41.7391C29.7061 41.7391 28.9409 40.9739 28.9409 40.0348V36.0643C28.9409 35.44 28.4348 34.9339 27.8104 34.9339C27.1861 34.9339 26.68 35.44 26.68 36.0643V40.0348C26.68 42.2209 28.4591 44 30.6452 44C32.8313 44 34.6104 42.2209 34.6104 40.0348V7.96523C34.6104 7.02436 35.3756 6.26087 36.3148 6.26087C37.2539 6.26087 38.0191 7.0261 38.0191 7.96523V33.2296C38.0191 35.4156 39.7982 37.1948 41.9843 37.1948C44.1704 37.1948 45.9496 35.4156 45.9496 33.2296V18.0904C45.9496 17.1496 46.7148 16.3861 47.6539 16.3861C48.593 16.3861 49.3582 17.1513 49.3582 18.0904V31.3565C49.3582 31.9809 49.8643 32.487 50.4887 32.487C51.113 32.487 51.6191 31.9809 51.6191 31.3565V18.0904C51.6191 15.9044 49.84 14.1252 47.6539 14.1252Z" fill="url(#paint0_linear_17_483)"/>
@@ -186,4 +186,4 @@ Please refer to our [Tool Calling Guide](https://huggingface.co/MiniMaxAI/MiniMa
 
 # Contact Us
 
-Contact us at [model@minimax.io](mailto:model@minimax.io).
+Contact us at [model@minimax.io](mailto:model@minimax.io).

__init__.py

@@ -0,0 +1,26 @@
+# SPDX-FileCopyrightText: 2024-2025 ModelCloud.ai
+# SPDX-FileCopyrightText: 2024-2025 qubitium@modelcloud.ai
+# SPDX-License-Identifier: Apache-2.0
+# Contact: qubitium@modelcloud.ai, x.com/qubitium
+#
+# """MiniMax M2 Hugging Face remote code support."""
+
+from .configuration_minimax_m2 import MiniMaxM2Config
+from .modeling_minimax_m2 import (
+    MiniMaxForCausalLM,
+    MiniMaxM2ForCausalLM,
+    MiniMaxM2Model,
+    MiniMaxM2PreTrainedModel,
+    MiniMaxModel,
+    MiniMaxPreTrainedModel,
+)
+
+__all__ = [
+    "MiniMaxM2Config",
+    "MiniMaxM2PreTrainedModel",
+    "MiniMaxM2Model",
+    "MiniMaxM2ForCausalLM",
+    "MiniMaxPreTrainedModel",
+    "MiniMaxModel",
+    "MiniMaxForCausalLM",
+]

config.json

@@ -79,7 +79,7 @@
   "layernorm_mlp_beta": 1.0,
   "max_position_embeddings": 196608,
   "mlp_intermediate_size": 8192,
-  "model_type": "mixtral",
+  "model_type": "minimax",
   "mtp_transformer_layers": 1,
   "num_attention_heads": 48,
   "num_experts_per_tok": 8,
@@ -105,5 +105,9 @@
   "use_qk_norm": true,
   "use_routing_bias": true,
   "vocab_size": 200064,
-  "torch_dtype": "bfloat16"
+  "torch_dtype": "bfloat16",
+  "auto_map": {
+    "AutoConfig": "configuration_minimax_m2.MiniMaxM2Config",
+    "AutoModelForCausalLM": "modeling_minimax_m2.MiniMaxM2ForCausalLM"
+  }
 }

configuration_minimax_m2.py

@@ -0,0 +1,147 @@
+# SPDX-FileCopyrightText: 2024-2025 ModelCloud.ai
+# SPDX-FileCopyrightText: 2024-2025 qubitium@modelcloud.ai
+# SPDX-License-Identifier: Apache-2.0
+# Contact: qubitium@modelcloud.ai, x.com/qubitium
+
+"""Configuration for the MiniMax M2 architecture."""
+
+from __future__ import annotations
+
+from typing import List, Optional, Union
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class _QuantizationConfigDict(dict):
+    """Ensure quantization config always exposes a `quant_method`."""
+
+    def __init__(self, data: Optional[dict] = None):
+        if data is None:
+            data = {}
+        super().__init__(data)
+        self.setdefault("quant_method", "none")
+
+    def to_dict(self):
+        return dict(self)
+
+
+class MiniMaxM2Config(PretrainedConfig):
+    model_type = "minimax"
+
+    def __init__(
+        self,
+        vocab_size: int = 200_064,
+        hidden_size: int = 3_072,
+        intermediate_size: int = 1_536,
+        mlp_intermediate_size: int = 8_192,
+        num_hidden_layers: int = 62,
+        num_attention_heads: int = 48,
+        num_key_value_heads: int = 8,
+        head_dim: Optional[int] = 128,
+        num_local_experts: int = 256,
+        num_experts_per_tok: int = 8,
+        attn_type_list: Optional[List[int]] = None,
+        attention_dropout: float = 0.0,
+        hidden_act: str = "silu",
+        rms_norm_eps: float = 1e-6,
+        max_position_embeddings: int = 196_608,
+        rope_theta: float = 5_000_000.0,
+        rotary_dim: int = 64,
+        rope_scaling: Optional[dict] = None,
+        use_qk_norm: bool = True,
+        qk_norm_type: str = "per_layer",
+        use_routing_bias: bool = True,
+        scoring_func: str = "sigmoid",
+        router_aux_loss_coef: float = 0.001,
+        router_jitter_noise: float = 0.0,
+        output_router_logits: bool = False,
+        use_grouped_topk: bool = True,
+        num_expert_group: Optional[int] = None,
+        topk_group: Optional[int] = None,
+        routed_scaling_factor: float = 1.0,
+        layernorm_full_attention_beta: float = 1.0,
+        layernorm_linear_attention_beta: float = 1.0,
+        layernorm_mlp_beta: float = 1.0,
+        shared_intermediate_size: int = 0,
+        shared_moe_mode: str = "sigmoid",
+        use_mtp: bool = True,
+        num_mtp_modules: int = 3,
+        mtp_transformer_layers: int = 1,
+        attn_window_size: Optional[Union[int, List[int]]] = None,
+        swa_rope_theta: float = -1.0,
+        sliding_window: Optional[int] = None,
+        initializer_range: float = 0.02,
+        tie_word_embeddings: bool = False,
+        max_model_len: Optional[int] = None,
+        bos_token_id: Optional[int] = None,
+        eos_token_id: Optional[int] = None,
+        pad_token_id: Optional[int] = None,
+        use_cache: bool = True,
+        **kwargs,
+    ) -> None:
+        quantization_config = kwargs.pop("quantization_config", None)
+        if quantization_config is None:
+            quantization_config = _QuantizationConfigDict()
+        elif not isinstance(quantization_config, _QuantizationConfigDict):
+            quantization_config = _QuantizationConfigDict(quantization_config)
+        transformers_version = kwargs.pop("transformers_version", None)
+
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            pad_token_id=pad_token_id,
+            **kwargs,
+        )
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.mlp_intermediate_size = mlp_intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim or hidden_size // num_attention_heads
+        self.num_local_experts = num_local_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.attn_type_list = attn_type_list or [1] * num_hidden_layers
+        self.attention_dropout = attention_dropout
+        self.hidden_act = hidden_act
+        self.rms_norm_eps = rms_norm_eps
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.rotary_dim = rotary_dim
+        self.rope_scaling = rope_scaling
+        self.use_qk_norm = use_qk_norm
+        self.qk_norm_type = qk_norm_type
+        self.use_routing_bias = use_routing_bias
+        self.scoring_func = scoring_func
+        self.router_aux_loss_coef = router_aux_loss_coef
+        self.router_jitter_noise = router_jitter_noise
+        self.output_router_logits = output_router_logits
+        self.use_grouped_topk = use_grouped_topk
+        self.num_expert_group = num_expert_group
+        self.topk_group = topk_group
+        self.routed_scaling_factor = routed_scaling_factor
+        self.layernorm_full_attention_beta = layernorm_full_attention_beta
+        self.layernorm_linear_attention_beta = layernorm_linear_attention_beta
+        self.layernorm_mlp_beta = layernorm_mlp_beta
+        self.shared_intermediate_size = shared_intermediate_size
+        self.shared_moe_mode = shared_moe_mode
+        self.use_mtp = use_mtp
+        self.num_mtp_modules = num_mtp_modules
+        self.mtp_transformer_layers = mtp_transformer_layers
+        self.attn_window_size = attn_window_size
+        self.swa_rope_theta = swa_rope_theta
+        self.sliding_window = sliding_window
+        self.initializer_range = initializer_range
+        self.max_model_len = max_model_len
+        self.use_cache = use_cache
+
+        # Convenient accessor used by rotary embedding helper
+        self.partial_rotary_factor = float(self.rotary_dim) / float(self.head_dim)
+        self.quantization_config = quantization_config
+        self.transformers_version = transformers_version
+
+
+__all__ = ["MiniMaxM2Config"]

modeling_minimax_m2.py

@@ -0,0 +1,765 @@
+# SPDX-FileCopyrightText: 2024-2025 ModelCloud.ai
+# SPDX-FileCopyrightText: 2024-2025 qubitium@modelcloud.ai
+# SPDX-License-Identifier: Apache-2.0
+# Contact: qubitium@modelcloud.ai, x.com/qubitium
+
+"""PyTorch implementation of the MiniMax M2 architecture for Hugging Face Transformers."""
+
+from __future__ import annotations
+
+import copy
+import time
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding, repeat_kv, rotate_half
+
+from .configuration_minimax_m2 import MiniMaxM2Config
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "MiniMaxM2Config"
+_CHECKPOINT_FOR_DOC = "MiniMaxAI/MiniMax-M2"
+
+
+def load_balancing_loss_func(
+    gate_logits: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
+    num_experts: int,
+    top_k: int,
+    attention_mask: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if gate_logits is None:
+        return torch.tensor(0.0)
+    if isinstance(gate_logits, torch.Tensor):
+        logits = gate_logits
+    else:
+        logits = torch.cat([layer_gate.to(gate_logits[0].device) for layer_gate in gate_logits], dim=0)
+
+    routing_weights = torch.softmax(logits, dim=-1, dtype=torch.float32)
+    _, selected = torch.topk(routing_weights, top_k, dim=-1)
+    expert_mask = torch.nn.functional.one_hot(selected, num_experts)
+
+    if attention_mask is None:
+        tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
+        router_prob_per_expert = torch.mean(routing_weights, dim=0)
+    else:
+        batch_size, seq_len = attention_mask.shape
+        num_layers = logits.shape[0] // (batch_size * seq_len)
+
+        expanded_mask = (
+            attention_mask[None, :, :, None, None]
+            .expand(num_layers, batch_size, seq_len, top_k, num_experts)
+            .reshape(-1, top_k, num_experts)
+            .to(logits.device)
+        )
+        tokens_per_expert = torch.sum(expert_mask.float() * expanded_mask, dim=0) / torch.sum(expanded_mask, dim=0)
+
+        router_mask = (
+            attention_mask[None, :, :, None]
+            .expand(num_layers, batch_size, seq_len, num_experts)
+            .reshape(-1, num_experts)
+            .to(logits.device)
+        )
+        router_prob_per_expert = torch.sum(routing_weights * router_mask, dim=0) / torch.sum(router_mask, dim=0)
+
+    loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
+    return loss * num_experts
+
+
+def apply_rotary_pos_emb_partial(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    rotary_dim: int,
+    unsqueeze_dim: int = 2,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    cos = cos.unsqueeze(unsqueeze_dim)[..., :rotary_dim]
+    sin = sin.unsqueeze(unsqueeze_dim)[..., :rotary_dim]
+    q_rot = q[..., :rotary_dim]
+    k_rot = k[..., :rotary_dim]
+
+    q_rot = (q_rot * cos) + (rotate_half(q_rot) * sin)
+    k_rot = (k_rot * cos) + (rotate_half(k_rot) * sin)
+
+    q = torch.cat((q_rot, q[..., rotary_dim:]), dim=-1)
+    k = torch.cat((k_rot, k[..., rotary_dim:]), dim=-1)
+    return q, k
+
+
+class MiniMaxM2RMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        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 MiniMaxM2MLP(nn.Module):
+    def __init__(self, config: MiniMaxM2Config) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+
+        self.w1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.w2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.w3 = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        gate = self.act_fn(self.w1(hidden_states))
+        up = self.w3(hidden_states)
+        hidden_states = gate * up
+        hidden_states = self.w2(hidden_states)
+        return hidden_states
+
+
+class MiniMaxM2SparseMoeBlock(nn.Module):
+    def __init__(self, config: MiniMaxM2Config) -> None:
+        super().__init__()
+        self.hidden_dim = config.hidden_size
+        self.experts = nn.ModuleList([MiniMaxM2MLP(config) for _ in range(config.num_local_experts)])
+        self.num_experts = config.num_local_experts
+        self.top_k = config.num_experts_per_tok
+        self.jitter_noise = config.router_jitter_noise
+        self.use_routing_bias = config.use_routing_bias
+        self.scoring_func = getattr(config, "scoring_func", "softmax")
+        self.use_grouped_topk = getattr(config, "use_grouped_topk", False)
+        self.num_expert_group = getattr(config, "num_expert_group", None)
+        self.topk_group = getattr(config, "topk_group", None)
+        self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
+
+        if self.use_grouped_topk:
+            if self.num_expert_group is None or self.num_expert_group <= 0:
+                self.num_expert_group = 1
+            if self.topk_group is None or self.topk_group <= 0:
+                self.topk_group = min(self.num_expert_group, self.top_k)
+        else:
+            self.num_expert_group = 1
+            self.topk_group = 1
+
+        self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False)
+        if self.use_routing_bias:
+            self.e_score_correction_bias = nn.Parameter(torch.zeros(self.num_experts, dtype=torch.float32))
+        else:
+            self.register_parameter("e_score_correction_bias", None)
+
+    def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch_size, seq_len, hidden_dim = hidden_states.shape
+        if self.training and self.jitter_noise > 0:
+            noise = torch.empty_like(hidden_states).uniform_(
+                1.0 - self.jitter_noise,
+                1.0 + self.jitter_noise,
+            )
+            hidden_states = hidden_states * noise
+
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        gate_dtype = self.gate.weight.dtype
+        router_logits = self.gate(hidden_states.to(gate_dtype)).to(torch.float32)
+        if self.e_score_correction_bias is not None:
+            # Bias is applied after scoring (see vLLM/SGLang implementations).
+            correction_bias = self.e_score_correction_bias.to(router_logits.device, router_logits.dtype)
+        else:
+            correction_bias = None
+
+        if self.scoring_func == "sigmoid":
+            scores = torch.sigmoid(router_logits)
+        elif self.scoring_func == "softmax":
+            scores = torch.softmax(router_logits, dim=-1)
+        else:
+            raise ValueError(f"Unsupported scoring function: {self.scoring_func}")
+
+        if correction_bias is not None:
+            original_scores = scores
+            scores = scores + correction_bias
+        else:
+            original_scores = scores
+        topk_scores: torch.Tensor
+        if self.use_grouped_topk and self.num_expert_group > 1:
+            experts_per_group = scores.size(-1) // self.num_expert_group
+            scores_grouped = scores.view(scores.size(0), self.num_expert_group, experts_per_group)
+            if correction_bias is not None:
+                topk_in_group = min(2, experts_per_group)
+                if topk_in_group > 0:
+                    group_scores = scores_grouped.topk(topk_in_group, dim=-1)[0].sum(dim=-1)
+                else:
+                    group_scores = torch.zeros_like(scores_grouped[..., 0])
+            else:
+                group_scores = scores_grouped.max(dim=-1).values
+            group_mask = torch.zeros_like(group_scores)
+            selected_groups = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=True).indices
+            group_mask.scatter_(1, selected_groups, 1.0)
+            mask = group_mask.unsqueeze(-1).expand(-1, -1, experts_per_group).reshape(scores.size())
+            masked_scores = scores.masked_fill(mask == 0, float("-inf"))
+            topk_scores, selected_experts = torch.topk(masked_scores, self.top_k, dim=-1, sorted=True)
+        else:
+            topk_scores, selected_experts = torch.topk(scores, self.top_k, dim=-1, sorted=True)
+
+        if correction_bias is not None:
+            routing_weights = original_scores.gather(1, selected_experts)
+        else:
+            routing_weights = topk_scores
+
+        routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True).clamp(min=1e-12)
+        if self.routed_scaling_factor != 1.0:
+            routing_weights = routing_weights * self.routed_scaling_factor
+        routing_weights = routing_weights.to(hidden_states.dtype)
+        selected_experts = selected_experts.to(torch.long)
+
+        final_hidden_states = torch.zeros_like(hidden_states)
+        expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
+        expert_hit = torch.nonzero(expert_mask.sum(dim=(-1, -2)) > 0, as_tuple=False).flatten()
+
+        for expert_idx in expert_hit.tolist():
+            expert_layer = self.experts[expert_idx]
+            idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0))
+            token_states = hidden_states.index_select(0, top_x)
+            expert_output = expert_layer(token_states) * routing_weights[top_x, idx].unsqueeze(-1)
+            final_hidden_states.index_add_(0, top_x, expert_output.to(final_hidden_states.dtype))
+
+        final_hidden_states = final_hidden_states.view(batch_size, seq_len, hidden_dim)
+        return final_hidden_states, router_logits
+
+
+class MiniMaxM2Attention(nn.Module):
+    def __init__(self, config: MiniMaxM2Config, layer_idx: int) -> None:
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+
+        self.head_dim = config.head_dim
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // max(1, self.num_key_value_heads)
+        self.rotary_dim = config.rotary_dim
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+
+        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
+        max_model_len = getattr(config, "max_model_len", None)
+        if max_model_len is not None:
+            max_position_embeddings = max(max_position_embeddings, max_model_len)
+
+        attn_window_size = getattr(config, "attn_window_size", None)
+        if isinstance(attn_window_size, list):
+            sliding_window = attn_window_size[layer_idx]
+        else:
+            sliding_window = attn_window_size
+        if sliding_window is not None and sliding_window <= 0:
+            sliding_window = None
+        self.sliding_window = sliding_window
+
+        swa_rope_theta = getattr(config, "swa_rope_theta", -1.0)
+        rope_theta = config.rope_theta
+        if self.sliding_window is not None and swa_rope_theta > 0:
+            rope_theta = swa_rope_theta
+
+        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
+
+        self.use_qk_norm = config.use_qk_norm
+        if self.use_qk_norm:
+            self.q_norm = MiniMaxM2RMSNorm(self.num_heads * self.head_dim, eps=config.rms_norm_eps)
+            self.k_norm = MiniMaxM2RMSNorm(self.num_key_value_heads * self.head_dim, eps=config.rms_norm_eps)
+
+        rope_config = copy.deepcopy(config)
+        rope_config.hidden_size = config.hidden_size
+        rope_config.num_attention_heads = config.num_attention_heads
+        rope_config.partial_rotary_factor = float(config.rotary_dim) / float(self.head_dim)
+        rope_config.rope_theta = rope_theta
+        rope_config.max_position_embeddings = max_position_embeddings
+        self.rotary_emb = LlamaRotaryEmbedding(rope_config)
+
+    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,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        bsz, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        if self.use_qk_norm:
+            q_flat = query_states.transpose(1, 2).reshape(bsz * q_len, -1)
+            k_flat = key_states.transpose(1, 2).reshape(bsz * q_len, -1)
+            q_flat = self.q_norm(q_flat)
+            k_flat = self.k_norm(k_flat)
+            query_states = q_flat.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+            key_states = k_flat.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        if position_embeddings is None:
+            cos, sin = self.rotary_emb(value_states, position_ids)
+        else:
+            cos, sin = position_embeddings
+
+        query_states, key_states = apply_rotary_pos_emb_partial(
+            query_states.transpose(1, 2), key_states.transpose(1, 2), cos, sin, self.rotary_dim
+        )
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+
+        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)
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        attn_weights = torch.matmul(query_states, key_states.transpose(-2, -1)) * self.scaling
+        if attention_mask is not None:
+            attn_weights = attn_weights + attention_mask
+
+        if self.sliding_window is not None and past_key_values is None:
+            query_positions = torch.arange(q_len, device=hidden_states.device).view(1, 1, q_len, 1)
+            key_positions = torch.arange(key_states.shape[-2], device=hidden_states.device).view(1, 1, 1, -1)
+            window_mask = key_positions < (query_positions - self.sliding_window)
+            if window_mask.any():
+                attn_weights = attn_weights.masked_fill(window_mask, float("-inf"))
+
+        attn_weights = torch.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        if self.training and self.attention_dropout > 0:
+            attn_weights = F.dropout(attn_weights, p=self.attention_dropout)
+
+        attn_output = torch.matmul(attn_weights, value_states)
+        attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, -1)
+        attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+        return attn_output, attn_weights
+
+
+class MiniMaxM2LogitsProcessor(nn.Module):
+    def __init__(self, config: MiniMaxM2Config) -> None:
+        super().__init__()
+        self.scale = getattr(config, "logits_scale", 1.0)
+
+    def forward(self, lm_head: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor:
+        logits = lm_head(hidden_states)
+        if self.scale != 1.0:
+            logits = logits * self.scale
+        return logits
+
+
+class MiniMaxM2DecoderLayer(nn.Module):
+    def __init__(self, config: MiniMaxM2Config, layer_idx: int) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = MiniMaxM2Attention(config, layer_idx)
+        self.block_sparse_moe = MiniMaxM2SparseMoeBlock(config)
+        self.input_layernorm = MiniMaxM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = MiniMaxM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    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,
+        output_attentions: bool = False,
+        residual: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], torch.Tensor]:
+        residual_input = hidden_states if residual is None else residual
+        hidden_states = self.input_layernorm(hidden_states)
+
+        attn_output, attn_weights = 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,
+            output_attentions=output_attentions,
+        )
+        hidden_states = residual_input + attn_output
+
+        residual_post_attn = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        moe_output, router_logits = self.block_sparse_moe(hidden_states)
+        hidden_states = residual_post_attn + moe_output
+
+        return hidden_states, hidden_states, router_logits, attn_weights
+
+
+class MiniMaxM2PreTrainedModel(PreTrainedModel):
+    config_class = MiniMaxM2Config
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MiniMaxM2DecoderLayer"]
+    _supports_flash_attn = False
+    _supports_sdpa = False
+    _supports_attention_backend = False
+
+    def _init_weights(self, module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+    def _remap_qkv_weights(self, state_dict):
+        num_q = self.config.num_attention_heads * self.config.head_dim
+        num_kv = self.config.num_key_value_heads * self.config.head_dim
+
+        for layer_idx in range(self.config.num_hidden_layers):
+            prefix = f"model.layers.{layer_idx}.self_attn"
+            weight_key = f"{prefix}.qkv_proj.weight"
+            if weight_key in state_dict:
+                qkv_weight = state_dict.pop(weight_key)
+                q_weight, k_weight, v_weight = qkv_weight.split([num_q, num_kv, num_kv], dim=0)
+                state_dict.setdefault(f"{prefix}.q_proj.weight", q_weight)
+                state_dict.setdefault(f"{prefix}.k_proj.weight", k_weight)
+                state_dict.setdefault(f"{prefix}.v_proj.weight", v_weight)
+
+    def load_state_dict(self, state_dict, strict: bool = True):
+        if not isinstance(state_dict, dict):
+            raise TypeError(f"Expected state_dict to be dict, got {type(state_dict)}")
+
+        filtered_state_dict = {}
+        drop_suffixes = ("weight_scale_inv", "weight_scale", "input_scale", "scales", "amax")
+        for key, value in state_dict.items():
+            if key.endswith(drop_suffixes) or "fp8" in key:
+                continue
+            filtered_state_dict[key] = value
+
+        self._remap_qkv_weights(filtered_state_dict)
+
+        if logger.isEnabledFor(logging.INFO):
+            logger.info(
+                "MiniMaxM2: loading %d tensors (filtered from %d original).",
+                len(filtered_state_dict),
+                len(state_dict),
+            )
+
+        load_start = time.perf_counter()
+        result = super().load_state_dict(filtered_state_dict, strict=strict)
+        load_elapsed = time.perf_counter() - load_start
+        if logger.isEnabledFor(logging.INFO):
+            logger.info("MiniMaxM2: state_dict load finished in %.2f seconds.", load_elapsed)
+
+        return result
+
+
+class MiniMaxM2Model(MiniMaxM2PreTrainedModel):
+    def __init__(self, config: MiniMaxM2Config) -> None:
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [MiniMaxM2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = MiniMaxM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.gradient_checkpointing = False
+
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        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.Tensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+        output_router_logits: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[MoeModelOutputWithPast, Tuple]:
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds.")
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        output_router_logits = (
+            output_router_logits if output_router_logits is not None else self.config.output_router_logits
+        )
+
+        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 self.config.sliding_window is not None:
+            causal_mask = create_sliding_window_causal_mask(
+                config=self.config,
+                input_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                cache_position=cache_position,
+                past_key_values=past_key_values,
+                position_ids=position_ids,
+            )
+        else:
+            causal_mask = create_causal_mask(
+                config=self.config,
+                input_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                cache_position=cache_position,
+                past_key_values=past_key_values,
+                position_ids=position_ids,
+            )
+
+        hidden_states = inputs_embeds
+
+        all_hidden_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+        all_router_logits = () if output_router_logits else None
+
+        residual = None
+        for decoder_layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=None,
+                output_attentions=output_attentions,
+                residual=residual,
+            )
+
+            hidden_states, residual, router_logits, attn_weights = layer_outputs
+
+            if output_router_logits:
+                all_router_logits = all_router_logits + (router_logits,)
+            if output_attentions:
+                all_attentions = all_attentions + (attn_weights,)
+
+        hidden_states = self.norm(hidden_states)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            outputs = (hidden_states, past_key_values)
+            if output_hidden_states:
+                outputs += (all_hidden_states,)
+            if output_attentions:
+                outputs += (all_attentions,)
+            if output_router_logits:
+                outputs += (all_router_logits,)
+            return outputs
+
+        return MoeModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+            hidden_states=all_hidden_states,
+            attentions=all_attentions,
+            router_logits=all_router_logits,
+        )
+
+
+class MiniMaxM2ForCausalLM(MiniMaxM2PreTrainedModel, GenerationMixin):
+    def __init__(self, config: MiniMaxM2Config) -> None:
+        super().__init__(config)
+        self.model = MiniMaxM2Model(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.router_aux_loss_coef = config.router_aux_loss_coef
+        self.num_experts = config.num_local_experts
+        self.num_experts_per_tok = config.num_experts_per_tok
+        self.logits_processor = MiniMaxM2LogitsProcessor(config)
+
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
+        self.lm_head = new_embeddings
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[Cache] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        **kwargs,
+    ):
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:]
+            if attention_mask is not None:
+                attention_mask = attention_mask[:, -past_key_values.get_seq_length() - 1 :]
+
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "past_key_values": past_key_values,
+            "inputs_embeds": inputs_embeds,
+        }
+
+    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.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+        output_router_logits: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+    ) -> Union[MoeCausalLMOutputWithPast, Tuple]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        output_router_logits = (
+            output_router_logits if output_router_logits is not None else self.config.output_router_logits
+        )
+
+        model_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,
+            cache_position=cache_position,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            output_router_logits=output_router_logits,
+            return_dict=True,
+        )
+
+        hidden_states = model_outputs.last_hidden_state
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) and logits_to_keep > 0 else slice(None)
+        logits = self.logits_processor(self.lm_head, hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(shift_logits.view(-1, self.vocab_size), shift_labels.view(-1))
+
+        aux_loss = None
+        if output_router_logits and model_outputs.router_logits is not None:
+            aux_loss = load_balancing_loss_func(
+                model_outputs.router_logits,
+                num_experts=self.num_experts,
+                top_k=self.num_experts_per_tok,
+                attention_mask=attention_mask,
+            )
+            if loss is not None:
+                loss = loss + self.router_aux_loss_coef * aux_loss.to(loss.device)
+
+        if not return_dict:
+            output = (logits,) + (model_outputs.past_key_values,)
+            if output_hidden_states:
+                output += (model_outputs.hidden_states,)
+            if output_attentions:
+                output += (model_outputs.attentions,)
+            if output_router_logits:
+                output += (model_outputs.router_logits,)
+            return ((loss,) + output) if loss is not None else output
+
+        return MoeCausalLMOutputWithPast(
+            loss=loss,
+            aux_loss=aux_loss,
+            logits=logits,
+            past_key_values=model_outputs.past_key_values,
+            hidden_states=model_outputs.hidden_states,
+            attentions=model_outputs.attentions,
+            router_logits=model_outputs.router_logits,
+        )
+
+# -----------------------------------------------------------------------------
+# Backward compatibility aliases
+# -----------------------------------------------------------------------------
+
+MiniMaxRMSNorm = MiniMaxM2RMSNorm
+MiniMaxSparseMoeBlock = MiniMaxM2SparseMoeBlock
+MiniMaxAttention = MiniMaxM2Attention
+MiniMaxDecoderLayer = MiniMaxM2DecoderLayer
+MiniMaxMLP = MiniMaxM2MLP
+MiniMaxPreTrainedModel = MiniMaxM2PreTrainedModel
+MiniMaxModel = MiniMaxM2Model
+
+
+class MiniMaxForCausalLM(MiniMaxM2ForCausalLM):
+    """Alias for compatibility with checkpoints exporting MiniMaxForCausalLM."""
+
+
+__all__ = [
+    "MiniMaxM2RMSNorm",
+    "MiniMaxM2SparseMoeBlock",
+    "MiniMaxM2Attention",
+    "MiniMaxM2DecoderLayer",
+    "MiniMaxM2Model",
+    "MiniMaxM2ForCausalLM",
+    "MiniMaxM2PreTrainedModel",
+    "MiniMaxRMSNorm",
+    "MiniMaxSparseMoeBlock",
+    "MiniMaxAttention",
+    "MiniMaxDecoderLayer",
+    "MiniMaxPreTrainedModel",
+    "MiniMaxModel",
+    "MiniMaxMLP",
+    "MiniMaxForCausalLM",
+]

test_minimax_m2_hf.py

@@ -0,0 +1,178 @@
+# SPDX-FileCopyrightText: 2024-2025 ModelCloud.ai
+# SPDX-FileCopyrightText: 2024-2025 qubitium@modelcloud.ai
+# SPDX-License-Identifier: Apache-2.0
+# Contact: qubitium@modelcloud.ai, x.com/qubitium
+
+"""
+MiniMax-M2 Hugging Face checkpoint sanity check with streaming output.
+
+Usage:
+    python test_minimax_m2_hf.py \
+        --model-path /monster/data/model/MiniMax-M2-bf16 \
+        --question "How many letter A are there in the word Alphabet? Reply with the number only."
+"""
+
+from __future__ import annotations
+
+import argparse
+import threading
+from pathlib import Path
+
+import torch.nn as nn
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+
+from gptqmodel.hf_minimax_m2.modeling_minimax_m2 import (
+    MiniMaxAttention,
+    MiniMaxDecoderLayer,
+    MiniMaxForCausalLM,
+    MiniMaxMLP,
+    MiniMaxM2Attention,
+    MiniMaxM2DecoderLayer,
+    MiniMaxM2ForCausalLM,
+    MiniMaxM2MLP,
+    MiniMaxM2RMSNorm,
+    MiniMaxM2SparseMoeBlock,
+    MiniMaxRMSNorm,
+    MiniMaxSparseMoeBlock,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="MiniMax-M2 HF checkpoint smoke test.")
+    parser.add_argument(
+        "--model-path",
+        type=str,
+        default="/monster/data/model/MiniMax-M2-bf16",
+        help="Path to the MiniMax-M2 Hugging Face checkpoint directory.",
+    )
+    parser.add_argument(
+        "--question",
+        type=str,
+        default="How many letter A are there in the word Alphabet? Reply with the number only.",
+        help="User question to send through the chat template.",
+    )
+    parser.add_argument(
+        "--max-new-tokens",
+        type=int,
+        default=512,
+        help="Maximum number of new tokens to sample from the model.",
+    )
+    return parser.parse_args()
+
+
+def build_prompt(tokenizer: AutoTokenizer, question: str) -> str:
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": question},
+    ]
+    return tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+
+
+def assert_module_types(model: MiniMaxM2ForCausalLM) -> None:
+    causal_lm_types = (MiniMaxM2ForCausalLM, MiniMaxForCausalLM)
+    decoder_layer_types = (MiniMaxM2DecoderLayer, MiniMaxDecoderLayer)
+    attention_types = (MiniMaxM2Attention, MiniMaxAttention)
+    moe_block_types = (MiniMaxM2SparseMoeBlock, MiniMaxSparseMoeBlock)
+    norm_types = (MiniMaxM2RMSNorm, MiniMaxRMSNorm)
+    mlp_types = (MiniMaxM2MLP, MiniMaxMLP)
+
+    assert isinstance(
+        model, causal_lm_types
+    ), f"Expected MiniMaxM2ForCausalLM/MiniMaxForCausalLM, received {type(model).__name__}"
+
+    decoder = getattr(model, "model", None)
+    assert decoder is not None, "Model is missing the `model` attribute with decoder layers."
+
+    for layer_idx, layer in enumerate(decoder.layers):
+        assert isinstance(
+            layer, decoder_layer_types
+        ), f"Layer {layer_idx}: expected MiniMax(M2)DecoderLayer, got {type(layer).__name__}"
+        assert isinstance(
+            layer.self_attn, attention_types
+        ), f"Layer {layer_idx}: unexpected self_attn type {type(layer.self_attn).__name__}"
+        assert isinstance(
+            layer.block_sparse_moe, moe_block_types
+        ), f"Layer {layer_idx}: unexpected MoE block type {type(layer.block_sparse_moe).__name__}"
+        assert isinstance(
+            layer.input_layernorm, norm_types
+        ), f"Layer {layer_idx}: unexpected input_layernorm type {type(layer.input_layernorm).__name__}"
+        assert isinstance(
+            layer.post_attention_layernorm, norm_types
+        ), f"Layer {layer_idx}: unexpected post_attention_layernorm type {type(layer.post_attention_layernorm).__name__}"
+
+        moe_block = layer.block_sparse_moe
+        assert isinstance(
+            moe_block.experts, nn.ModuleList
+        ), f"Layer {layer_idx}: expected experts to be a ModuleList, got {type(moe_block.experts).__name__}"
+        for expert_idx, expert in enumerate(moe_block.experts):
+            assert isinstance(
+                expert, mlp_types
+            ), f"Layer {layer_idx} expert {expert_idx}: expected MiniMax(M2)MLP, got {type(expert).__name__}"
+
+
+def main() -> None:
+    args = parse_args()
+    model_path = Path(args.model_path).expanduser().resolve()
+
+    print(f"Loading tokenizer from {model_path}...")
+    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+
+    print(f"Loading model from {model_path}...")
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        dtype="bfloat16",
+        device_map="auto",
+        trust_remote_code=True,
+    )
+
+    # Uncomment to enforce module type checks.
+    # print("Validating module types...")
+    # assert_module_types(model)
+
+    prompt = build_prompt(tokenizer, args.question)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    print("Running generation (streaming)...\n")
+    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=False)
+    eos_ids = model.generation_config.eos_token_id
+    if eos_ids is None:
+        eos_ids = []
+    elif isinstance(eos_ids, int):
+        eos_ids = [eos_ids]
+    think_end_id = tokenizer.convert_tokens_to_ids("</think>")
+    if think_end_id is not None and think_end_id not in eos_ids:
+        eos_ids = eos_ids + [think_end_id]
+
+    generation_kwargs = dict(
+        **inputs,
+        max_new_tokens=args.max_new_tokens,
+        streamer=streamer,
+        eos_token_id=eos_ids if eos_ids else None,
+    )
+
+    generation_thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
+    generation_thread.start()
+
+    completion = []
+    first_chunk = True
+    seen_end_reasoning = False
+    for text in streamer:
+        if first_chunk:
+            print("<think>", end="", flush=True)
+            completion.append("<think>")
+            first_chunk = False
+        print(text, end="", flush=True)
+        completion.append(text)
+        if "</think>" in text:
+            seen_end_reasoning = True
+
+    generation_thread.join()
+    print("\n\n=== Completed Response ===")
+    final_text = "".join(completion).strip()
+    print(final_text or "<empty response>")
+    if not seen_end_reasoning:
+        print("\n[warning] No </think> token detected in streamed output.", flush=True)
+
+
+if __name__ == "__main__":
+    main()