Add full Jamba GGUF MoE 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': 3, 'supported': 56}; importable: 54; quantized import: {'rejected': 12, 'supported': 135}; runtime: {'deferred': 144, 'rejected': 3} |
+| Architectures | 147 | graph verdicts: {'deferred': 88, 'rejected': 3, 'supported': 56}; importable: 54; quantized import: {'rejected': 11, 'supported': 136}; 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` |
@@ -411,7 +411,7 @@ before graph construction or durable output.
 | `internlm2` | — | model=`internlm2`; tensor=`llama` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=rejected | 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. 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. |
 | `jais` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | JAIS combines fused biased QKV, causal ALiBi, 1/head_dim attention scaling, parallel SwiGLU, and converter-baked MuP embedding/output scales. Reusing Falcon or Llama would lose required value transforms. |
 | `jais2` | — | none (fails before config extraction) | audited-direct-loader-conditional-union | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | JAIS2 is a distinct RoPE, bias-bearing LayerNorm decoder with split Q/K/V and a non-gated ReLU-squared FFN. It is not the ALiBi/SwiGLU JAIS graph and has no exact Mobius tensor recipe. |
-| `jamba` | — | model=`jamba`; tensor=`jamba` | 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. |
+| `jamba` | — | model=`jamba`; tensor=`jamba` | not claimed | config=supported; tensor_map=supported; graph=supported; runtime=deferred; quantized_import=supported | Exact mixed attention/Mamba and dense/routed-MoE schedules, strict tensor closure and shapes, GGUF value transforms, compatible projection quantization, value-checked expert ordering, reduced Transformers parity, and multi-token ORT state threading, reorder, and replay are covered. Generic ORT GenAI runtime packaging remains deferred because its released cache schema cannot represent heterogeneous KV, convolution, and recurrent state slots; tracked by #605. |
 | `jina-bert-v2` | — | none (fails before config extraction) | not claimed | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | JinaBERT v2 uses ALiBi, optional full-width Q/K norms, an extra attention norm, and either separate or fused GeGLU inputs. Mobius has no graph with that exact combination. |
 | `jina-bert-v3` | — | none (fails before config extraction) | not claimed | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | JinaBERT v3 uses RoPE and may alternate dense GELU and routed MoE layers. BertModel has absolute positions and no MoE path. |
 | `kimi-k3` | — | none (fails before config extraction) | not claimed | config=deferred; tensor_map=deferred; graph=deferred; runtime=deferred; quantized_import=supported | Kimi-K3 alternates KDA recurrent and NoPE MLA layers and requires convolution plus matrix state, sigmoid routed latent MoE with SiTU experts, optional shared experts, and cross-layer residual banks. Kimi-Linear is not an alias: its gates, expert activation, latent projections, and residual contract differ. Mobius has no exact graph or mixed-state task for either ABI. |
@@ -559,19 +559,26 @@ independent full-logit and generation parity.
 ### Second hybrid cohort
 
 `jamba`, `nemotron_h`, and `granitehybrid` have graph-import support for exact
-dense subsets only; runtime packaging remains deferred pending independent
+pinned subsets. Jamba includes routed MoE layers; Nemotron-H and GraniteHybrid
+remain dense-only. Runtime packaging remains deferred pending independent
 real-artifact full-logit and stateful-generation parity.
 
 - Schedules come from suffix-exact per-layer metadata. Jamba and GraniteHybrid
   use `attention.head_count_kv` (`0` selects Mamba/Mamba2). Nemotron-H combines
   that array with per-layer `feed_forward_length` to select exactly one of
   Mamba2, attention, or dense ReLU² MLP.
-- Jamba requires `ssm.inner_size == 2 * embedding_length`. Nemotron-H rejects
-  MTP and all MoE files. GraniteHybrid rejects routed-MoE files until 3-D expert
-  fusion, ordering, and quantized preservation have independent value tests.
+- Jamba requires `ssm.inner_size == 2 * embedding_length`, SiLU experts,
+  Mamba-1 with biased depthwise convolution and bias-free projections, and
+  softmax-first top-k routing without post-top-k renormalization. Routed layers
+  are inferred exactly from `ffn_gate_inp`; stacked expert tensors are split in
+  numeric order. There are no shared experts. Nemotron-H rejects MTP and all MoE
+  files. GraniteHybrid rejects routed-MoE files until 3-D expert fusion,
+  ordering, and quantized preservation have independent value tests.
 - Every layer must provide exactly its pinned loader tensor family. Missing,
   wrong-mixer, partial, auxiliary, scale/input-scale, and out-of-range tensors
-  are rejected before graph construction. GGUF Mamba decay values are inverted
+  are rejected before graph construction. Compatible attention, dense-FFN, and
+  expert MatMul weights may remain quantized; Mamba and other state-sensitive
+  tensors are dequantized. GGUF Mamba decay values are inverted
   from `-exp(A_log)`; convolution and grouped Mamba2 tensors are restored to
   graph shapes.
 - State inputs and outputs are caller-owned. Mamba1 uses conv

src/mobius/_configs/_base.py

@@ -2592,6 +2592,9 @@ class JambaConfig(ArchitectureConfig):
     attn_layer_offset: int = 4
     expert_layer_period: int = 2
     expert_layer_offset: int = 1
+    # GGUF serializes the resolved schedule through tensor presence rather than
+    # preserving the source period/offset pair.
+    expert_layer_indices: list[int] | None = None
 
     @classmethod
     def from_transformers(cls, config, parent_config=None) -> JambaConfig:
@@ -2601,18 +2604,33 @@ def from_transformers(cls, config, parent_config=None) -> JambaConfig:
         n = base.num_hidden_layers
         attn_period = getattr(config, "attn_layer_period", 8)
         attn_offset = getattr(config, "attn_layer_offset", 4)
+        expert_period = getattr(config, "expert_layer_period", 2)
+        expert_offset = getattr(config, "expert_layer_offset", 1)
+        if attn_period <= 0 or not 0 <= attn_offset < attn_period:
+            raise ValueError("Jamba attn_layer_offset must be in [0, attn_layer_period)")
+        if expert_period <= 0 or not 0 <= expert_offset < expert_period:
+            raise ValueError("Jamba expert_layer_offset must be in [0, expert_layer_period)")
         layer_types = []
         for i in range(n):
-            if (i - attn_offset) % attn_period == 0:
+            if i % attn_period == attn_offset:
                 layer_types.append("full_attention")
             else:
                 layer_types.append("mamba")
 
         num_experts = getattr(config, "num_experts", 16)
         num_experts_per_tok = getattr(config, "num_experts_per_tok", 2)
+        dt_rank = getattr(config, "mamba_dt_rank", "auto")
+        if dt_rank == "auto":
+            dt_rank = math.ceil(base.hidden_size / 16)
 
         # Exclude fields we set explicitly below to avoid duplicate keyword args
-        _exclude = {"layer_types", "num_local_experts", "num_experts_per_tok"}
+        _exclude = {
+            "layer_types",
+            "num_local_experts",
+            "num_experts_per_tok",
+            "norm_topk_prob",
+            "rope_type",
+        }
         base_fields = {k: v for k, v in _shallow_fields(base).items() if k not in _exclude}
         return cls(
             **base_fields,
@@ -2623,13 +2641,16 @@ def from_transformers(cls, config, parent_config=None) -> JambaConfig:
             mamba_d_state=getattr(config, "mamba_d_state", 16),
             mamba_d_conv=getattr(config, "mamba_d_conv", 4),
             mamba_expand=getattr(config, "mamba_expand", 2),
-            mamba_dt_rank=getattr(config, "mamba_dt_rank", 256),
+            mamba_dt_rank=int(dt_rank),
             mamba_conv_bias=getattr(config, "mamba_conv_bias", True),
             mamba_proj_bias=getattr(config, "mamba_proj_bias", False),
             attn_layer_period=attn_period,
             attn_layer_offset=attn_offset,
-            expert_layer_period=getattr(config, "expert_layer_period", 2),
-            expert_layer_offset=getattr(config, "expert_layer_offset", 1),
+            expert_layer_period=expert_period,
+            expert_layer_offset=expert_offset,
+            expert_layer_indices=[i for i in range(n) if i % expert_period == expert_offset],
+            norm_topk_prob=False,
+            rope_type=None,
         )
 
 

src/mobius/components/_ssm.py

@@ -188,6 +188,60 @@ def _project_ssm_params(self, op, x_db):
         c_mat = self.c_layernorm(op, c_mat)
         return dt_raw, b_mat, c_mat
 
+    def _repeat_for_channels(self, op: OpBuilder, value: ir.Value) -> ir.Value:
+        # Jamba Mamba-1 shares token-dependent B/C across all expanded channels.
+        value = op.Unsqueeze(value, [2])  # (B, T, 1, d_state)
+        value = op.Tile(value, [1, 1, self.d_inner, 1])
+        return op.Reshape(value, [0, 0, self.d_inner * self.d_state])
+
+    def forward(
+        self,
+        op: OpBuilder,
+        x: ir.Value,
+        ssm_state: ir.Value,
+        padding_mask: ir.Value | None = None,
+    ):
+        """Run the complete multi-token Jamba selective scan."""
+        dt_raw, b_mat, c_mat = self._project_ssm_params(op, self.x_proj(op, x))
+        dt = op.Softplus(op.Cast(self.dt_proj(op, dt_raw), to=ir.DataType.FLOAT))
+        del padding_mask
+
+        x_f32 = op.Cast(x, to=ir.DataType.FLOAT)
+        decay = op.Mul(
+            op.Unsqueeze(dt, [-1]),
+            op.Unsqueeze(
+                op.Neg(op.Exp(op.Cast(self.A_log, to=ir.DataType.FLOAT))),
+                [0, 1],
+            ),
+        )
+        decay = op.Reshape(decay, [0, 0, self.d_inner * self.d_state])
+        value = op.Mul(dt, x_f32)
+        internal_state = op.Unsqueeze(
+            op.Cast(ssm_state, to=ir.DataType.FLOAT),
+            [-1],
+        )
+        output, present_state = op.LinearAttention(
+            self._repeat_for_channels(op, op.Cast(c_mat, to=ir.DataType.FLOAT)),
+            self._repeat_for_channels(op, op.Cast(b_mat, to=ir.DataType.FLOAT)),
+            value,
+            internal_state,
+            decay,
+            scale=1.0,
+            q_num_heads=self.d_inner,
+            kv_num_heads=self.d_inner,
+            update_rule="gated",
+            _domain="com.microsoft",
+            _outputs=2,
+        )
+        output = op.Add(
+            output,
+            op.Mul(x_f32, op.Cast(self.D, to=ir.DataType.FLOAT)),
+        )
+        return (
+            op.CastLike(output, x),
+            op.CastLike(op.Squeeze(present_state, [-1]), ssm_state),
+        )
+
 
 class Mamba2Scan(nn.Module):
     """Multi-head selective scan for Mamba2/SSD architecture.

src/mobius/integrations/gguf/_arch_registry.py

@@ -261,6 +261,15 @@
     "stateful ORT generation. Runtime packaging remains deferred until that evidence exists."
 )
 
+_JAMBA_RUNTIME_VALIDATION_PENDING = (
+    "Exact mixed attention/Mamba and dense/routed-MoE schedules, strict tensor closure "
+    "and shapes, GGUF value transforms, compatible projection quantization, value-checked "
+    "expert ordering, reduced Transformers parity, and multi-token ORT state threading, "
+    "reorder, and replay are covered. Generic ORT GenAI runtime packaging remains deferred "
+    "because its released cache schema cannot represent heterogeneous KV, convolution, and "
+    "recurrent state slots; tracked by #605."
+)
+
 _RWKV_GRAPH_REASONS = {
     "rwkv6": (
         "RWKV6 carries two F32 states per layer (two token-shift vectors and a per-head "
@@ -1063,8 +1072,7 @@
             "ssm.time_step_rank",
         ),
         runtime=Support.DEFERRED,
-        quantized_import=Support.REJECTED,
-        reason=_RECURRENT_RUNTIME_VALIDATION_PENDING + " " + _NO_QUANTIZED_PROJECTION_REASON,
+        reason=_JAMBA_RUNTIME_VALIDATION_PENDING,
     ),
     GGUFArchitectureSpec(
         gguf_arch="nemotron_h",

src/mobius/integrations/gguf/_arch_registry_test.py

@@ -162,6 +162,7 @@
         "gpt2",
         "granitemoe",
         "hunyuan-dense",
+        "jamba",
         "lfm2",
         "llada",
         "llada-moe",
@@ -270,7 +271,6 @@ def test_every_float_importable_architecture_has_a_quantized_verdict(self) -> No
             "chatglm",
             "granitehybrid",
             "internlm2",
-            "jamba",
             "lfm2moe",
             "mamba",
             "mamba2",