fix(core): bound the declared quantization mode string, not just group_size and bits #973

Background

PR #972 (issue #958, merged as db6cce9dc) bounded the declared quantization parameters at the point they are stored. It added mlxcel_core::layers::validate_quantization_params(group_size, bits) (src/lib/mlxcel-core/src/layers.rs:1292) and called it from the shared loaders plus roughly twenty family-local expert-plane constructors, and it made QuantizedEmbedding::new and QuantizedMultiLinear::new return Result. The defect class it exists to close is this one: 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, gather_qmm, dequantize and quantized_embedding cross the cxx bridge as UniquePtr<MlxArray> rather than Result (src/lib/mlxcel-core/src/lib.rs:1092, :1108, :1121, :1136), so the C++ throw is an uncatchable std::terminate that catch_unwind cannot contain. This issue is the third value in the same quantization block, the mode string, which #972 deliberately left out of scope. The body of PR #972 records that both review passes on it raised the unbounded mode independently, which is why it was filed rather than folded in.

What is unbounded

validate_quantization_params bounds bits to 1..=32 and group_size to 1..=MAX_QUANT_GROUP_SIZE (1 << 20, src/lib/mlxcel-core/src/layers.rs:1419). The mode string is bounded in exactly one place in the tree, and that place is family-local: SUPPORTED_QUANT_SCHEMES (src/models/gemma4.rs:3741) is ["affine", "mxfp4", "nvfp4", "mxfp8"], checked by validate_quantization_scheme (src/models/gemma4.rs:3768) and wired into gemma4's own two loaders only (src/models/gemma4.rs:3868 and src/models/gemma4.rs:4134). That function also inspects quant_method, which is the #467 external-format guard and a superset of what a core mode check needs. Nothing in mlxcel-core bounds the mode at all.

What MLX accepts

MLX accepts exactly four modes. string_to_quantization_mode (ml-explore/mlx mlx/primitives.cpp) matches affine, mxfp4, mxfp8 and nvfp4, and on anything else throws std::invalid_argument with the message [<tag>] Invalid quantization mode '<mode>'.. quantized_matmul, gather_qmm, dequantize and quantize all reach it, most through validate_mode_with_type in mlx/ops.cpp. This is a closed four-value enum in C++, not a numeric range.

What happens today with an unrecognized mode

It is neither rejected nor silently defaulted, and the failure is an uncatchable abort rather than a clean error or a wrong-but-finite result. Traced end to end on current main:

1. Outside gemma4, no load-time check looks at the string at all.
2. On the UnifiedLinear / UnifiedEmbedding path it reaches reconcile_quantization_layout (src/lib/mlxcel-core/src/layers.rs:1467), whose only use of mode is the two-way branch at src/lib/mlxcel-core/src/layers.rs:1502, if mode == "affine" { ... } else { ... }. Any unrecognized string therefore takes the block-float branch. On a genuine affine 4-bit tensor that branch computes in_features = packed_in * (32 / bits) and re-derives group_size = in_features / num_groups, which lands back on the declared value, so reconciled is false and not even the divergence warning on the mlxcel::quant tracing target fires. The bad string is then stored verbatim on QuantizedWeight.mode (src/lib/mlxcel-core/src/layers.rs:41) or QuantizedEmbedding.mode (:230).
3. The stored string is passed verbatim into quantized_matmul / gather_qmm / dequantize on every forward pass (roughly two dozen sites in layers.rs alone, for example :1869, :2492, :2821, :3016).
4. MLX throws inside string_to_quantization_mode. That throw crosses a UniquePtr<MlxArray> bridge function, so it becomes std::terminate and the process dies with SIGABRT at the first forward pass. The model loaded without a warning.

The neighbouring failure modes in validate_mode_with_type are the same abort class and worth designing against together: a mode that is valid but wrong for the tensor throws Biases must be provided for affine quantization (declared affine, no .biases plane), Scale type must be uint8 but received type <t> (declared block-float on affine scales), or Biases must be null for quantization mode '<mode>' (declared block-float with a bias plane present). There is no silent-wrong-output path here; every mode mismatch aborts.

Where the declared string actually reaches MLX today

This is narrower than it looks, and the narrowing is load-bearing for how the fix should be scoped.

