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

[inureyes]

Summary

#929 (PR #956) bounded the declared quantization pair at the two shared loaders, reconcile_quantization_layout and SwitchLinear::from_stacked_parts. The families that keep their own quantized expert type never reach either, so a config-derived "bits": 0 was still stored verbatim and handed to gather_qmm, which reaches the same w.shape(-1) * 32 / bits division inside extract_quantized_matmul_dims. gather_qmm, quantized_matmul and dequantize cross the cxx bridge as UniquePtr<MlxArray> rather than Result, so the C++ throw is an uncatchable std::terminate at the first routed forward pass rather than a load error.

Every one of those call sites is now bounded at load.

Approach: why the storage point rather than the load boundary

The issue asked me to evaluate whether a single call at each family's load boundary covers the same ground more cheaply. I traced it, and it does not. The bound goes at the point the pair is stored, and that choice is forced rather than stylistic:

• The load boundary is duplicated in at least a dozen places. The pipeline stage executors for glm4, glm4_moe_lite, deepseek_v3 and llama4 call TransformerBlock::from_weights / DecoderLayer::from_weights directly in a loop, never entering the family's Model::from_weights. DeepSeekV32StageModel::from_filtered_weights, Qwen3NextStageModel::from_filtered_weights, Qwen35StageModel::from_filtered_weights and GptOssStageModel::from_filtered_weights are second in-file boundaries that each re-derive group_size / bits independently. The VLM text wrappers glm4v_moe, ernie4_5_moe_vl and loading/vlm_step3p7.rs build expert planes through their own chains, and glm_moe_dsa::to_dsv32_args feeds a GLM config into the deepseek_v32 loader.
• Three families build the quantized variant as a bare struct literal from a different module, which no boundary check can reach at all: qwen3_vl_moe.rs builds qwen3_moe::SwitchLinear::Quantized directly, ernie4_5_moe_vl.rs builds ernie4_5_moe::SwitchLinear::Quantized directly, and nemotron_h.rs has no named constructor for QuantizedSwitchLinear at all.
• The pair arrives from more than one config type per expert enum, so bounding the producer would not have covered it either. qwen3_next::SwitchLinear is fed by Qwen3NextConfig, by Qwen35Config through to_qwen3next_config, and by a config synthesized from Qwen3OmniSpeechConfig in audio/qwen3_omni_moe/talker.rs.
• The cost is roughly equal anyway. There are about as many load boundaries as storage points, so the boundary version buys nothing in call-site count while requiring a whole-program reachability argument that a new pub caller silently invalidates.

Two of those config types also default to group_size: 0, bits: 0 rather than 64/4 when no quantization block was inherited (Qwen3VLMoeConfig and Ernie45MoeVlTextConfig), which is the hostile pair itself, latent behind the fact that such a checkpoint currently has no .scales.

New switch_layers::validate_expert_quantization_params(prefix, group_size, bits) delegates to mlxcel_core::layers::validate_quantization_params and prefixes the tensor name, so every family reports the same bound with the offending plane named. Only the quantized arms are gated; a bf16 expert plane carries no packing and stays loadable at any declared pair.

Should QuantizedEmbedding::new be fallible? Yes

QuantizedEmbedding::new and QuantizedMultiLinear::new now return Result<Self, String>. They are the only ways to build a quantized layer without going through from_weights_with_mode, and therefore without the reconciler and the bound it carries. A caller that hand-builds one is exactly the caller whose declared pair has never been looked at.

That is also why mamba / mamba2 needed only a ? rather than being routed through from_weights_with_mode. Routing them would not have been behavior-preserving: both take the table out of the map by remove, under either the backbone.embeddings or the model.embed_tokens spelling, so neither can address a single-prefix map loader; from_weights_with_mode additionally reconciles against the tensor shapes and stores the reconciled pair rather than the declared one, treats .biases as optional where mamba requires both, and would need the .or_else() prefix fallback re-expressed at the call site. Making new fallible fixes both families with one line each and closes the door for the next family, which routing them would not have done. QuantizedMultiLinear::new had no callers at all, so making it fallible was free and pre-emptive.

