Add dedicated LFM2MoE GGUF graph 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': 91, 'rejected': 3, 'supported': 53}; importable: 51; quantized import: {'rejected': 11, 'supported': 136}; runtime: {'deferred': 144, 'rejected': 3} |
+| Architectures | 147 | graph verdicts: {'deferred': 90, 'rejected': 3, 'supported': 54}; importable: 52; quantized import: {'rejected': 12, 'supported': 135}; runtime: {'deferred': 144, 'rejected': 3} |
 | 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 exact-copy only with a validated embedded `tokenizer.huggingface.json` |
@@ -418,7 +418,7 @@ before graph construction or durable output.
 | `kimi-linear` | — | none (fails before config extraction) | not claimed | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | Kimi-Linear alternates KDA recurrent and NoPE MLA layers, carrying three rolling convolution histories and a per-head matrix state in addition to attention cache. Its two-stage decay/output gates and sigmoid correction-bias MoE routing are not represented by any Mobius graph or state task; aliasing it to Kimi-K3, Mamba, or ordinary attention would change the model. |
 | `laguna` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | Laguna combines per-head-or-element softplus attention gates, dual-RoPE interleaved sliding-window attention, a dense prefix, and sigmoid correction-biased routed/shared experts. Mobius has no exact graph or iSWA cache contract. |
 | `lfm2` | — | model=`lfm2`; tensor=`lfm2` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=supported | Config extraction, exact pinned tensor-name closure, GGUF value transforms, and synthetic recurrent-state execution are covered, but no representative real-weight GGUF has yet passed independent full-logit parity and deterministic multi-token stateful ORT generation. Runtime packaging remains deferred until that evidence exists. |
-| `lfm2moe` | — | none (fails before config extraction) | not claimed | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | LFM2MoE has an arbitrary per-layer attention/short-convolution schedule plus a dense-to-sigmoid-routed-MoE transition. Recurrent layers use F32 rolling convolution state with copy-on-write sequence reorder and bounded rollback snapshots, while attention layers use KV cache. The dense LFM2 graph and ordinary KV/static-cache tasks do not own that mixed state or the expert correction-bias semantics. |
+| `lfm2moe` | — | model=`lfm2_moe`; tensor=`lfm2`+`lfm2_moe_extras` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=rejected | Config extraction, exact pinned tensor-name closure, GGUF value transforms, and synthetic recurrent-state execution are covered, but no representative real-weight GGUF has yet passed independent full-logit parity and deterministic multi-token stateful ORT generation. Runtime packaging remains deferred until that evidence exists. The mobius graph uses floating Linear modules for this architecture, so no MatMulNBits or BlockQuantizedMatMul target can consume preserved GGUF projection weights. Use keep_quantized=False for explicit float import. |
 | `llada` | — | model=`llada`; tensor=`llama` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=supported | Config extraction, suffix-exact tensor closure, masked-diffusion task dispatch, and synthetic full-sequence execution are covered, but no pinned real GGUF has passed independent Hugging Face/llama.cpp masked-step logit parity and deterministic multi-step generation parity. Runtime packaging remains deferred until both exist. |
 | `llada-moe` | — | model=`llada`; module=`llada_moe`; tensor=`llama`+`diffusion_fused_qkv`+`moe_qk_norm_extras`+`moe_extras` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=supported | Config extraction, suffix-exact tensor closure, masked-diffusion task dispatch, and synthetic full-sequence execution are covered, but no pinned real GGUF has passed independent Hugging Face/llama.cpp masked-step logit parity and deterministic multi-step generation parity. Runtime packaging remains deferred until both exist. |
 | `llama` | `mistral` | model=`llama`; tensor=`llama` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=supported | Config extraction, exact tensor-name closure, and a full synthetic GGUF graph build are covered, but no representative real-weight GGUF has yet passed ORT parity or generation validation. Runtime packaging remains deferred until that evidence exists. |

src/mobius/_configs/__init__.py

@@ -41,6 +41,7 @@
     JambaConfig,
     JetMoeConfig,
     Lfm2Config,
+    Lfm2MoeConfig,
     Lfm2VlConfig,
     LongcatFlashConfig,
     Mamba2Config,
@@ -117,6 +118,7 @@
     "JambaConfig",
     "JetMoeConfig",
     "Lfm2Config",
+    "Lfm2MoeConfig",
     "Lfm2VlConfig",
     "LongcatFlashConfig",
     "Mamba2Config",

src/mobius/_configs/_base.py

@@ -1604,6 +1604,23 @@ def from_transformers(cls, config, parent_config=None) -> Lfm2Config:
         )
 
 
