Add dedicated Falcon-H1 model and GGUF import

docs/api/build_from_gguf.md

@@ -15,7 +15,7 @@ from mobius import build_from_gguf
 
 | Census | Total | Closure |
 |---|---:|---|
-| Architectures | 147 | graph verdicts: {'deferred': 90, 'rejected': 3, 'supported': 54}; importable: 52; quantized import: {'rejected': 12, 'supported': 135}; runtime: {'deferred': 144, 'rejected': 3} |
+| Architectures | 147 | graph verdicts: {'deferred': 89, 'rejected': 3, 'supported': 55}; importable: 53; 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` |
@@ -382,7 +382,7 @@ before graph construction or durable output.
 | `exaone-moe` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | EXAONE-MoE serializes a dense trailing NextN block after an iSWA routed/shared expert trunk, but the pinned loader skips appended blocks. Mobius has no exact SWA schedule, remapped global-MTP transform, or executable head contract. |
 | `exaone4` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | EXAONE4 serializes attention/FFN post-norm trailing blocks and NextN tensors with optional synthetic Llama3 RoPE factors, but the pinned loader skips them. No Mobius task owns those preserved-only semantics. |
 | `falcon` | — | model=`falcon`; tensor=`falcon` | 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. |
-| `falcon-h1` | — | none (fails before config extraction) | not claimed | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | Falcon-H1 executes attention and Mamba2 in parallel in every block and requires KV and recurrent states simultaneously; FalconCausalLMModel is not compatible with that graph or state ABI. |
+| `falcon-h1` | — | model=`falcon_h1`; tensor=`falcon_h1` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=supported | The dedicated graph and GGUF importer preserve parallel Attention+Mamba2 layers and their four-state ABI, but runtime packaging remains deferred pending heterogeneous-state schema support (onnxruntime/mobius#605) and real full-logit plus deterministic stateful-generation evidence. |
 | `gemma` | — | model=`gemma`; 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. |
 | `gemma-embedding` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | Gemma Embedding is a bidirectional stateless embedding graph with alternating sliding-window attention, four norm sites, sqrt(hidden) embedding scaling, and optional pooling/dense modules. Causal Gemma3 text/VLM tasks expose the wrong ABI. |
 | `gemma2` | — | model=`gemma2`; tensor=`llama`+`gemma2_extras` | 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. |
@@ -583,13 +583,20 @@ 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`, `plamo2`, and `falcon-h1` are deferred before config extraction.
-MiniMax-01's pinned Lightning schedule and decay/scaling semantics do not match
-the current graph; PLaMo2 needs a dedicated fused-QKV Mamba1/attention graph;
-Falcon-H1 executes attention and Mamba2 in parallel in every block and therefore
-needs KV plus conv/SSM states simultaneously. It is not an alias of ordinary
-Falcon. `nemotron_h_moe` remains rejected because its folded MTP attention+MoE
-head has no equivalent package contract.
+`minimax-01` and `plamo2` are deferred before config extraction. MiniMax-01's
+pinned Lightning schedule and decay/scaling semantics do not match the current
+graph; PLaMo2 needs a dedicated fused-QKV Mamba1/attention graph.
+
+Falcon-H1 graph/import support is dedicated rather than aliased to Falcon. Every
+layer exposes `(key, value, conv_state, ssm_state)` in that order. Static cache is
+rejected, and quantized-source imports require `--dequantize` until the mixed graph
+can preserve only exact attention/FFN MatMul roles. Runtime packaging remains
+deferred pending the heterogeneous-state schema tracked by
+[`onnxruntime/mobius#605`](https://github.com/onnxruntime/mobius/issues/605) and
+real full-logit plus deterministic stateful-generation evidence.
+
+`nemotron_h_moe` remains rejected because its folded MTP attention+MoE head has
+no equivalent package contract.
 
 ### Audio/TTS/codec cohort
 

src/mobius/__main__.py

@@ -174,6 +174,11 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
             from mobius._registry import _TEXT_ONLY_MODEL_TYPE
 
             model_type = _TEXT_ONLY_MODEL_TYPE.get(model_type, model_type)
+        if model_type == "falcon_h1":
+            raise ValueError(
+                "--static-cache cannot represent Falcon-H1's per-layer K, V, "
+                "convolution, and SSM states"
+            )
         if model_type == "gemma4":
             from mobius.tasks._gemma4 import Gemma4Task
 

src/mobius/_configs/__init__.py

@@ -31,6 +31,7 @@
     DFlashConfig,
     Eagle3Config,
     EncoderConfig,
+    FalconH1Config,
     Gemma2Config,
     Gemma3nConfig,
     Gemma3nMultiModalConfig,
@@ -107,6 +108,7 @@
     "DepthAnythingConfig",
     "Eagle3Config",
     "EncoderConfig",
+    "FalconH1Config",
     "Gemma2Config",
     "Gemma3nAudioConfig",
     "Gemma3nConfig",

src/mobius/_configs/_base.py

@@ -2691,6 +2691,150 @@ def from_transformers(cls, config, parent_config=None) -> BambaConfig:
         )
 
 
