Add exact MiniMax-01 GGUF support

docs/api/build_from_gguf.md

@@ -15,7 +15,7 @@ from mobius import build_from_gguf
 
 | Census | Total | Closure |
 |---|---:|---|
-| Architectures | 147 | graph verdicts: {'deferred': 88, 'rejected': 2, 'supported': 57}; importable: 55; quantized import: {'rejected': 11, 'supported': 136}; runtime: {'deferred': 144, 'rejected': 2, 'supported': 1} |
+| Architectures | 147 | graph verdicts: {'deferred': 87, 'rejected': 2, 'supported': 58}; importable: 56; quantized import: {'rejected': 11, 'supported': 136}; runtime: {'deferred': 144, 'rejected': 2, 'supported': 1} |
 | Active stored qtypes | 25 | 24 have an import route; 1 are explicitly deferred with no route |
 | Serialized projector strings | 60 | {'graph-importable': 2, 'runtime-supported': 0} |
 | Tokenizer pre identifiers | 87 | 56 semantic groups; all default to deferred and become materializable only from a validated embedded `tokenizer.huggingface.json` or an exact pinned source in runtime evidence |
@@ -428,7 +428,7 @@ before graph construction or durable output.
 | `mimo2` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | MiMo2 requires fused-QKV dense MTP blocks, attention sinks, interleaved sliding KV cache, and three chained heads selected by offsets. Mobius permits one head and cannot preserve that state or FP8 converter transform. |
 | `minicpm` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | MiniCPM requires architecture-specific embedding, residual, and logit scales, Q/K permutation, optional long/short RoPE tensors, and a conditional dense-or-MoE loader. The existing MiniCPM graph does not prove this complete GGUF contract. |
 | `minicpm3` | — | none (fails before config extraction) | not claimed | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | The pinned MiniCPM3 graph uses MLA Q/KV LoRA projections, separate NoPE/RoPE query and key channels, and embedding, residual, and LM-head scales. The current Mobius MiniCPM graph does not represent that exact topology or its scales. |
-| `minimax-01` | — | none (fails before config extraction) | not claimed | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | The pinned loader schedule is not periodic and its Lightning Attention decay, scaling, residual multipliers, and recurrent rollback semantics are not represented by the current MiniMax graph. |
+| `minimax-01` | — | model=`minimax`; tensor=`minimax` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=supported | Graph import is exact, but released ORT GenAI packaging cannot represent the heterogeneous KV/recurrent state slots or bounded rollback snapshots; runtime packaging remains tracked by #605. |
 | `minimax-m2` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | MiniMax-M2 uses full-vector Q/K norms, partial RoPE, and all-layer correction-biased routed experts under metadata-selected gating. Mobius has no exact graph or suffix-safe expert import for that topology. |
 | `minimax-m3` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | MiniMax-M3 adds F32 sparse-indexer tensors and a second index-key cache with position/cell maps, block masks, rollback, and reorder semantics alongside main K/V state. Mobius has no MSA cache task or sparse-index operators; dense fallback would change the model. |
 | `mistral3` | — | none (fails before config extraction) | exact-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | The pinned Mistral3 loader selects dense or routed-expert text blocks from metadata and applies architecture-specific output temperature scaling. A VLM package additionally requires the deferred Pixtral clip sidecar and exact patch/merge/token contract. The existing Hugging Face Mistral3 graph does not cover that conditional GGUF closure. |
@@ -593,8 +593,11 @@ real-artifact full-logit and stateful-generation parity.
 - `static_cache=True` and non-hybrid task dispatch are rejected for these mixed
   state ABIs.
 
-`minimax-01` remains deferred before config extraction because its pinned
-Lightning schedule and decay/scaling semantics do not match the current graph.
+`minimax-01` graph import supports its exact pinned Lightning schedule,
+decay/scaling semantics, recurrent state, and mixed full-attention cache.
+Runtime packaging remains deferred because the released schema cannot represent
+that heterogeneous state ABI or bounded rollback snapshots; this is tracked by
+[`onnxruntime/mobius#605`](https://github.com/onnxruntime/mobius/issues/605).
 PLaMo2 has a dedicated alternating Mamba1/attention graph and strict GGUF
 tensor closure. Its mixed per-layer recurrent/KV runtime package remains
 deferred to [`onnxruntime/mobius#605`](https://github.com/onnxruntime/mobius/issues/605).

src/mobius/_configs/__init__.py

@@ -47,6 +47,7 @@
     LongcatFlashConfig,
     Mamba2Config,
     MambaConfig,
+    MiniMaxConfig,
     MllamaConfig,
     MMSConfig,
     MoonshineConfig,
@@ -127,6 +128,7 @@
     "Mamba2Config",
     "MambaConfig",
     "MllamaConfig",
+    "MiniMaxConfig",
     "MMSConfig",
     "MoonshineConfig",
     "MuseGlimmerConfig",

src/mobius/_configs/_base.py

@@ -486,6 +486,16 @@ class ArchitectureConfig(BaseModelConfig):
     topk_method: str = "greedy"
     first_k_dense_replace: int = 0
     n_shared_experts: int | None = None
+    disable_qmoe: bool = False
+
+    # MiniMax-01 hybrid attention and normalized-residual scaling.
+    lightning_norm_eps: float | None = None
+    full_attn_alpha_factor: float = 1.0
+    full_attn_beta_factor: float = 1.0
+    linear_attn_alpha_factor: float = 1.0
+    linear_attn_beta_factor: float = 1.0
+    mlp_alpha_factor: float = 1.0
+    mlp_beta_factor: float = 1.0
 
     # Multi-head Latent Attention (MLA) config — DeepSeek-V2/V3
     q_lora_rank: int | None = None
@@ -1435,6 +1445,64 @@ class CausalLMConfig(ArchitectureConfig):
     """
 
 
+@dataclasses.dataclass
+class MiniMaxConfig(CausalLMConfig):
+    """Exact configuration for MiniMax-Text-01 and MiniMax-M1 backbones."""
+
+    @classmethod
+    def from_transformers(cls, config, parent_config=None) -> MiniMaxConfig:
+        base = ArchitectureConfig.from_transformers(config, parent_config)
+        raw_schedule = getattr(config, "attn_type_list", None)
+        if raw_schedule is None:
+            raise ValueError("MiniMax-01 config requires an explicit attn_type_list")
+        if len(raw_schedule) != base.num_hidden_layers:
+            raise ValueError(
+                "MiniMax-01 attn_type_list must contain exactly "
+                f"{base.num_hidden_layers} entries, got {len(raw_schedule)}"
+            )
+        if any(value not in (0, 1, False, True) for value in raw_schedule):
+            raise ValueError("MiniMax-01 attn_type_list entries must be 0 or 1")
+        if not bool(getattr(config, "postnorm", True)):
+            raise ValueError("MiniMax-01 requires postnorm=true")
+        if int(getattr(config, "shared_intermediate_size", 0) or 0):
+            raise ValueError(
+                "MiniMax-01 shared experts are not supported by the pinned GGUF architecture"
+            )
+
+        beta_names = (
+            "layernorm_full_attention_beta",
+            "layernorm_linear_attention_beta",
+            "layernorm_mlp_beta",
+        )
+        betas = {name: float(getattr(config, name, 1.0)) for name in beta_names}
+        if any(not math.isclose(value, 1.0) for value in betas.values()):
+            raise ValueError(f"MiniMax-01 beta residual factors must all equal 1.0: {betas}")
+
+        fields = _shallow_fields(base)
+        fields.update(
+            model_type="minimax",
+            layer_types=[
+                "full_attention" if int(value) == 1 else "lightning_attention"
+                for value in raw_schedule
+            ],
+            hidden_act="silu",
+            norm_topk_prob=True,
+            disable_qmoe=True,
+            lightning_norm_eps=float(getattr(config, "lightning_norm_eps", 1e-6)),
+            full_attn_alpha_factor=float(
+                getattr(config, "layernorm_full_attention_alpha", 1.0)
+            ),
+            full_attn_beta_factor=betas["layernorm_full_attention_beta"],
+            linear_attn_alpha_factor=float(
+                getattr(config, "layernorm_linear_attention_alpha", 1.0)
+            ),
+            linear_attn_beta_factor=betas["layernorm_linear_attention_beta"],
+            mlp_alpha_factor=float(getattr(config, "layernorm_mlp_alpha", 1.0)),
+            mlp_beta_factor=betas["layernorm_mlp_beta"],
+        )
+        return cls(**fields)
+
+
 @dataclasses.dataclass
 class EncoderConfig(ArchitectureConfig):
     """Configuration for encoder-only models (BERT, ViT, etc.)."""

src/mobius/_configs/_base_test.py

@@ -7,7 +7,48 @@
 
 import types
 
-from mobius._configs import ArchitectureConfig, NemotronParseConfig
+import pytest
+
+from mobius._configs import ArchitectureConfig, MiniMaxConfig, NemotronParseConfig
+
+
+def test_minimax_config_extracts_exact_schedule_head_geometry_and_residuals():
+    config = types.SimpleNamespace(
+        model_type="MiniMaxText01",
+        hidden_size=48,
+        intermediate_size=32,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=16,
+        rotary_dim=8,
+        rope_theta=10_000_000.0,
+        vocab_size=64,
+        attn_type_list=[0, 1],
+        num_local_experts=2,
+        num_experts_per_tok=1,
+        rms_norm_eps=1e-5,
+        postnorm=True,
+        shared_intermediate_size=0,
+        layernorm_full_attention_alpha=3.5,
+        layernorm_full_attention_beta=1.0,
+        layernorm_linear_attention_alpha=3.5,
+        layernorm_linear_attention_beta=1.0,
+        layernorm_mlp_alpha=3.5,
+        layernorm_mlp_beta=1.0,
+    )
+
+    extracted = MiniMaxConfig.from_transformers(config)
+
+    assert extracted.model_type == "minimax"
+    assert extracted.head_dim == 16
+    assert extracted.partial_rotary_factor == pytest.approx(0.5)
+    assert extracted.layer_types == ["lightning_attention", "full_attention"]
+    assert extracted.lightning_norm_eps == pytest.approx(1e-6)
+    assert extracted.full_attn_alpha_factor == pytest.approx(3.5)
+    assert extracted.linear_attn_alpha_factor == pytest.approx(3.5)
+    assert extracted.mlp_alpha_factor == pytest.approx(3.5)
+    assert extracted.disable_qmoe
 
 
 def test_nemotron_parse_maps_raw_mbart_decoder_attention_heads():

src/mobius/_registry.py

@@ -32,6 +32,7 @@
     Lfm2Config,
     Lfm2MoeConfig,
     Lfm2VlConfig,
+    MiniMaxConfig,
     MMSConfig,
     MoonshineConfig,
     MuseGlimmerConfig,
@@ -602,7 +603,9 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None:
     "granitemoeshared": ModelRegistration(GraniteMoECausalLMModel),
     "hunyuan_v1_moe": ModelRegistration(HunYuanMoEV1CausalLMModel),
     "jetmoe": ModelRegistration(JetMoeCausalLMModel),
-    "minimax": ModelRegistration(MiniMaxCausalLMModel),
+    "minimax": ModelRegistration(MiniMaxCausalLMModel, config_class=MiniMaxConfig),
+    "MiniMaxText01": ModelRegistration(MiniMaxCausalLMModel, config_class=MiniMaxConfig),
+    "minimax_text_01": ModelRegistration(MiniMaxCausalLMModel, config_class=MiniMaxConfig),
     "mixtral": ModelRegistration(MoECausalLMModel),
     "olmoe": ModelRegistration(MoECausalLMModel),
     "phimoe": ModelRegistration(Phi3MoECausalLMModel),

src/mobius/components/_lightning_attention.py

@@ -59,31 +59,38 @@ class LightningAttention(nn.Module):
         layer_idx: Zero-based layer index, used to compute slope_rate.
     """
 
-    def __init__(self, config: ArchitectureConfig, layer_idx: int):
+    def __init__(
+        self,
+        config: ArchitectureConfig,
+        layer_idx: int,
+        linear_class: type | None = None,
+    ):
         super().__init__()
+        linear_class = linear_class or Linear
         self.num_heads = config.num_attention_heads
-        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.head_dim = config.head_dim
         self.hidden_size = config.hidden_size
         self._dtype = config.dtype
 
         # Fused QKV projection — SiLU applied to all 3*H*d_k before split
-        self.qkv_proj = Linear(
+        self.qkv_proj = linear_class(
             config.hidden_size,
             self.num_heads * self.head_dim * 3,
             bias=False,
         )
-        self.out_proj = Linear(
+        self.o_proj = linear_class(
             self.num_heads * self.head_dim,
             config.hidden_size,
             bias=False,
         )
         # output_gate: sigmoid gate applied to the normalized attention output
-        self.output_gate = Linear(
+        self.output_gate = linear_class(
             config.hidden_size,
             self.num_heads * self.head_dim,
             bias=False,
         )
-        self.norm = RMSNorm(self.num_heads * self.head_dim, eps=config.rms_norm_eps)
+        norm_eps = config.lightning_norm_eps or config.rms_norm_eps
+        self.norm = RMSNorm(self.num_heads * self.head_dim, eps=norm_eps)
 
         # Per-head log-space decay values (negative, so exp < 1)
         # HF: slope_rate[h] = base^(h+1) * factor
@@ -97,6 +104,7 @@ def forward(
         op: OpBuilder,
         hidden_states: ir.Value,
         recurrent_state: ir.Value,
+        attention_mask: ir.Value | None = None,
     ):
         """Lightning Attention forward.
 
@@ -117,14 +125,28 @@ def forward(
         qkv = self.qkv_proj(op, hidden_states)
         qkv = op.Swish(qkv)
 
-        # Split into Q, K, V: each (B, T, num_heads * head_dim)
-        head_total = self.num_heads * self.head_dim
+        # MiniMax stores Q/K/V adjacent within each head:
+        # (B, T, H * 3D) -> (B, T, H, 3D) -> three (B, T, H, D) tensors.
+        qkv = op.Reshape(qkv, [0, 0, self.num_heads, 3 * self.head_dim])
         query, key, value = op.Split(
             qkv,
-            op.Constant(value_ints=[head_total, head_total, head_total]),
+            op.Constant(value_ints=[self.head_dim, self.head_dim, self.head_dim]),
             axis=-1,
             _outputs=3,
         )
+        query = op.Reshape(query, [0, 0, -1])
+        key = op.Reshape(key, [0, 0, -1])
+        value = op.Reshape(value, [0, 0, -1])
+        if attention_mask is not None:
+            # Only newly processed tokens contribute to the recurrent state.
+            current_mask = op.Slice(
+                attention_mask,
+                op.Neg(seq_dim),
+                op.Constant(value_ints=[9223372036854775807]),
+                op.Constant(value_ints=[1]),
+            )
+            current_mask = op.Unsqueeze(op.CastLike(current_mask, value), [2])
+            value = op.Mul(value, current_mask)
 
         # Static decay tensor: (B, T, num_heads) with constant per-head values.
         # Each decay[h] = -slope_rate[h] in log-space → exp(decay[h]) < 1.
@@ -144,8 +166,6 @@ def forward(
         decay = op.Expand(decay_1, expand_to)  # (B, T, num_heads)
 
         # LinearAttention "gated": S_t = exp(g_t) * S_{t-1} + k_t ⊗ v_t
-        # scale = 1/sqrt(head_dim) for proper scaling
-        scale = 1.0 / math.sqrt(self.head_dim)
         attn_out, new_state = op.LinearAttention(
             query,
             key,
@@ -155,7 +175,7 @@ def forward(
             update_rule="gated",
             q_num_heads=self.num_heads,
             kv_num_heads=self.num_heads,
-            scale=scale,
+            scale=1.0,
             _domain=DOMAIN,
             _outputs=2,
         )
@@ -166,16 +186,22 @@ def forward(
         gate = op.Sigmoid(self.output_gate(op, hidden_states))
         attn_out = op.Mul(gate, attn_out)
 
-        output = self.out_proj(op, attn_out)
+        output = self.o_proj(op, attn_out)
         return output, new_state
 
 
 def _compute_decay_log(layer_idx: int, num_layers: int, num_heads: int) -> list[float]:
     """Compute per-head log-space decay values for Lightning Attention.
 
     Returns negative values so that exp(decay[h]) = exp(-slope_rate[h]) < 1.
-    Matches HF ``MiniMaxLightningAttention.get_slope_rate()``.
+    Matches the pinned llama.cpp MiniMax-01 slope calculation.
     """
-    base = 1.0 / (2.0 ** (8.0 / num_heads))
-    factor = 1.0 - layer_idx / (num_layers - 1.0 + 1e-5) + 1e-5
-    return [-(base ** (h + 1)) * factor for h in range(num_heads)]
+    if num_layers <= 1:
+        raise ValueError("MiniMax Lightning Attention requires at least two layers")
+
+    if num_heads <= 0:
+        raise ValueError("MiniMax Lightning Attention requires at least one head")
+    start = 2.0 ** (-(2.0 ** -(math.log2(num_heads) - 3.0)))
+    slopes = [start ** (head + 1) for head in range(num_heads)]
+    factor = 1.0 - layer_idx / (num_layers - 1.0) + 1e-5
+    return [-slope * factor for slope in slopes]

src/mobius/components/_lightning_attention_test.py

@@ -0,0 +1,26 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+import math
+
+import pytest
+
+from mobius.components._lightning_attention import _compute_decay_log
+
+
+def test_minimax_decay_matches_pinned_64_head_formula():
+    actual = _compute_decay_log(layer_idx=7, num_layers=80, num_heads=64)
+    factor = 1.0 - 7.0 / 79.0 + 1e-5
+    expected = [-(2.0 ** (-(head + 1) / 8.0)) * factor for head in range(64)]
+
+    assert actual == pytest.approx(expected)
+
+
+def test_minimax_decay_supports_non_power_of_two_head_counts():
+    actual = _compute_decay_log(layer_idx=0, num_layers=2, num_heads=6)
+    start = 2.0 ** (-(2.0 ** -(math.log2(6) - 3.0)))
+    expected = [-(start ** (head + 1)) * 1.00001 for head in range(6)]
+
+    assert actual == pytest.approx(expected)

src/mobius/components/_moe.py

@@ -429,7 +429,11 @@ def __init__(
         assert config.num_experts_per_tok is not None
         self.num_experts = config.num_local_experts
         self.top_k = config.num_experts_per_tok
-        self._qmoe_quantization = _supported_qmoe_quantization(config.quantization)
+        self._qmoe_quantization = (
+            None
+            if getattr(config, "disable_qmoe", False)
+            else _supported_qmoe_quantization(config.quantization)
+        )
         # Clipped-SwiGLU attributes (QMoE's ``activation_alpha``/``activation_beta``/
         # ``swiglu_limit``). Left ``None`` by default so existing callers get a
         # byte-identical QMoE call (the attributes are simply omitted, even though