• src/models/gpt_oss.rs is the ONLY family in the tree that reads quantization.mode out of config.json and threads it into a loader. Quantization::params_for (src/models/gpt_oss.rs:101) reads the per-prefix mode with .unwrap_or("affine") and no validation; Quantization::defaults (:151) does the same for the top-level block. Both return an owned String.
• That string reaches MLX two ways. It is stored on ExpertLinear::Quantized { mode: String } (src/models/gpt_oss.rs:335, assigned at :444) and handed to mlxcel_core::gather_qmm as the last argument in ExpertLinear::forward (:360), and it is passed to UnifiedLinear::from_weights_with_mode at seven sites (:634 router, :758 / :763 / :768 / :773 attention projections, :1056 and :1219 lm_head).
• Quantization::validate (src/models/gpt_oss.rs:133), the early diagnostic PR fix(models): bound quantization params in the family-local MoE expert loaders #972 added, destructures as let (gs, bits, _) = ... and discards the mode on both the top-level pair and every per-prefix override.
• Every other quantizing family never sees the declared string. UnifiedLinear::from_weights (src/lib/mlxcel-core/src/layers.rs:1755) and QuantizedEmbedding::from_weights (:281) derive the mode from bias presence via infer_quantization_mode (:1430), which returns one of the four valid &'static str values by construction, and the shared switch_layers::SwitchLinear::from_weights (src/models/switch_layers.rs:288) hardcodes "affine". mlxcel_core::layers::QuantizedWeight::new_with_mode, QuantizedEmbedding::from_weights_with_mode and switch_layers::SwitchLinear::from_weights_with_mode are all pub and all take an unbounded &str, but no current caller passes a config-derived value into them.
• models::config::QuantizationArgs (src/models/config.rs:244) already carries an Option<String> mode with get_mode() returning the declared string verbatim and defaulting to "affine" (:268). It has zero production callers today. It is the obvious thing for the next family to reach for, and it would inherit the hole silently.

So the live reachable failure is gpt-oss only, and the reason to put the check in the core rather than in gpt_oss.rs is the second bullet list above: four pub entry points accept an unbounded mode string and one unused shared config helper hands one out.

Correction to the premise this issue was filed with

The original framing said a bad mode scoped to the expert prefix would leave the top-level defaults valid, mirroring what #958 said about bits. The code is the mirror image, exactly as PR #972 had to correct for bits: MLPBlock::from_weights reads args.quantization.defaults() for the experts (src/models/gpt_oss.rs:637), not quant_for, so a per-prefix override on the expert prefix never reaches them. The live gap is a hostile TOP-LEVEL mode reaching the experts while every guarded sibling loader sees a good per-prefix mode. That is the shape every real gpt-oss checkpoint ships. Verified against models/gpt-oss-20b-mxfp4/config.json, which declares top-level {"group_size": 32, "bits": 4, "mode": "mxfp4"} plus 122 per-prefix override objects, of which 98 are {"group_size": 64, "bits": 4, "mode": "affine"} (including model.embed_tokens) and 24 are {"group_size": 64, "bits": 8} with no mode key at all. Note that last group: params_for defaults a mode-less override to "affine" rather than inheriting the top-level "mxfp4", which happens to be right for those tensors but is an implicit behavior any fix must not disturb.

Related case: affine mode with no biases

ExpertLinear::from_weights (src/models/gpt_oss.rs:397 onward) enters the quantized branch on .scales presence alone and makes quant_biases an Option that becomes a null pointer when absent. An expert plane that ships .scales and no .biases under an explicitly declared affine mode is therefore constructed without complaint and aborts in gather_qmm with Biases must be provided for affine quantization. The shared loaders avoid this by ignoring the declaration and inferring from bias presence (infer_quantization_mode, src/lib/mlxcel-core/src/layers.rs:1430); the family-local path does not.

Proposed Solution

An allowlist is the correct shape here, and the reasoning that made validate_quantization_params a bounds check rather than an allowlist does not transfer. That reasoning (recorded in the doc comment at src/lib/mlxcel-core/src/layers.rs:1274 onward) is that mlxcel deliberately re-derives an effective bit width from tensor shapes, so an allowlist of {2,3,4,5,6,8} would reject legitimate mixed-precision exports. There is no analogous re-derivation from the declared mode string: infer_quantization_mode derives a mode from bias presence, independently of what was declared, and the accepted set is a closed four-value C++ enum with no interior values. Mirroring MLX's own parser is the check.