+@dataclasses.dataclass
+class FalconH1Config(ArchitectureConfig):
+    """Configuration for Falcon-H1 parallel Attention + Mamba2 decoder layers."""
+
+    mamba_d_ssm: int = 1024
+    mamba_n_heads: int = 128
+    mamba_d_head: int = 8
+    mamba_n_groups: int = 1
+    mamba_d_state: int = 256
+    mamba_d_conv: int = 4
+    mamba_expand: int = 2
+    mamba_chunk_size: int = 256
+    mamba_conv_bias: bool = True
+    mamba_proj_bias: bool = False
+    mamba_norm_before_gate: bool = True
+    mamba_rms_norm: bool = False
+    time_step_limit: tuple[float, float] = (0.0, float("inf"))
+    attention_bias: bool = False
+    projectors_bias: bool = False
+    lm_head_multiplier: float = 1.0
+    embedding_multiplier: float = 1.0
+    mlp_multipliers: tuple[float, float] = (1.0, 1.0)
+    key_multiplier: float = 1.0
+    attention_out_multiplier: float = 1.0
+    attention_in_multiplier: float = 1.0
+    ssm_multipliers: tuple[float, float, float, float, float] = (
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+        1.0,
+    )
+    ssm_in_multiplier: float = 1.0
+    ssm_out_multiplier: float = 1.0
+
+    def __post_init__(self) -> None:
+        self.attn_qkv_bias = self.attention_bias
+        self.attn_o_bias = self.attention_bias
+        if self.hidden_act != "silu":
+            raise ValueError("Falcon-H1 supports only hidden_act='silu'")
+        if (
+            self.head_dim <= 0
+            or self.num_attention_heads <= 0
+            or self.hidden_size != self.num_attention_heads * self.head_dim
+        ):
+            raise ValueError("Falcon-H1 hidden_size must equal num_attention_heads * head_dim")
+        if (
+            self.num_key_value_heads <= 0
+            or self.num_attention_heads % self.num_key_value_heads
+        ):
+            raise ValueError("Falcon-H1 num_key_value_heads must divide num_attention_heads")
+        if self.mamba_d_ssm <= 0 or self.mamba_n_heads <= 0:
+            raise ValueError("Falcon-H1 Mamba dimensions must be positive")
+        if self.mamba_d_ssm % self.mamba_n_heads:
+            raise ValueError("mamba_n_heads must divide mamba_d_ssm")
+        if self.mamba_d_head * self.mamba_n_heads != self.mamba_d_ssm:
+            raise ValueError("mamba_d_head * mamba_n_heads must equal mamba_d_ssm")
+        if (
+            self.mamba_n_groups <= 0
+            or self.mamba_n_heads % self.mamba_n_groups
+            or self.mamba_d_ssm % self.mamba_n_groups
+        ):
+            raise ValueError("mamba_n_groups must divide both mamba_n_heads and mamba_d_ssm")
+        if self.mamba_d_state <= 0 or self.mamba_d_conv <= 0 or self.mamba_chunk_size <= 0:
+            raise ValueError("Falcon-H1 state, convolution, and chunk sizes must be positive")
+        if len(self.mlp_multipliers) != 2:
+            raise ValueError("mlp_multipliers must contain exactly two values")
+        if len(self.ssm_multipliers) != 5:
+            raise ValueError("ssm_multipliers must contain exactly five values")
+        if len(self.time_step_limit) != 2:
+            raise ValueError("time_step_limit must contain exactly two values")
+        time_step_min, time_step_max = self.time_step_limit
+        if time_step_min < 0 or time_step_max < time_step_min:
+            raise ValueError("time_step_limit must be ordered and non-negative")
+        multipliers = (
+            self.embedding_multiplier,
+            self.lm_head_multiplier,
+            self.attention_in_multiplier,
+            self.attention_out_multiplier,
+            self.key_multiplier,
+            self.ssm_in_multiplier,
+            self.ssm_out_multiplier,
+            *self.mlp_multipliers,
+            *self.ssm_multipliers,
+        )
+        if not all(math.isfinite(value) for value in multipliers):
+            raise ValueError("Falcon-H1 multipliers must all be finite")
+
+    @classmethod
+    def from_transformers(cls, config, parent_config=None) -> FalconH1Config:
+        base = ArchitectureConfig.from_transformers(config, parent_config)
+        fields = _shallow_fields(base)
+        fields.pop("embedding_multiplier", None)
+        fields["mlp_bias"] = bool(getattr(config, "mlp_bias", False))
+        d_ssm = getattr(config, "mamba_d_ssm", None)
+        if d_ssm is None:
+            d_ssm = int(getattr(config, "mamba_expand", 2)) * base.hidden_size
+        n_heads = int(getattr(config, "mamba_n_heads", 128))
+        d_head = getattr(config, "mamba_d_head", "auto")
+        if d_head == "auto":
+            if d_ssm % n_heads:
+                raise ValueError("mamba_n_heads must divide mamba_d_ssm")
+            d_head = d_ssm // n_heads
+        time_step_limit = getattr(config, "time_step_limit", None) or (
+            0.0,
+            float("inf"),
+        )
+        return cls(
+            **fields,
+            mamba_d_ssm=int(d_ssm),
+            mamba_n_heads=n_heads,
+            mamba_d_head=int(d_head),
+            mamba_n_groups=int(getattr(config, "mamba_n_groups", 1)),
+            mamba_d_state=int(getattr(config, "mamba_d_state", 256)),
+            mamba_d_conv=int(getattr(config, "mamba_d_conv", 4)),
+            mamba_expand=int(getattr(config, "mamba_expand", 2)),
+            mamba_chunk_size=int(getattr(config, "mamba_chunk_size", 256)),
+            mamba_conv_bias=bool(getattr(config, "mamba_conv_bias", True)),
+            mamba_proj_bias=bool(getattr(config, "mamba_proj_bias", False)),
+            mamba_norm_before_gate=bool(getattr(config, "mamba_norm_before_gate", True)),
+            mamba_rms_norm=bool(getattr(config, "mamba_rms_norm", False)),
+            time_step_limit=tuple(float(value) for value in time_step_limit),
+            attention_bias=bool(getattr(config, "attention_bias", False)),
+            projectors_bias=bool(getattr(config, "projectors_bias", False)),
+            lm_head_multiplier=float(getattr(config, "lm_head_multiplier", 1.0)),
+            embedding_multiplier=float(getattr(config, "embedding_multiplier", 1.0)),
+            mlp_multipliers=tuple(
+                float(value)
+                for value in (getattr(config, "mlp_multipliers", None) or (1.0, 1.0))
+            ),
+            key_multiplier=float(getattr(config, "key_multiplier", 1.0)),
+            attention_out_multiplier=float(getattr(config, "attention_out_multiplier", 1.0)),
+            attention_in_multiplier=float(getattr(config, "attention_in_multiplier", 1.0)),
+            ssm_multipliers=tuple(
+                float(value)
+                for value in (
+                    getattr(config, "ssm_multipliers", None) or (1.0, 1.0, 1.0, 1.0, 1.0)
+                )
+            ),
+            ssm_in_multiplier=float(getattr(config, "ssm_in_multiplier", 1.0)),
+            ssm_out_multiplier=float(getattr(config, "ssm_out_multiplier", 1.0)),
+        )
+
+
 @dataclasses.dataclass
 class GraniteMoeHybridConfig(BambaConfig):
     """Configuration for GraniteMoeHybrid: Mamba2+Attention hybrid with MoE on all layers.

src/mobius/_registry.py

@@ -25,6 +25,7 @@
 from mobius._configs import (
     BaseModelConfig,
     Eagle3Config,
+    FalconH1Config,
     Gemma3nMultiModalConfig,
     Gemma4AssistantConfig,
     Gemma4Config,
@@ -140,6 +141,7 @@
     FalconCausalLMModel,
     MPTCausalLMModel,
 )
+from mobius.models.falcon_h1 import FalconH1ForCausalLM
 from mobius.models.fun_asr import FunASRForConditionalGeneration
 from mobius.models.gemma3n import Gemma3nCausalLMModel, Gemma3nMultiModalModel
 from mobius.models.glm_asr import GlmAsrForConditionalGeneration
@@ -450,6 +452,13 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None:
     "ernie4_5": ModelRegistration(ErnieCausalLMModel),
     "exaone4": ModelRegistration(ExaOne4CausalLMModel),
     "falcon": ModelRegistration(FalconCausalLMModel),
+    "falcon_h1": ModelRegistration(
+        FalconH1ForCausalLM,
+        task="falcon-h1-text-generation",
+        config_class=FalconH1Config,
+        test_model_id="tiiuae/Falcon-H1-Tiny-90M-Base",
+        family="falcon-h1",
+    ),
     "gemma": ModelRegistration(GemmaCausalLMModel),
     "gemma2": ModelRegistration(Gemma2CausalLMModel),
     "gemma3": ModelRegistration(Gemma3MultiModalModel, task="vision-language"),

src/mobius/components/_attention.py

@@ -263,6 +263,7 @@ def __init__(
         self.num_attention_heads = config.num_attention_heads
         self.num_key_value_heads = config.num_key_value_heads
         self.scaling = scale if scale is not None else self.head_dim**-0.5
+        self._key_multiplier = getattr(config, "key_multiplier", 1.0)
         # NoPE models leave ``partial_rotary_factor`` as ``None``; treat as
         # the inert 1.0 for the purpose of computing ``rotary_embedding_dim``
         # (which will itself be 0, i.e. no partial-RoPE splitting). The
@@ -325,6 +326,8 @@ def forward(
         query_states = self.q_proj(op, hidden_states)
         key_states = self.k_proj(op, hidden_states)
         value_states = self.v_proj(op, hidden_states)
+        if not math.isclose(self._key_multiplier, 1.0):
+            key_states = op.Mul(key_states, self._key_multiplier)
 
         if self.q_norm is not None and self.k_norm is not None:
             if self._qk_norm_full: