fix(models): bound quantization params in the family-local MoE expert loaders #958

Problem / Background

#929 (PR #956) bounded the declared quantization params at the two shared loaders, reconcile_quantization_layout (every UnifiedLinear / UnifiedEmbedding) and SwitchLinear::from_stacked_parts (the shared MoE loader). Eighteen other gather_qmm call sites under src/models/ are behind a family-local quantized expert type that reaches MLX without passing through either.

PR #956 originally argued these stayed covered in practice because every such family loads embed_tokens through UnifiedEmbedding::from_weights before it builds an expert plane. That argument does not hold, and the PR body was corrected before merge. It fails three ways:

1. UnifiedEmbedding::from_weights only enters the quantized branch when {prefix}.scales exists. A checkpoint with a float embed_tokens and float attention but quantized expert planes never touches a guarded loader. That is not exotic, it is the canonical gpt-oss shape (MXFP4 experts, bf16 everything else).
2. src/models/gpt_oss.rs resolves quantization params per weight prefix (params_for), with no bounds check. A "bits": 0 scoped to the expert prefix leaves the global defaults valid, so every guarded loader sees a good pair while ExpertLinear::Quantized stores the 0 and divides by it.
3. src/models/mamba.rs and src/models/mamba2.rs construct QuantizedEmbedding::new directly rather than through from_weights_with_mode, so they bypass the reconciler entirely and the embedding claim is untrue for them even with a quantized table.

The consequence is the defect class #929 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. quantized_matmul and gather_qmm cross the cxx bridge as UniquePtr<MlxArray> rather than Result, so the C++ throw is an uncatchable std::terminate, not something catch_unwind contains.

Scope

Families with a local quantized expert type that stores args.group_size() / args.bits() verbatim: qwen3_moe.rs (and qwen3_vl_moe.rs, which imports its types), deepseek.rs, deepseek_v2.rs, deepseek_v3.rs, deepseek_v32.rs, glm4_moe.rs, glm4_moe_lite.rs, ernie4_5_moe.rs, exaone_moe.rs, hunyuan_moe.rs, jamba.rs, llama4.rs, nemotron_h.rs, qwen3_next.rs (and qwen3_5.rs, which reuses them), step3p5.rs, and gpt_oss.rs. Plus mamba.rs and mamba2.rs for the direct QuantizedEmbedding::new path.

Proposed Solution

Call mlxcel_core::layers::validate_quantization_params(group_size, bits) on the values each local constructor is about to store, the way SwitchLinear::from_stacked_parts now does. Several of those constructors return Self rather than Result<Self, String>, so the change is not purely additive for every family and the signature cascade is the bulk of the work; check whether a single call at each family's load boundary covers the same ground more cheaply.

gpt_oss needs the check inside params_for rather than once per model, because its params are per prefix. mamba / mamba2 need it where they build the QuantizedEmbedding by hand, or the hand-built construction should be routed through from_weights_with_mode instead.

Consider also whether QuantizedEmbedding::new and the equivalent by-hand constructors should be made fallible, so this cannot be reintroduced by the next family that builds one directly.

Acceptance Criteria

• A config.json declaring "bits": 0, "bits": 33 or "group_size": 0 fails at load on every family listed under Scope, including one whose only quantized tensors are the expert planes. The nemotron_h test is exactly that shape: float backbone.embeddings and float lm_head, quantized expert planes only. Jamba is the one Scope entry with nothing to guard, because it never constructs its quantized variant; filed as fix(models): Jamba never builds a quantized MoE expert plane #974.
• The gpt-oss per-prefix case is covered, in both directions. Note the premise needed correcting: the experts read Quantization::defaults() rather than quant_for, so the live gap is a hostile top-level pair behind valid per-prefix overrides (the shape every real gpt-oss checkpoint ships). ModelArgs::validate_quantization walks the top-level pair and every override in one pass.
• mamba and mamba2 no longer reach QuantizedEmbedding without a bounds check. QuantizedEmbedding::new and QuantizedMultiLinear::new are now fallible, which answers the design question below rather than patching the two call sites.
• Regression tests drive the bad values through the real load path, not a pure helper: each family's own loader, and for mamba / mamba2 / nemotron_h the whole model constructor. The SIGABRT expectation does not hold and is better not held: the tests assert on the load Result and never run a forward pass, so a lost guard fails cleanly at the assertion. SIGABRT is the production outcome being prevented. Verified empirically by disabling one guard and confirming the test fails on the intended assertion.
• Each hostile case is paired with a positive control that loads first, plus a non-quantized control wherever the loader has a float branch, so a guard that gated the bf16 path also fails.
• cargo clippy --release --lib --tests --features metal,accelerate clean and cargo fmt --all -- --check passes. CI green on PR fix(models): bound quantization params in the family-local MoE expert loaders #972 (cargo-fmt, cargo-deny, cross-repo refs, 2-host logical).
• Validated against at least one real quantized MoE checkpoint to confirm normal loading is unaffected. Not run in the implementation session (no release build there); left for the maintainer's real-checkpoint pass. Suggested coverage: gpt-oss-20b-mxfp4 (per-prefix), deepseek-v2-lite-4bit (the shared QuantizedMultiLinear MLA path), and one of qwen3-30b-a3b-4bit / mixtral-8x7b-4bit / nemotron-h-30b-4bit.

Refs

Follow-up to #929 (PR #956). Same abort class as #822, #427, #382, and docs/adr/0003-release-panic-unwind-with-core-thread-abort.md.

Implemented by PR #972. Adjacent defects found during the trace and filed separately rather than folded in: #973 (the mode string is bounded nowhere), #974 (Jamba never builds a quantized expert plane), #975 (deepseek::ModelArgs drops an inherited quantization block), #976 (mamba / mamba2 treat a scales-without-biases embedding as non-quantized).