1. Hoist the list into mlxcel_core::layers next to validate_quantization_params, as something like validate_quantization_mode(mode: &str) -> Result<(), String>, with an error naming the four accepted values and pointing at the offending prefix where one is available. Decide explicitly whether to trim and lowercase before comparing (gemma4's check_obj does, at src/models/gemma4.rs:3809) and whether an empty string means "not declared" (gemma4 treats it as acceptable).
2. Call it from reconcile_quantization_layout right after validate_quantization_params, which covers every UnifiedLinear and UnifiedEmbedding in one edit and makes the mode == "affine" branch at src/lib/mlxcel-core/src/layers.rs:1502 total rather than a silent fallthrough.
3. Call it from the four pub entry points that accept an unbounded mode and do not reach the reconciler: QuantizedWeight::new_with_mode, switch_layers::SwitchLinear::from_weights_with_mode, ExpertLinear::from_weights in src/models/gpt_oss.rs, and models::config::QuantizationArgs::get_mode (or make that one return Result, since it has no callers to break).
4. In src/models/gpt_oss.rs, check inside params_for / defaults or extend Quantization::validate to stop discarding the mode, so a hostile per-prefix mode is refused during config validation the way a hostile per-prefix pair now is.
5. Make src/models/gemma4.rs delegate its mode comparison to the shared function so the two lists cannot drift. gemma4 keeps its own validate_quantization_scheme for the quant_method and ModelOpt-NVFP4 policy, which is genuinely gemma4-shaped and covers fix: externally quantized 4-bit Gemma 4 (OptiQ) decodes to degenerate output that is trapped entirely in reasoning_content #467.
6. Decide whether the affine-without-biases case is rejected at load or repaired. Rejecting is the smaller change and matches the abort-at-load intent; repairing means routing the family-local constructor through infer_quantization_mode the way QuantizedEmbedding::from_weights does, which would additionally fix a block-float export mislabelled affine in config.json.

Inspection and reproduction

Confirm the current state without building anything:

# The only family that reads the declared mode string out of config.json.
grep -n 'get("mode")' src/models/gpt_oss.rs
# Everything else derives it from bias presence instead.
grep -rn "infer_quantization_mode" src/lib/mlxcel-core/src/layers.rs
# The only bound that exists, and it is gemma4-local.
grep -n "SUPPORTED_QUANT_SCHEMES\|fn validate_quantization_scheme" src/models/gemma4.rs
# MLX's own accepted set (present after any release build).
grep -n -A 18 "QuantizationMode string_to_quantization_mode" \
  target/release/build/mlxcel-core-*/out/build/_deps/mlx-src/mlx/primitives.cpp

Reproduce the abort. No local checkpoint declares an unsupported mode, so it has to be synthesized. Symlink rather than copy: models/gpt-oss-20b-mxfp4 is 10 GB.

mkdir -p /tmp/gpt-oss-badmode
for f in /Volumes/Monolith/Development/rust/mlxcel/models/gpt-oss-20b-mxfp4/*; do ln -sf "$f" /tmp/gpt-oss-badmode/; done
rm /tmp/gpt-oss-badmode/config.json
python3 - <<'PY'
import json
src = "/Volumes/Monolith/Development/rust/mlxcel/models/gpt-oss-20b-mxfp4/config.json"
c = json.load(open(src))
c["quantization"]["mode"] = "gptq"          # top-level: the pair the experts read
json.dump(c, open("/tmp/gpt-oss-badmode/config.json", "w"), indent=2)
PY
DEVELOPER_DIR=/Applications/Xcode-26.6.0.app/Contents/Developer \
  cargo build --release --features metal,accelerate
./target/release/mlxcel generate -m /tmp/gpt-oss-badmode -p "Hello" -n 8

Expected today: the model loads, then the process dies with SIGABRT and libc++abi: terminating due to uncaught exception of type std::invalid_argument: [gather_qmm] Invalid quantization mode 'gptq'. on the first forward pass. Expected after the fix: a load error naming the four supported modes. Repeat with the mode moved onto a per-prefix override (c["quantization"]["model.layers.0.self_attn.q_proj"]["mode"] = "gptq") to exercise the UnifiedLinear surface, which aborts in quantized_matmul instead.

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. gpt-oss-20b-mxfp4 is the direct positive control for this issue, since it is the only local checkpoint whose declared mode actually reaches a kernel, and it exercises both a valid non-affine top-level mode and valid per-prefix affine overrides in one load. minicpm-v-4.6-mxfp4 (top-level mode: mxfp4) and gemma-4-31b-it-nvfp4 (ModelOpt NVFP4, quant_method: modelopt) are the other two non-affine local checkpoints and are worth including so the new check does not regress the block-float and #467 paths. qwen3.5-0.8b-optiq-4bit is present but declares mode: affine, so it is not a hostile-mode case.

Not possible here. deepseek-v2-lite-4bit currently produces degenerate output for an unrelated reason tracked as #991, so it is not a usable positive control until that is fixed. No local checkpoint declares an unsupported mode, so the hostile case is synthetic in every form (edited config or unit test) and there is no real-world artifact to validate the rejection message against.

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 regression tests 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 across sixteen files.

Acceptance Criteria

• A config.json declaring an unsupported quantization.mode (for example "optiq", "gptq", "int4", or a string that is empty after trimming) fails at load with a message naming the four supported modes, on every quantizing family rather than gemma4 alone.
• The check is an allowlist mirroring MLX's string_to_quantization_mode, and the decision about trimming, case folding, and empty-string handling is made explicitly and matches whatever src/models/gemma4.rs does, so the two cannot disagree on the same config.
• The gpt-oss per-prefix case is covered in both directions: a hostile top-level mode is refused even when every per-prefix override is valid (the live gap, because the experts read Quantization::defaults), and a hostile per-prefix mode is refused even when the top-level mode is valid.
• affine, mxfp4, mxfp8 and nvfp4 still load, including the mixed per-prefix form models/gpt-oss-20b-mxfp4 ships (mxfp4 experts on the top-level defaults, affine overrides for the embedding, attention and router) and the mode-less override form in the same file, which must keep resolving to affine.
• The four pub entry points that accept an unbounded mode string and never reach reconcile_quantization_layout are covered: QuantizedWeight::new_with_mode, QuantizedEmbedding::from_weights_with_mode, switch_layers::SwitchLinear::from_weights_with_mode, and models::config::QuantizationArgs::get_mode.
• An expert plane carrying .scales with no .biases under a declared affine mode either fails at load or is loaded under the inferred block-float mode, rather than reaching gather_qmm and aborting.
• src/models/gemma4.rs no longer owns a private copy of the mode list; SUPPORTED_QUANT_SCHEMES is gone or delegates, and gemma4 keeps only its quant_method / ModelOpt-NVFP4 policy.
• Regression tests drive the bad mode through a real load path (a family loader or model constructor), not the helper in isolation, and assert on the returned Result without running a forward pass. Follow the HOSTILE_QUANT_PARAMS pattern at src/models/switch_layers.rs:217 with an equivalent hostile-mode table.
• Each hostile case is paired with a positive control that loads first, so a check that passes by rejecting everything quantized also fails.
• cargo clippy --release --lib --tests --features metal,accelerate -- -D warnings is clean and cargo fmt --all -- --check passes.
• Validated against models/gpt-oss-20b-mxfp4 (the only local checkpoint whose declared mode reaches a kernel), plus models/minicpm-v-4.6-mxfp4 and models/gemma-4-31b-it-nvfp4 for the block-float and fix: externally quantized 4-bit Gemma 4 (OptiQ) decodes to degenerate output that is trapped entirely in reasoning_content #467 paths, and one affine 4-bit checkpoint, confirming normal loading and output are unchanged.

Refs

Follow-up to #958 (PR #972, merged as db6cce9dc), which bounded group_size and bits and left mode unbounded. Same uncatchable-abort class as #929, #822, #427, #382, and docs/adr/0003-release-panic-unwind-with-core-thread-abort.md. The existing gemma4 allowlist was added for #467. Sibling follow-ups from the same trace: #974 and #975.