Verified inventory versus the issue's scope list

The issue's list is accurate except for one entry, and it misses four call sites.

Confirmed as listed (17): qwen3_moe, qwen3_vl_moe, deepseek, deepseek_v2, deepseek_v3, deepseek_v32, glm4_moe, glm4_moe_lite, ernie4_5_moe, exaone_moe, hunyuan_moe, llama4, nemotron_h, qwen3_next (and qwen3_5 through it), step3p5, gpt_oss, plus mamba / mamba2.

Does not match: jamba. Its SwitchLinear enum is #[allow(dead_code)] and the Quantized variant is never constructed anywhere in the tree; only Regular is built, at three sites. There is no constructor to bound and no reachable gather_qmm. Left unchanged and recorded here rather than given a guard with no call site.

Missed by the issue (4), all the same defect and all fixed here:

• mlxcel_core::layers::QuantizedMultiLinear::from_weights is a shared loader in mlxcel-core that stores the declared pair verbatim and hands it to quantized_matmul on every MLA forward, without ever calling the reconciler. It sits in exactly the blind spot from_stacked_parts occupied before fix(core): an unvalidated config.json quantization block reaches MLX and terminates the process #929, and it is reached by DeepSeek V3, DeepSeek V3.2, GLM4 MoE Lite and LongCat Flash NGram through their embed_q / unembed_out. This is the highest-value item in the diff: one fix covers four families that fix(core): reject quantization params MLX would abort on, at load #956 and this issue both missed.
• src/models/ernie4_5_moe_vl.rs has its own load_switch_linear that does not call the ernie4_5_moe constructor the issue lists, so fixing that family alone would have left the VLM sibling open.
• src/models/kimi_linear.rs keeps a private MultiLinear (distinct from the mlxcel-core one) that feeds quantized_matmul directly.
• longcat_flash_ngram::sanitize_weights and youtu_vl_lm_sanitize::sanitize_text_weights derive a pair from kv_lora_rank and a tensor axis during MLA kv_b_proj decomposition and hand it to dequantize. Same uncatchable abort, during weight sanitization rather than at first forward. longcat's sanitizer becomes fallible for the check, a two-caller cascade.

The gpt-oss case is real but not for the reason the issue gives

The issue says a hostile pair scoped to the expert prefix leaves the global defaults valid. The code is the mirror image of that: MLPBlock::from_weights reads Quantization::defaults() for the experts, not quant_for, so a per-prefix override on the expert prefix never reaches them. The gap is the other direction, and it is worse. The canonical gpt-oss-20b-mxfp4 config carries valid per-prefix overrides for model.embed_tokens, the four attention projections and the router, with the mxfp4 expert parameters living in the top-level defaults. So a hostile top-level pair reaches the experts alone while every reconciler-guarded loader in the model sees a good per-prefix pair. Verified against the real config in models/gpt-oss-20b-mxfp4/config.json.

Handled two ways: the authoritative bound at ExpertLinear::from_weights, plus a family-level early diagnostic ModelArgs::validate_quantization that walks the top-level pair and every per-prefix override in one pass at both gpt_oss load boundaries. The test covers both directions.

On scope: the issue's Scope section is incomplete, not wrong

Every family the issue names is real and every one except jamba got a guard. What the Scope section does not capture is that the same defect exists on a second axis it never looks at: the MLA kv_b_proj decomposition. Five sanitizers (DeepSeek V3, DeepSeek V3.2, KimiLinear, LongCat Flash NGram, Youtu-VL) each carry a byte-identical copy of arithmetic that solves a quantization pair from kv_lora_rank and a tensor axis and hands it straight to dequantize. That is the same uncatchable-abort class from the same untrusted config.json, differing only in that the pair is derived rather than declared.

It belongs in this PR rather than a follow-up for three reasons. First, four of the five files were already being modified here for their expert loaders, so splitting would have meant two PRs touching the same functions. Second, the first commit's guard on two of those sites was placed after the division it needed to protect, which is a defect this PR introduced and therefore has to fix. Third, the fix is a single shared helper that all five call, so landing it piecemeal would have left copies diverging mid-series.

What was deliberately NOT folded in is filed rather than left as an observation: #973 (the mode string is bounded nowhere and reaches gather_qmm as an unvalidated string, flagged independently by both reviewers), #974 (Jamba never builds a quantized expert plane at all, which is why the issue named it but there was nothing to guard), #975 (deepseek::ModelArgs silently drops an inherited quantization block), and #976 (mamba / mamba2 treat a scales-without-biases embedding as non-quantized).

Review follow-ups applied

Two review passes over the first two commits found four things, all fixed in the third and fourth commits:

• Three more MLA kv_b_proj sites had the identical defect and were missed. DeepSeekV3Model::sanitize_weights, DeepSeekV32Model::sanitize_weights and KimiLinearModel::sanitize_weights each solve the packed pair from kv_lora_rank and a tensor axis and hand it to dequantize, exactly like the LongCat and Youtu sanitizers the first commit converted. There were five copies of that arithmetic, not two. All five now share infer_mla_quantization_params, and the three new ones cascade fallibility through sanitize_weights_with_args into loading/vlm_kimi_vl.rs, loading/special.rs, models/glm_moe_dsa.rs and the deepseek_v3 and glm_moe_dsa pipeline stage executors.
• The bound was placed after the division it needed to protect. inferred_bits = (packed_in * 32) / kv_lora_rank divides by a config.json field and inferred_gs = kv_lora_rank / num_groups by a checkpoint axis. Rust integer division by zero panics, so a kv_lora_rank of 0 or a zero-length scales axis took the process down before the guard could reject anything, and packed_in * 32 wraps i32 in release. The shared helper now checks each divisor before dividing, evaluates the multiply in i64, rejects a rank-0 tensor instead of indexing out of bounds, and additionally verifies the solved pair against MLX's own predicate, because both divisions truncate and an in-range pair is not yet a describable one.
• FusedQKVLinear::from_weights_separate_with_mode is another unguarded shared loader. It never calls the reconciler, and infer_quantization_bits early-returns the caller's bits unchanged whenever group_size <= 0, so a declared (0, 0) was stored verbatim on the fused QuantizedWeight. Shadowed in every current family by a sibling UnifiedLinear that sees the same pair and rejects it first, so this was not a live hole, but it is the same shape as the one fix(core): an unvalidated config.json quantization block reaches MLX and terminates the process #929 left open and it is one line.
• The tests' SIGABRT claim was false and is corrected. The comment in sixteen files said a regression would take the test binary down with SIGABRT. That is what happens in production; these tests assert on the load result and never run a forward pass, so a regression fails cleanly at the assertion. The wording would have misled whoever next debugs a failure here. The QuantizedEmbedding::new doc made a similar overclaim (the fallible signature does not make the bound impossible to skip, because the fields are pub) and now says so.

Also from review: the MultiLinear enum Used by: list gained LongCat Flash NGram, and FusedQKVLinear::from_weights_separate_with_mode had a stale list naming a Phi3 caller that does not exist while omitting eleven that do. Both corrected against grep.

A finalizer pass then found two more things, fixed in the sixth commit:

