fix(core): reject quantization params MLX would abort on, at load

[inureyes]

Summary

An unvalidated quantization block in config.json reached MLX and terminated the process. reconcile_quantization_layout returned a non-positive group_size / bits pair unchanged on a trust-the-caller basis, and MLX reconstructs a quantized matrix's unpacked width as w.shape(-1) * 32 / bits in validate_quantized_input, so a declared "bits": 0 is a division by zero and anything above 32 collapses the quotient to something no real scales width can match. quantized_matmul, gather_qmm and dequantize all cross the cxx bridge as UniquePtr<MlxArray> rather than Result, so the C++ throw is an uncatchable std::terminate at the first forward pass rather than a load error: a repository with real .scales and "bits": 0 loaded cleanly and took mlxcel-server down on the first request that touched a quantized projection. The second half of the issue is the same shape: validate_embedding_table skipped the width check for a quantized table, so a table honestly packed for a different model width kept exactly the right row count, passed load, and then fed a wrong-width hidden state into fast::layer_norm / fast::rms_norm.

Both are now rejected at load, on every quantizing family, dense and MoE.

What changed

src/lib/mlxcel-core/src/layers.rs

• New validate_quantization_params(group_size, bits): rejects bits outside 1..=32 and group_size < 1, with a message naming the offending field and value. It is a bounds check rather than an allowlist of the widths MLX supports, because mlxcel deliberately re-derives bit width from tensor shapes when the declared value disagrees, and an allowlist of {2,3,4,5,6,8} would reject the mixed-precision exports that behavior serves (the Qwen3.5/3.6 MoE per-path overrides, the minicpm-v mxfp4 group-size case).
• reconcile_quantization_layout calls it before anything else. The old early return folded "degenerate shapes" and "degenerate params" into one condition, which is what let an empty shape launder a hostile pair; the shape case stays permissive, the param case does not.
• New validate_quantized_packing plus QuantizedTensorShapes, hoisted verbatim (messages included) out of bailing_moe.rs. It reconstructs the input width the way MLX does, as scales.shape(-1) * group_size using the reconciled group size, and checks .biases shape equals .scales shape. Works for a 2-D projection and a 3-D stacked expert plane alike.
• New infer_quantization_mode(has_biases, group_size, bits), deduplicating the identical mode-inference block that sat in UnifiedLinear::from_weights, QuantizedEmbedding::from_weights and bailing_moe::quant_mode.
• New UnifiedEmbedding::quantized() accessor, so a load-time validator can reach the scales shape and the reconciled params that is_quantized() alone cannot give it.

src/models/switch_layers.rs (the scope correction the issue records)

SwitchLinear::from_stacked_parts never calls the reconciler, re-derives bits itself and accepts the result only in 2..=8, and never infers group_size at all, so a declared 0 zeroed the denominator, skipped inference entirely, and was stored on SwitchLinear::Quantized and handed to gather_qmm, which reaches the same /bits line through the same extract_quantized_matmul_dims. A reconciler-only guard would have left every quantized MoE checkpoint exposed. It is now fallible, takes the prefix so the error names the tensor, bounds the declared pair itself, and additionally refuses a stored triple that fails MLX's own predicate from validate_quantized_input, evaluated in i64 so the multiply cannot wrap on a large expert plane. That second check matters because the existing fallback to the declared bits whenever the shapes solve outside 2..=8 guaranteed a throw. The mixed-precision inference the fallback exists for is untouched (a genuinely 8-bit plane under a 4-bit config default still reconciles), degenerate shapes evaluate to 0 == 0 and stay permissive, and a non-quantized expert plane is not gated on quantization params at all.

The families that carry a local quantized expert type are unchanged and are not in this diff, and they are not covered by this PR. An earlier version of this body claimed they were, on the grounds that each loads embed_tokens through UnifiedEmbedding::from_weights first. That is wrong and has been retracted: UnifiedEmbedding::from_weights only reaches the reconciler when .scales exists, so a float embedding with quantized expert planes (the canonical gpt-oss shape) never touches a guarded loader; gpt_oss resolves quantization params per weight prefix, so a hostile pair scoped to the experts leaves every guarded loader with a valid one; and mamba / mamba2 build QuantizedEmbedding by hand, bypassing the reconciler entirely. Those eighteen call sites are filed as #958.