+@dataclasses.dataclass
+class Lfm2MoeConfig(CausalLMConfig):
+    """Configuration for LFM2MoE's dense-prefix and routed-expert feed-forwards."""
+
+    num_dense_layers: int = 2
+    use_expert_bias: bool = True
+
+    @classmethod
+    def from_transformers(cls, config, parent_config=None) -> Lfm2MoeConfig:
+        base = ArchitectureConfig.from_transformers(config, parent_config)
+        return cls(
+            **_shallow_fields(base),
+            num_dense_layers=getattr(config, "num_dense_layers", 2),
+            use_expert_bias=getattr(config, "use_expert_bias", True),
+        )
+
+
 @dataclasses.dataclass
 class Lfm2VlConfig(Lfm2Config):
     """Configuration for LiquidAI LFM2-VL (SigLIP2 NaFlex + LFM2 decoder).

src/mobius/_registry.py

@@ -29,6 +29,7 @@
     Gemma4AssistantConfig,
     Gemma4Config,
     Lfm2Config,
+    Lfm2MoeConfig,
     Lfm2VlConfig,
     MMSConfig,
     MoonshineConfig,
@@ -80,6 +81,7 @@
     InternLM2CausalLMModel,
     LayerNormCausalLMModel,
     Lfm2CausalLMModel,
+    Lfm2MoECausalLMModel,
     Lfm2VlForConditionalGeneration,
     LLaDAModel,
     LLaDAMoEModel,
@@ -499,6 +501,7 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None:
     "internlm2": ModelRegistration(InternLM2CausalLMModel),
     "llama4_text": ModelRegistration(Llama4CausalLMModel),
     "lfm2": ModelRegistration(Lfm2CausalLMModel, config_class=Lfm2Config),
+    "lfm2_moe": ModelRegistration(Lfm2MoECausalLMModel, config_class=Lfm2MoeConfig),
     "dream": ModelRegistration(DreamModel, task="masked-diffusion"),
     "Dream": ModelRegistration(DreamModel, task="masked-diffusion"),
     "llada": ModelRegistration(LLaDAModel, task="masked-diffusion"),

src/mobius/integrations/gguf/_arch_registry.py

@@ -322,14 +322,6 @@
     "ordinary attention would change the model."
 )
 
-_LFM2MOE_GRAPH_REASON = (
-    "LFM2MoE has an arbitrary per-layer attention/short-convolution schedule plus a "
-    "dense-to-sigmoid-routed-MoE transition. Recurrent layers use F32 rolling convolution "
-    "state with copy-on-write sequence reorder and bounded rollback snapshots, while "
-    "attention layers use KV cache. The dense LFM2 graph and ordinary KV/static-cache "
-    "tasks do not own that mixed state or the expert correction-bias semantics."
-)
-
 _ENCODER_RUNTIME_VALIDATION_PENDING = (
     "Config extraction, exact pinned tensor closure, encoder-only task dispatch, and "
     "synthetic ORT execution are covered, but no pinned real GGUF artifact has passed "
@@ -1138,11 +1130,22 @@
     ),
     GGUFArchitectureSpec(
         gguf_arch="lfm2moe",
-        config=Support.DEFERRED,
-        tensor_map=Support.DEFERRED,
-        graph=Support.DEFERRED,
+        model_type="lfm2_moe",
+        tensor_map_recipe=("lfm2", "lfm2_moe_extras"),
+        config_postprocessor="lfm2moe",
+        required_metadata=(
+            "attention.head_count_kv",
+            "attention.layer_norm_rms_epsilon",
+            "rope.freq_base",
+            "shortconv.l_cache",
+            "expert_count",
+            "expert_used_count",
+            "expert_feed_forward_length",
+            "expert_gating_func",
+        ),
         runtime=Support.DEFERRED,
-        reason=_LFM2MOE_GRAPH_REASON,
+        quantized_import=Support.REJECTED,
+        reason=_RECURRENT_RUNTIME_VALIDATION_PENDING + " " + _NO_QUANTIZED_PROJECTION_REASON,
     ),
     GGUFArchitectureSpec(
         gguf_arch="minimax-01",

src/mobius/integrations/gguf/_arch_registry_test.py

@@ -57,7 +57,7 @@
 #: Number of importable architectures. Pinned so that adding support is a
 #: deliberate act that also updates the documented support matrix, and so that
 #: accidentally losing an architecture is a failure rather than a silence.
-_EXPECTED_SUPPORTED_COUNT = 51
+_EXPECTED_SUPPORTED_COUNT = 52
 _FINAL_CENSUS_CLOSURE = frozenset(
     {
         "afmoe",
@@ -268,6 +268,7 @@ def test_every_float_importable_architecture_has_a_quantized_verdict(self) -> No
             "granitehybrid",
             "internlm2",
             "jamba",
+            "lfm2moe",
             "mamba",
             "mamba2",
             "nemotron_h",
@@ -899,14 +900,12 @@ class TestPinnedRemainingHybridCohort:
         "deepseek4",
         "kimi-k3",
         "kimi-linear",
-        "lfm2moe",
     )
     _EXPECTED_TENSOR_COUNTS: ClassVar[dict[str, int]] = {
         "bailingmoe3": 41,
         "deepseek4": 44,
         "kimi-k3": 46,
         "kimi-linear": 40,
-        "lfm2moe": 24,
     }
 
     @pytest.mark.parametrize("architecture", _ARCHITECTURES)
@@ -954,7 +953,6 @@ def test_no_unpinned_alias_or_config_mutation_is_reachable(
             ("deepseek4", ("compressed-cache", "rollback", "ordinary KV")),
             ("kimi-k3", ("matrix state", "latent MoE", "residual banks")),
             ("kimi-linear", ("convolution histories", "matrix state", "correction-bias")),
-            ("lfm2moe", ("rolling convolution", "reorder", "rollback")),
         ],
     )
     def test_state_and_schedule_mismatch_is_explicit(
@@ -965,6 +963,19 @@ def test_state_and_schedule_mismatch_is_explicit(
         for term in state_terms:
             assert term in reason
 
+    def test_lfm2moe_graph_and_import_advance_but_runtime_stays_deferred(self) -> None:
+        spec = try_get_arch_spec("lfm2moe")
+        assert spec is not None
+        assert spec.model_type == "lfm2_moe"
+        assert spec.config is Support.SUPPORTED
+        assert spec.tensor_map is Support.SUPPORTED
+        assert spec.graph is Support.SUPPORTED
+        assert spec.runtime is Support.DEFERRED
+        assert spec.quantized_import is Support.REJECTED
+        assert spec.reason is not None
+        assert "representative real-weight GGUF" in spec.reason
+        assert "keep_quantized=False" in spec.reason
+
     def test_hugging_face_deepseek_v4_registration_remains_valid(self) -> None:
         assert "deepseek_v4" in _REGISTRATIONS
         assert try_get_arch_spec("deepseek4").model_type is None

src/mobius/integrations/gguf/_builder.py

@@ -906,7 +906,7 @@ def _raise_for_invalid_hybrid_tensor_contract(gguf_model) -> None:
     if architecture in {"jamba", "nemotron_h", "granitehybrid"}:
         _raise_for_invalid_mamba_hybrid_tensor_contract(gguf_model)
         return
-    if architecture not in {"lfm2", "qwen35", "qwen35moe", "qwen3next"}:
+    if architecture not in {"lfm2", "lfm2moe", "qwen35", "qwen35moe", "qwen3next"}:
         return
 
     metadata = gguf_model.metadata
@@ -934,7 +934,7 @@ def _raise_for_invalid_hybrid_tensor_contract(gguf_model) -> None:
                 f"{match.group(1)} (block_count={total_layers})"
             )
 
-    if architecture == "lfm2":
+    if architecture in {"lfm2", "lfm2moe"}:
         required_global = {"token_embd.weight", "token_embd_norm.weight"}
         auxiliary = sorted(
             name
@@ -946,13 +946,7 @@ def _raise_for_invalid_hybrid_tensor_contract(gguf_model) -> None:
                 "lfm2 causal-LM import does not support embedding/ColBERT head "
                 f"tensor(s): {auxiliary}"
             )
-        common_suffixes = {
-            "attn_norm.weight",
-            "ffn_norm.weight",
-            "ffn_gate.weight",
-            "ffn_up.weight",
-            "ffn_down.weight",
-        }
+        common_suffixes = {"attn_norm.weight", "ffn_norm.weight"}
         full_suffixes = {
             "attn_q.weight",
             "attn_k.weight",
@@ -1019,6 +1013,30 @@ def _raise_for_invalid_hybrid_tensor_contract(gguf_model) -> None:
         required.update(
             full_suffixes if layer_type == "full_attention" else recurrent_suffixes
         )
+        if architecture == "lfm2":
+            required.update({"ffn_gate.weight", "ffn_up.weight", "ffn_down.weight"})
+        elif architecture == "lfm2moe":
+            dense_ffn = {"ffn_gate.weight", "ffn_up.weight", "ffn_down.weight"}
+            routed_ffn = {
+                "ffn_gate_inp.weight",
+                "ffn_gate_exps.weight",
+                "ffn_up_exps.weight",
+                "ffn_down_exps.weight",
+                "exp_probs_b.bias",
+            }
+            dense_layers = int(metadata.get("lfm2moe.leading_dense_block_count", 0))
+            if layer < dense_layers:
+                required.update(dense_ffn)
+                wrong_ffn = sorted(layer_names & routed_ffn)
+            else:
+                required.update(routed_ffn)
+                wrong_ffn = sorted(layer_names & dense_ffn)
+            if wrong_ffn:
+                raise ValueError(
+                    f"lfm2moe layer {layer} contains tensor(s) from the wrong "
+                    f"{'routed' if layer < dense_layers else 'dense'} FFN family: "
+                    f"{wrong_ffn}"
+                )
 
         if architecture in {"qwen35moe", "qwen3next"}:
             fused = "ffn_gate_up_exps.weight" in layer_names
@@ -1062,7 +1080,7 @@ def _raise_for_invalid_hybrid_tensor_contract(gguf_model) -> None:
                 required.update(
                     {"attn_qkv.weight", "attn_gate.weight"} if modern else {"ssm_in.weight"}
                 )
-            elif architecture != "lfm2":
+            elif architecture not in {"lfm2", "lfm2moe"}:
                 required.update({"attn_qkv.weight", "attn_gate.weight"})
 
         missing = sorted(required - layer_names)
@@ -2070,6 +2088,7 @@ def build_from_gguf(
         )
     if gguf_arch in {
         "lfm2",
+        "lfm2moe",
         "qwen35",
         "qwen35moe",
         "qwen3next",
@@ -2753,28 +2772,42 @@ def _normalize_gguf_weights(
         # Fused stacked gate/up experts [num_experts, 2*out, ...] are split
         # before ordinary stacked-expert unpacking. This handles float weights
         # and packed MatMulNBits companions without dropping either half.
-        fused_marker = ".mlp.experts.gate_up_proj."
-        if fused_marker in key and value.dim() >= 3:
+        fused_marker = next(
+            (
+                marker
+                for marker in (
+                    ".mlp.experts.gate_up_proj.",
+                    ".feed_forward.experts.gate_up_proj.",
+                )
+                if marker in key
+            ),
+            None,
+        )
+        if fused_marker is not None and value.dim() >= 3:
             prefix, suffix = key.rsplit(fused_marker, 1)
+            container = fused_marker.removesuffix(".gate_up_proj.")
             if value.shape[1] % 2:
                 raise ValueError(
                     f"Fused expert tensor {key!r} has odd gate/up width {value.shape[1]}"
                 )
             gate, up = value.chunk(2, dim=1)
             for i in range(value.shape[0]):
-                result[f"{prefix}.mlp.experts.{i}.gate_proj.{suffix}"] = gate[i]
-                result[f"{prefix}.mlp.experts.{i}.up_proj.{suffix}"] = up[i]
+                result[f"{prefix}{container}.{i}.gate_proj.{suffix}"] = gate[i]
+                result[f"{prefix}{container}.{i}.up_proj.{suffix}"] = up[i]
             continue
 
         # Stacked expert weights [num_experts, out, ...] → per-expert.
         unpacked = False
         for proj in ("gate_proj", "up_proj", "down_proj"):
-            marker = f".mlp.experts.{proj}."
-            if marker in key and value.dim() >= 3:
-                prefix, suffix = key.rsplit(marker, 1)
-                for i in range(value.shape[0]):
-                    result[f"{prefix}.mlp.experts.{i}.{proj}.{suffix}"] = value[i]
-                unpacked = True
+            for container in (".mlp.experts", ".feed_forward.experts"):
+                marker = f"{container}.{proj}."
+                if marker in key and value.dim() >= 3:
+                    prefix, suffix = key.rsplit(marker, 1)
+                    for i in range(value.shape[0]):
+                        result[f"{prefix}{container}.{i}.{proj}.{suffix}"] = value[i]
+                    unpacked = True
+                    break
+            if unpacked:
                 break
         if unpacked:
             continue
@@ -3960,8 +3993,12 @@ def _validate_moe_weight_shape(
     if num_experts is None:
         return
     expert_size = getattr(config, "moe_intermediate_size", None) or config.intermediate_size
-    if ".mlp.experts." in name:
-        projection = name.rsplit(".mlp.experts.", 1)[1].split(".", 1)[0]
+    expert_marker = next(
+        (marker for marker in (".mlp.experts.", ".feed_forward.experts.") if marker in name),
+        None,
+    )
+    if expert_marker is not None:
+        projection = name.rsplit(expert_marker, 1)[1].split(".", 1)[0]
         if projection not in {"gate_proj", "up_proj", "down_proj"}:
             return
         expected = (
@@ -3973,7 +4010,7 @@ def _validate_moe_weight_shape(
             raise ValueError(
                 f"Invalid stacked expert shape for {name}: expected {expected}, got {shape}"
             )
-    elif name.endswith(".mlp.gate.weight"):
+    elif name.endswith((".mlp.gate.weight", ".feed_forward.gate.weight")):
         expected = (num_experts, config.hidden_size)
         if shape != expected:
             raise ValueError(