• The five MLA sanitizers were covered only by the pure-helper unit test, which contradicts this PR's own stated test philosophy. All five now have a test driving the real sanitize_weights / sanitize_text_weights. LongCat Flash NGram had no #[cfg(test)] block at all before this, and DeepSeek V3's inline module had no quantized-MLA case. One trap is worth recording: unlike the expert-plane tests, which stop at the loader's Result, these drive a sanitizer that actually evaluates the dequantize, so the packed fixture must be UINT32. A float fixture aborts the whole test binary with SIGABRT on the positive control, which is what the first draft of all five did and what running them caught. The reason is now recorded at each fixture.
• docs/adding-models.md never mentioned the convention, so the next family ported with a local quantized expert type would silently reintroduce the bug. It gains a Quantization Parameter Bounds section stating the rule and why the load boundary does not dominate, a table of what already carries the bound so a porter inherits it for free, and a checklist item pointing at it.

Both reviewers independently flagged that quantization.mode is bounded nowhere and reaches gather_qmm as an unvalidated string, which is the same abort class. That is a different config field and a different fix (an allowlist, which already exists at gemma4::SUPPORTED_QUANT_SCHEMES but is wired only into gemma4), so it is filed as a follow-up rather than folded in here.

Test plan

Both acceptance criteria on testing hold, and the second one needs a correction to how the issue phrases it.

Real load path, not a pure helper. Every test calls the family's own loader (SwitchLinear::from_weights, load_switch_linear, ExpertLinear::from_weights, Step3p5SwitchGLU::from_weights, MultiLinear::from_weights) over a synthetic WeightMap carrying an honestly packed expert plane, so the guard is exercised at the point a real checkpoint reaches it. Four go further and drive the whole model constructor: MambaModel::from_weights, Mamba2Model::from_weights and NemotronHModel::from_weights (the last including its C++ model registration), plus the gpt_oss config walk over the real gpt-oss-20b-mxfp4 quantization shape.

Positive control first, always. Each test loads the honest pair before it tries any hostile one, and every test whose loader has a non-quantized branch additionally asserts that a bf16 plane with no .scales still loads under a hostile pair. A guard that rejected everything quantized, or that gated the float path, fails on those two assertions before it reaches the hostile loop.

The issue expects a regression to surface as a SIGABRT from the test binary. It does not, and that is better rather than worse: these tests assert on the load Result and never run a forward pass, so a lost guard fails cleanly at the assertion instead of aborting. SIGABRT is what happens in production, which is the thing being prevented. The doc comments that repeated the issue's phrasing were corrected in the fourth commit, because they would have misled whoever next debugs a failure here.

• cargo test --release --lib --features metal,accelerate quantization_params (19 passed: the 18 new per-family guards plus the existing switch_layers one)
• cargo test --release --lib --features metal,accelerate models::gpt_oss (3 passed, including the per-prefix config walk over the real gpt-oss-20b-mxfp4 quantization shape)
• cargo test --release --lib --features metal,accelerate models::ernie4_5_moe_vl (5), models::mamba (8), models::nemotron_h (19)
• cargo test --release -p mlxcel-core --lib --features metal,accelerate layers::tests (65 passed, including the two tests for the fallible by-hand constructors and the shared MultiLinear loader, plus the MLA divisor and overflow test)
• Affected-family sweep, all green: deepseek (28), deepseek_v2 (1), deepseek_v3 (24), deepseek_v32 (12), glm4_moe (7), glm4_moe_lite (1), ernie4_5_moe (6), exaone_moe (3), hunyuan_moe (1), jamba (4), llama4 (1), qwen3_next (8), qwen3_5 (5), step3p5 (12), qwen3_moe (1), qwen3_vl_moe (2), kimi_linear (1), longcat_flash_ngram, youtu_vl_lm (3)
• MoE regression sweep for the families this does NOT change, to confirm the shared-helper additions are inert for them: switch_layers (22), mixtral, qwen2_moe, olmoe, phimoe (5), dots1 (9), gemma4 (90), minimax (15), bailing_moe (42)
• cargo test --release --lib --features metal,accelerate kv_b_proj (5 passed: one per MLA sanitizer, each driving the real sanitize_weights / sanitize_text_weights rather than the shared helper)
• cargo test --release --lib --features metal,accelerate models:: (774 passed, 0 failed), loading:: (220), distributed:: (1268)
• Anti-vacuity check, run empirically rather than argued: the guard in qwen3_moe::SwitchLinear::from_weights was temporarily replaced with let _ = validate_expert_quantization_params(..) and the suite re-run. The test failed on exactly the intended assertion ((group_size 64, bits 0) must be refused at load, not stored for gather_qmm at qwen3_moe_tests.rs:86) rather than passing for some other reason, and passed again once the guard was restored. No hostile pair in HOSTILE_QUANT_PARAMS is rejected by any pre-existing check, so these tests detect a regression rather than merely accompanying one.
• cargo clippy --release --lib --tests --features metal,accelerate clean
• cargo fmt --all -- --check
• Real quantized checkpoint load, to confirm normal loading is unaffected. Not run here: the release binaries are not built in this session and the reference MoE checkpoints are large. Left for the follow-up validation pass, which should cover gpt-oss-20b-mxfp4 (the per-prefix case), deepseek-v2-lite-4bit (the shared QuantizedMultiLinear MLA path), and one of qwen3-30b-a3b-4bit / mixtral-8x7b-4bit / nemotron-h-30b-4bit.