src/models/gpt2.rs

validate_embedding_table no longer skips the width check for a quantized table. The dequantized width is recoverable at load, so the quantized case goes through the shared validate_quantized_packing and the four families behind that helper (BailingMoe, Gpt2, GptBigCode, GptNeoX) get the check instead of one family owning a private copy.

src/models/bailing_moe.rs, src/models/gpt_neox.rs, src/models/helium.rs

The three duplicated per-family bounds checks are folded into the shared guard and kept only as delegating early diagnostics that name the family and fire during config validation, before any tensor is touched. bailing_moe::validate_quantized_packing becomes a thin WeightMap adapter over the shared function. The tree-wide grep for bits < 1 | bits > 32 | group_size < 1 now matches only layers.rs, down from three model files plus the shared crate.

Notes

bits == 0 is C++ undefined behavior and the targets diverge: on AArch64 the divide returns 0 so std::invalid_argument fires, while on x86-64 the hardware raises SIGFPE and kills the process before any exception exists. Only the AArch64 outcome is reachable from this machine; the x86-64 path is recorded in the doc comment rather than tested.

The regression tests deliberately drive the bad values through the real load path rather than the pure shape helpers, so a regression aborts the test binary with SIGABRT instead of failing cleanly. Every hostile case is paired with a positive control, so no guard can pass by rejecting everything quantized.

Review follow-ups applied

A review of the first commit found two correctness holes in the guards this PR adds, both fixed in the second commit:

• validate_quantized_packing compared only the scales side against the config width, which is not MLX's test. MLX compares the scales side against w.shape(-1) * 32 / bits, and the block-float reconciler keeps the declared group_size in both of its fallbacks, so a triple could satisfy the config check while the packed weight described a different width and MLX still threw. Both sides are now compared, in i64.
• group_size was bounded below but not above. A declared i32::MAX survived the block-float fallback and reached quantized_matmul, where MLX multiplies it by the scales width in C++ int: signed overflow, and therefore undefined behavior reached before the shape check that would have thrown. It is now bounded at 2^20, an overflow bound rather than an allowlist, with the arithmetic that justifies it recorded on the constant.

Also from that review: the i32 num_groups * group_size in from_stacked_parts moved to i64 (it wrapped in release and panicked in an overflow-checked build), a zero or absent last axis is now rejected instead of satisfying 0 == 0 and aborting later, the MoE path gained the biases-shape check that is a second validate_quantized_input throw, and several documentation and message-wording items. The weight dtype condition from validate_quantized_input is deliberately still unchecked, with the reason recorded in the doc comment: the in-tree fixtures build packed weights as f32.

Test plan

• cargo test --release --lib --features metal,accelerate models::switch_layers (21 passed, including the two new MoE guards over both the pre-stacked and per-expert layouts)
• cargo test --release --lib --features metal,accelerate models::gpt2 (27 passed, including the quantized position-table width guard and the shared-guard load rejection)
• cargo test --release --lib --features metal,accelerate models::bailing_moe (42 passed, unchanged messages after the hoist)
• cargo test --release --lib --features metal,accelerate models::gpt_neox (32), models::helium (27), models::gpt_bigcode (20)
• cargo test --release -p mlxcel-core --lib --features metal,accelerate layers::tests (61 passed, including the four new quantization tests)
• MoE regression sweep: models::{mixtral,qwen2_moe,olmoe,phimoe,deepseek,dots1,gemma4,minimax} all green
• cargo clippy --release --lib --tests --features metal,accelerate clean apart from the two warnings already on main in src/multimodal/host_preprocessor.rs and host_preprocessor_export.rs
• cargo fmt --all -- --check
• cargo test --release -p mlxcel-core --lib layers::tests re-run after the review fixes (62 passed, including the new weight-side packing test and the group_size upper bound), plus the full model sweep above
• Real quantized checkpoint load, to confirm normal loading is unaffected. Not run in this environment: the release binaries are not built here and the MoE reference checkpoint is 33.6 GB. Left for the follow-up validation pass. This matters more than usual for the embedding-table change, which converts one class of currently-loading-but-garbage-decoding tables (an affine table packed at group 32 under a config default of 64) into a load error.

Closes #929