fix(models): Jamba never builds a quantized MoE expert plane #974

Background

PR #972 (issue #958, merged as db6cce9dc) bounded the declared quantization parameters where they are stored, adding mlxcel_core::layers::validate_quantization_params(group_size, bits) (src/lib/mlxcel-core/src/layers.rs:1292) and calling it at the shared loaders plus roughly twenty family-local expert-plane constructors, and making QuantizedEmbedding::new and QuantizedMultiLinear::new return Result. The defect class it exists to close: a config-derived value passes every Rust-side check, violates an undocumented MLX precondition, and terminates the process at the FIRST forward pass rather than failing at load, because quantized_matmul and gather_qmm cross the cxx bridge as UniquePtr<MlxArray> rather than Result (src/lib/mlxcel-core/src/lib.rs:1092, :1108), so the C++ throw is an uncatchable std::terminate that catch_unwind cannot contain.

#958 named eighteen families in scope. While implementing it, the inventory was checked family by family with grep, and Jamba was the one entry of the eighteen that did not match: it has a Quantized expert variant, but nothing anywhere constructs it, so there was no storage point to bound. PR #972 left src/models/jamba.rs unchanged and recorded the finding rather than adding a guard with no call site. This issue is that finding.

The finding, confirmed on current main

What exists. src/models/jamba.rs:337 declares a family-local SwitchLinear enum with two variants. The Quantized variant is defined at :338 and carries weight, scales, biases, group_size and bits. Its forward arm at :353 calls mlxcel_core::gather_qmm correctly, with the mode hardcoded to "affine" at :376. The Regular variant is defined at :345 and its forward arm at :380 does mlxcel_core::take followed by mlxcel_core::matmul. The enum carries #[allow(dead_code)] at :336, under a comment at :335 saying it is kept local because it uses take plus matmul rather than gather_mm.

What is missing. There is no construction site for Quantized anywhere in the tree. JambaModel::from_weights (:1102) takes the MoE branch at :1222 when config.is_expert_layer(i) holds, reads the three expert planes by name, and hands each straight to SwitchLinear::Regular at :1247, :1250 and :1251:

switch_mlp: SwitchGLU {
    gate_proj: SwitchLinear::Regular { weight: gate_weight },
    up_proj: SwitchLinear::Regular { weight: up_weight },
    down_proj: SwitchLinear::Regular { weight: down_weight },
},

There is no .scales lookup anywhere in that branch. group_size and bits are computed at :1115 and :1116 and are used only by the embedding at :1119 and the dense UnifiedLinear path at :1258 onward; they never reach the MoE branch. Every other MoE family in the tree branches on scales presence, either through the shared switch_layers::SwitchLinear::from_weights or through a family-local equivalent such as src/models/deepseek.rs:847. #[allow(dead_code)] at :336 is therefore suppressing the symptom of a missing path rather than documenting a deliberate choice.

The per-expert to stacked remap in sanitize_weights (:1053, MoE loop at :1073) has the same shape: the inner while let Some(w) = weights.remove(&format!("{}.experts.{}.{}.weight", base, e, proj)) at :1086 matches only .weight, and the stacked insert at :1094 writes only switch_mlp.{proj}.weight. A genuine per-expert quantized export would arrive at the loader with its .scales and .biases left behind in the map, unstacked and unused.

The two readings, and what distinguishes them

Reading A: the quantized path is genuinely dead code and should be removed, because no real Jamba checkpoint ships quantized expert planes.

Reading B: Jamba silently fails to use quantized experts on a checkpoint that ships them, which is a real functional gap.

What settles it is whether a public quantized Jamba checkpoint with num_experts > 1 exists and what its expert tensors look like. Evidence gathered so far, all of which the implementer should re-check rather than take on trust:

• models/jamba-v0.1-4bit cannot distinguish them. Its config.json declares "num_experts": 1, and JambaConfig::is_expert_layer (:215) requires num_experts > 1, so the MoE branch is never entered and all 28 layers load through the dense UnifiedLinear path at :1258. Its 804 tensors contain 252 feed_forward.* and zero switch_mlp.* or .experts.*. The dense projections are quantized (171 .scales and 171 .biases across the checkpoint), so this is a real 4-bit checkpoint, just not a MoE one. The directory name is misleading: its config.json is byte-identical to mlx-community/AI21-Jamba-Reasoning-3B-4bit (28 layers, hidden 2560, 1.6 GB on disk), and mlx-community/Jamba-v0.1-4bit does not resolve on the Hub. mlx-community currently publishes only the Reasoning-3B conversions, all of which declare num_experts: 1.
• A public MoE conversion does exist outside mlx-community. ncls-p/AI21-Jamba2-Mini-mlx-4Bit declares model_type: jamba, num_experts: 16, num_experts_per_tok: 2, 32 layers, hidden 4096, and quantization: {"group_size": 64, "bits": 4, "mode": "affine"}. Its index lists 1015 tensors, of which 144 are switch_mlp: 48 .weight, 48 .scales and 48 .biases, which is 16 expert layers times three projections, already pre-stacked. That is exactly the layout the missing Quantized branch would consume. It is roughly 29 GB and is not present locally.

If that checkpoint behaves as the code predicts, reading B is correct. The prediction, stated so it can be falsified: sanitize_weights finds no .experts.N keys (the planes are already stacked), so its remap loop is a no-op and the .scales / .biases survive untouched in the map; from_weights enters the MoE branch on the odd layers, weights.remove(".../switch_mlp.gate_proj.weight") returns the packed uint32 plane, and SwitchLinear::Regular is built from it; the first forward pass then reaches mlxcel_core::matmul on a uint32 operand, which MLX rejects with [matmul] Only inexact types are supported but ... (ml-explore/mlx mlx/ops.cpp), and that throw crosses the bridge as an uncatchable abort rather than a load error. Confirming or refuting that first run is the highest-value step in this issue.

Note the confound to rule out before concluding: Jamba2 may differ from Jamba in ways unrelated to quantization, so a load failure on AI21-Jamba2-Mini-mlx-4Bit is not by itself evidence about the expert path. If the checkpoint is unusable for an unrelated reason, fall back to driving JambaModel::from_weights over a synthetic weight map that carries a pre-stacked quantized expert plane, which is the pattern PR #972 used across sixteen families through the insert_stacked_quantized_expert_plane helper in src/models/switch_layers.rs.

Proposed Solution

Contingent on the determination above. If reading B holds:

1. In the MoE branch at src/models/jamba.rs:1225 onward, look up {prefix}.feed_forward.switch_mlp.{proj}.scales and .biases alongside .weight, and select SwitchLinear::Quantized when scales are present.
2. Prefer routing through the shared switch_layers::SwitchLinear over the family-local enum if the shape handling allows it. The comment at :335 justifies the local enum by the Regular arm using take plus matmul rather than gather_mm, which does not by itself justify a private quantized arm, since the quantized arm at :353 is already an ordinary gather_qmm call.
3. Whichever path is taken, run switch_layers::validate_expert_quantization_params(prefix, group_size, bits) (src/models/switch_layers.rs:171) on the declared values before storing them, so the family does not reintroduce the fix(models): bound quantization params in the family-local MoE expert loaders #958 abort class the moment the path becomes live.
4. Extend the per-expert stacking at :1073 to carry .scales and .biases, or route it through switch_layers::stack_individual_experts (src/models/switch_layers.rs:535), which already handles all three planes.
5. Remove #[allow(dead_code)] from :336 once the variant is constructed.
6. The local Quantized arm hardcodes "affine" at :376 and takes a non-optional biases, which is right for the affine export above but cannot represent a block-float one. Decide whether to support that here or to fail at load with a clear message. Related: fix(core): bound the declared quantization mode string, not just group_size and bits #973 covers the mode string being unbounded generally.

If reading A holds instead, delete the Quantized variant and the #[allow(dead_code)] attribute together, and record in the file why a quantized Jamba MoE plane is out of scope, so the next reader does not re-add a half-path.

Inspection and reproduction

Confirm the gap without building anything:

# The variant exists.
grep -n "enum SwitchLinear" -A 16 src/models/jamba.rs
# Nothing constructs it. Only Regular appears, at three sites.
grep -n "SwitchLinear::Quantized\|SwitchLinear::Regular" src/models/jamba.rs
# The MoE branch never looks for scales.
sed -n '1222,1258p' src/models/jamba.rs
# The per-expert remap carries only .weight.
sed -n '1073,1096p' src/models/jamba.rs

Confirm the local checkpoint cannot exercise the path:

python3 -c "import json; c=json.load(open('models/jamba-v0.1-4bit/config.json')); print({k: c[k] for k in ['model_type','num_experts','num_experts_per_tok','num_hidden_layers','hidden_size']})"
python3 -c "import json; m=json.load(open('models/jamba-v0.1-4bit/model.safetensors.index.json'))['weight_map']; print('total', len(m), '| switch_mlp', len([k for k in m if 'switch_mlp' in k]), '| experts.', len([k for k in m if '.experts.' in k]))"

Settle the determination (roughly 29 GB download, and check free space first):

DEVELOPER_DIR=/Applications/Xcode-26.6.0.app/Contents/Developer \
  cargo build --release --features metal,accelerate
hf download ncls-p/AI21-Jamba2-Mini-mlx-4Bit --local-dir models/jamba2-mini-4bit
./target/release/mlxcel generate -m models/jamba2-mini-4bit -p "The capital of France is" -n 20 --no-chat-template

Local validation: possible and not possible

Possible here. The quantized MoE checkpoints available under models/ in the repo root are gpt-oss-20b-mxfp4, deepseek-v2-lite-4bit, qwen3-30b-a3b-4bit, qwen3-moe-4bit, mixtral-8x7b-4bit, jamba-v0.1-4bit, nemotron-h-30b-4bit and llama-4-scout-17b-4bit. models/jamba-v0.1-4bit is the regression control for this change and nothing else: it proves the dense path is untouched, since it never enters the MoE branch. Unit-level validation of the new quantized branch is fully possible locally by driving JambaModel::from_weights over a synthetic pre-stacked expert plane, following the sixteen-family pattern PR #972 established.

Not possible here without a download. No local checkpoint has num_experts > 1, so neither the positive case (a quantized Jamba MoE plane loads and generates coherently) nor the negative case (today's Regular arm aborts on packed weights) can be observed against real weights until ncls-p/AI21-Jamba2-Mini-mlx-4Bit or an equivalent is fetched. deepseek-v2-lite-4bit is not a usable cross-family control right now either: it produces degenerate output for an unrelated reason tracked as #991.

Machine notes

• Pass DEVELOPER_DIR=/Applications/Xcode-26.6.0.app/Contents/Developer on every cargo invocation. Without it cmake fails with xcrun: error: unable to find utility "metal".
• Only the release profile is warm (target/release exists, target/debug does not). A plain debug cargo check triggers a cold MLX C++ build. Use cargo check --release --features metal,accelerate and cargo test --release.
• A test that actually trips an MLX C++ throw aborts the whole test binary with SIGABRT rather than failing cleanly, because the bridge functions return UniquePtr<MlxArray> and cannot propagate the exception. Write the bounds-check regression to assert on the load Result and never run a forward pass, which is the convention PR fix(models): bound quantization params in the family-local MoE expert loaders #972 settled on.

Acceptance Criteria

• The two readings are resolved with evidence, and the outcome is recorded in the issue before the code change lands: either a public quantized Jamba checkpoint with num_experts > 1 is shown to mis-load today, or it is shown that none exists and the variant is dead.
• If the path is implemented: a Jamba checkpoint with num_experts > 1 and quantized switch_mlp planes loads through SwitchLinear::Quantized and generates coherent output, verified against an mlx-lm reference on the same weights and the same raw prompt (--no-chat-template on both sides).
• If the path is removed instead: the Quantized variant and its gather_qmm arm are gone, and the file records why.
• #[allow(dead_code)] is removed from the SwitchLinear enum at src/models/jamba.rs:336 either way, and the build stays clean under -D warnings.
• The declared group_size / bits are bounds-checked before they are stored, so a hostile "bits": 0 or "group_size": 0 fails at load rather than aborting at the first forward pass. Use switch_layers::validate_expert_quantization_params and the shared HOSTILE_QUANT_PARAMS table at src/models/switch_layers.rs:217.
• A per-expert (unstacked) quantized layout stacks its .scales and .biases rather than dropping them, or the loader rejects that layout with a clear message instead of silently building a Regular plane from packed weights.
• Regression tests drive the values through JambaModel::from_weights rather than a helper in isolation, assert on the returned Result, and never run a forward pass. Each hostile case is paired with a positive control that loads first.
• models/jamba-v0.1-4bit continues to load and generate identically to main, confirming the dense path is untouched.
• cargo clippy --release --lib --tests --features metal,accelerate -- -D warnings is clean and cargo fmt --all -- --check passes.

Refs

Follow-up to #958 (PR #972, merged as db6cce9dc). #958 listed jamba.rs under Scope; PR #972 found no expert loader to guard because the quantized construction site does not exist. Same uncatchable-abort class as #929, #822. Sibling follow-ups from the same trace: #973 and #975.