fix(core): an unvalidated config.json quantization block reaches MLX and terminates the process #929

Problem / Background

A quantization block in config.json is read verbatim and threaded into MLX's quantized kernels without any bounds check. A checkpoint that declares "bits": 0 (or any value above 32, or a non-positive group_size) alongside real .scales tensors loads without complaint and then kills the process on the first request that touches a quantized projection. This is not a Rust panic, so it is not containable: per docs/adr/0003-release-panic-unwind-with-core-thread-abort.md the release build unwinds and the worker paths can catch_unwind, but a C++ std::terminate is not a panic, and that ADR already names MLX FFI exceptions as the documented residual abort vector.

This surfaced during the security review of #928 (implementing #836, GPT-NeoX), but the exposure is tree-wide and not specific to that family.

Confirmed chain

1. mlxcel deliberately passes the bad value through. src/lib/mlxcel-core/src/layers.rs:976-984, in reconcile_quantization_layout (declared at line 969):

// Insufficient shape info: trust the caller, mirroring the historical
// early-return in `infer_quantization_bits` for empty / scalar shapes.
if weight_shape.is_empty() || scales_shape.is_empty() || group_size <= 0 || bits <= 0 {
    return Ok(ReconciledQuant { bits, group_size, reconciled: false });
}

A non-positive group_size or bits is returned unchanged on a "trust the caller" basis. A bits value above 32 is not caught here either: the affine branch calls infer_quantization_bits, which re-derives from shapes, and the block-float branch gates on 32 % bits == 0 and otherwise returns the declared pair as-is. The declared value reaches the kernel exactly as written.

2. MLX rejects it by throwing. mlx/ops.cpp:115 in validate_quantized_input (declared at line 83), pinned commit b7c3dd6d27f45b5365b08a840310187dc503f1db:

  if (w.shape(-1) * 32 / bits != scales.shape(-1) * group_size) {
    std::ostringstream msg;
    msg << "[" << tag << "] The shapes of the weight and scales are "
        << "incompatible based on bits and group_size. w.shape() == "
        << w.shape() << " and scales.shape() == " << scales.shape()
        << " with group_size=" << group_size << " and bits=" << bits;
    throw std::invalid_argument(msg.str());
  }

At bits == 0 this is an integer division by zero. At bits > 32 the quotient collapses toward zero and can match no real scales width. A non-positive group_size likewise makes the right-hand side unmatchable. The exception type is std::invalid_argument.

validate_quantized_input is not reached only from quantized_matmul. It is called from extract_quantized_matmul_dims (mlx/ops.cpp:125, the call at line 134), which is the shared entry gate for quantized_matmul (ops.cpp:4628) and gather_qmm (ops.cpp:5417) alike, plus a direct call at ops.cpp:4706. Every quantized matmul flavor in the tree, dense and MoE, divides by bits at that one line.

One platform nuance worth confirming on the CUDA build: division by zero is undefined behavior in C++, and the two targets diverge in practice. On AArch64 the divide instruction returns zero rather than trapping, so the comparison fails and the std::invalid_argument above is what fires. On x86-64 the hardware raises SIGFPE, which kills the process outright before any exception exists. Both outcomes are process death; only the diagnostic differs.

3. The throw is uncatchable. quantized_matmul is declared in the cxx bridge at src/lib/mlxcel-core/src/lib.rs:1108 as returning UniquePtr<MlxArray>, not Result<UniquePtr<MlxArray>>:

unsafe fn quantized_matmul(
    x: &MlxArray,
    w: &MlxArray,
    scales: &MlxArray,
    biases: *const MlxArray,
    transpose: bool,
    group_size: i32,
    bits: i32,
    mode: &str,
) -> UniquePtr<MlxArray>;

cxx generates a noexcept shim for any bridge function that does not return Result, so a C++ throw crossing it calls std::terminate. The C++ side calls mlx::core::quantized_matmul directly with no try/catch in all four relevant places: the quantized_matmul shim itself (src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:2926), dequantize (:2954), and the two fused wrappers quantized_linear_forward (:1063) and quantized_linear_forward_global_scale (:1099). There is no try_quantized_matmul; only 7 declarations in the entire bridge return Result (try_array_to_raw_bytes, try_eval, try_matmul, try_conv2d, try_async_eval, try_conv1d, mlx_load_safetensors), and dequantize (lib.rs:1121) has the same infallible shape as quantized_matmul.

4. Almost no model family guards it, and the per-family guards are now duplicating. Of 163 files under src/models/, 62 declare pub quantization: Option<Quantization>, and three validate it:

$ grep -ln "bits < 1\|bits > 32\|bits <= 0\|group_size < 1\|group_size <= 0" src/models/*.rs
src/models/bailing_moe.rs
src/models/gpt_neox.rs
src/models/helium.rs

That count was 1 when this issue was filed. #930 (Helium) and #946 (Bailing MoE) each landed a near-identical private copy of the same check: validate_quantization at helium.rs:385 and bailing_moe.rs:830 are line-for-line the same bits < 1 || bits > 32 plus group_size < 1 pair that #928 added at gpt_neox.rs:355 (called from validate at gpt_neox.rs:295), down to the wording of the error message. That is the argument for the shared guard getting stronger, not weaker: the check is spreading by copy-paste while the other 59 declaring families stay unguarded.

src/models/gpt2.rs and src/models/gpt_bigcode.rs remain confirmed examples of the exposure: both declare the field (gpt2.rs:84, gpt_bigcode.rs:129) and both read bits back out for the loader (gpt2.rs:211, gpt_bigcode.rs:324) with no range check anywhere.

The defect class

The shape of this recurs and is worth naming, because per-instance fixes keep missing the next one: a config-derived value passes Rust-side validation, violates an undocumented precondition of an MLX C++ entry point, and terminates the process at first inference rather than at load. #928 fixed an earlier instance of exactly this shape, where an odd rotary dims was accepted at load and then aborted inside fast_rope. Anything that only fails once a tensor reaches MLX is a load-time validation gap, not a runtime error path.

Proposed Solution

Scope item 1: bound the quantization params in the shared reconciler

Validate in reconcile_quantization_layout rather than adding a per-family check to each of the 160-plus files under src/models/. The two production call sites, QuantizedEmbedding::from_weights_with_mode (layers.rs:313) and UnifiedLinear::from_weights_with_mode (layers.rs:1183), both go through reconcile_quantization_layout_logged (layers.rs:1059), so one guard covers every family's quantized embeddings and dense projections at once, and covers families added later by default.

One loader is not covered by that guard, and needs its own. SwitchLinear::from_stacked_parts (src/models/switch_layers.rs:281) is the MoE expert loader and never calls the reconciler. It re-derives the bit width itself from the invariant packed_in * 32 == bits * num_groups * group_size, accepts the inferred value only when it lands in 2..=8, and otherwise falls back to the declared bits unchecked. group_size is not inferred at all: a declared group_size of 0 makes the denominator zero, the inference is skipped entirely, and the zero is stored on SwitchLinear::Quantized and handed to the kernel. The MoE path reaches MLX through gather_qmm, which goes through the same extract_quantized_matmul_dims and therefore the same /bits division as the dense path. A fix that guards only reconcile_quantization_layout leaves every quantized MoE checkpoint exposed, so the switch-layer loader needs the equivalent bounds check on the values it ends up storing.

The guard must be a bounds check, not an allowlist of the widths MLX actually supports. This is a real constraint, not a stylistic preference: mlxcel intentionally re-derives bit width from tensor shapes when the declared value disagrees with the stored tensors. That is the entire purpose of the reconciled flag, and it is what makes mixed-precision exports load correctly today (the Qwen3.5/3.6 MoE per-path bit overrides, the minicpm-v mxfp4 group-size case). An allowlist of {2,3,4,5,6,8} would reject checkpoints that work now. Only values that can describe no packing at all should be refused: bits outside 1..=32, and group_size below 1.

Turning the current "trust the caller" early return into an Err is the natural landing spot, since the function already returns Result<ReconciledQuant, String> and already hard-fails for unsupported affine layouts. Callers already surface that Err as a load error. Check whether any caller relies on the permissive early return for a legitimate shape; the empty/scalar-shape case it was originally written for is a separate condition from the non-positive-param case and should stay permissive if so.

Keep the per-family validate_quantization in gpt_neox.rs, helium.rs and bailing_moe.rs, or fold the three copies into the shared guard, but do not leave the shared path unguarded on the grounds that three families check.

Scope item 2: check the width of quantized embedding tables

validate_embedding_table (src/models/gpt2.rs:848, shared by Bailing MoE, GPT-2, GPT-BigCode and GPT-NeoX per its Used by: comment at line 847, called from gpt2.rs:780, gpt_bigcode.rs:761,772, gpt_neox.rs:1127 and bailing_moe.rs:2001) skips the width check when the table is quantized (gpt2.rs:874):

if !table.is_quantized() && !dim_eq(*cols, expected_cols) {

The doc comment at lines 843-845 gives the reason: a quantized table is packed along the last axis, so its stored width is a function of bit depth rather than the model width. That is correct as far as it goes, but the consequence is that a quantized embedding whose dequantized width disagrees with hidden_size is accepted at load and then aborts on the first forward. The hidden state reaches fast::layer_norm / fast::rms_norm at the wrong width and MLX throws [layer_norm] weight must have the same size as the last dimension of ... (mlx/fast.cpp:214, and :75 for the rms_norm equivalent), which crosses the same infallible bridge and terminates the same way.

The dequantized width is recoverable at load from the packed width and the (now validated) bit depth, so the check can be restored for the quantized case rather than skipped.

#946 already landed this reconstruction, but privately in one family. validate_quantized_packing (src/models/bailing_moe.rs:934) reconstructs the input width as scales.shape(-1) * group_size using the group_size that reconcile_quantization_layout settled on, compares it against the config width, and separately rejects a .biases whose shape differs from .scales. It is applied to the token table at bailing_moe.rs:1226, to every dense projection including the output head at bailing_moe.rs:1044 (inside validate_projection), and to the stacked MoE experts at bailing_moe.rs:1120 (inside validate_stacked_experts). It is a private fn in that one file. Hoisting it to a shared location and calling it from validate_embedding_table for the quantized case satisfies this scope item for all four families that use the shared helper instead of just Bailing MoE, and gives the MoE gap from scope item 1 a ready-made check as well.

Both items are the same "reject at load rather than terminate at inference" theme, and a fix should cover both.

Acceptance Criteria

• reconcile_quantization_layout returns Err for bits outside 1..=32 and for group_size < 1, with a message naming the offending field and value, instead of returning the declared pair unchanged.
• The guard is a bounds check, not an allowlist: a declared bit width that disagrees with the stored tensor shapes still reconciles as it does today, and the existing mixed-precision reconciliation tests pass unchanged.
• A checkpoint whose config.json declares "bits": 0 (and separately "bits": 33, and "group_size": 0) fails at load with an actionable error naming the field, on every quantizing family, not only gpt_neox, helium and bailing_moe.
• The quantized MoE path is covered too: SwitchLinear::from_stacked_parts cannot store a non-positive group_size or an out-of-range bits on SwitchLinear::Quantized, so a bad config fails at load rather than aborting inside gather_qmm.
• validate_embedding_table checks the dequantized width of a quantized embedding table against the model width, and rejects a mismatch at load.
• A quantized embedding table whose dequantized width disagrees with hidden_size fails at load rather than reaching fast::layer_norm / fast::rms_norm.
• Regression tests exist for both scope items that drive the bad values through the real load path, to the point where MLX would be called. These are self-enforcing: if a guard regresses, the C++ throw crosses the infallible cxx bridge and aborts the test binary rather than producing a normal test failure. Note that cargo test always builds with panic = "unwind", but std::terminate is not a panic, so the harness cannot contain it either. A test that exercises only the pure shape helper without reaching a real call site does not satisfy this criterion.
• The tree-wide count is re-verified after the fix: no per-family workaround is needed for a family to be covered, and the three existing copies in gpt_neox.rs, helium.rs and bailing_moe.rs are either removed or explicitly kept as a redundant early diagnostic.
• cargo clippy --all-targets -- -D warnings and cargo fmt --check pass.
• Validated against at least one real quantized checkpoint to confirm normal loading is unaffected, per the repository's real-checkpoint validation rule. Not run in PR fix(core): reject quantization params MLX would abort on, at load #956: the release binaries are not built in that environment and the MoE reference checkpoint is 33.6 GB. Left for the follow-up validation pass.

Technical Considerations

• The mlx/ops.cpp and mlx/fast.cpp line numbers above are against pinned commit b7c3dd6d27f45b5365b08a840310187dc503f1db, re-confirmed against the local _deps/mlx-src checkout at that commit. Re-check them after an MLX bump; the precondition is what matters, not the line number.
• A complementary hardening is to make quantized_matmul and dequantize fallible at the cxx boundary, the way try_eval and try_matmul already are. That is defense in depth of the same kind fix(server): catch MLX C++ aborts at the decode FFI boundary so a backend throw fails the request, not the process #822 landed for the decode loop (PR fix(server): survive MLX decode-eval throws instead of aborting the worker #825) and fix: make conv2d/conv1d fallible at the FFI boundary so conv shape faults don't abort the server #427 landed for conv2d/conv1d, and it would convert this class of abort into a request-level error. It is not a substitute for load-time validation: failing the request on every forward pass is still the wrong outcome for a config that could have been rejected once at load.
• The SIGFPE-versus-throw divergence between AArch64 and x86-64 noted above means the CUDA build may fail even less gracefully than the Metal build. Worth a one-line confirmation during implementation.

Refs

Originating work: #928, #836. Related abort-class hardening: #822, #427, #382, and docs/adr/0003-release-panic-unwind-with-core-thread-abort.md.