The pre-existing multimodal::host_preprocessor::tests::xla_loader_keeps_text_and_unqualified_vlm_image_capability_false failure on main is tracked as #966 and is unrelated to this change.

Notes

Two adjacent findings turned up during the trace that are not fixed here, because both are separate defects rather than this one. Both are filed rather than left as observations, as #975 and #976 respectively:

• deepseek::ModelArgs has no quantization field, only flat top-level group_size / bits. The three DeepSeek-OCR loaders in loading/vlm_deepseekocr.rs inject the root quantization object into language_config before deserializing, so serde drops it and group_size() / bits() silently fall back to 64/4 regardless of what the checkpoint declares. DeepSeekModel::load has the same shape. The contrast is vlm_deepseek_vl2.rs, which does the identical inheritance into deepseek_v2::ModelArgs and works because that config has the nested field.
• mamba / mamba2 treat an embedding table with .scales but no .biases as non-quantized and fall through to Embedding::new, which would read packed uint32 data as a float table. Only reachable for a block-float embedding export, which no current mamba checkpoint ships.

Closes #958

[inureyes]

Real-checkpoint validation on Apple M1 Ultra (Metal + Accelerate, release build), covering the acceptance criterion left unticked for this pass.

Every quantized MoE checkpoint available locally loads and generates byte-identically to main with this PR applied, so the new bounds reject nothing legitimate and change no numerics. A/B run at the same prompt and token count against main at 30f1ca82f:

Checkpoint	Path exercised	Result
gpt-oss-20b-mxfp4	per-prefix params	identical to main
qwen3-30b-a3b-4bit	family-local expert plane	identical to main
mixtral-8x7b-4bit	family-local expert plane	identical to main
nemotron-h-30b-4bit	family-local expert plane	identical to main
deepseek-v2-lite-4bit	shared QuantizedMultiLinear MLA	identical to main, see below

One finding, and it is not this PR. deepseek-v2-lite-4bit emits degenerate output (}'_}'_}'_... repeated) with and without the chat template. It does so byte-identically on main, so it is pre-existing and unrelated to this change; the checkpoint simply had not been exercised before. It is filed separately. The comparison mattered here rather than being a formality: the same symptom under a PR that touches the MLA quantization path would otherwise have read as a regression, and merging without the A/B would have let a real defect pass unrecorded.

The kv_b_proj decomposition itself remains unexercised against real weights, since that needs a DeepSeek V3 / V3.2 / KimiLinear / LongCat checkpoint and none is available here. That limit is worth carrying forward rather than treating this pass as complete coverage of the MLA change.