Repository navigation
fix(core): QuantizedMultiLinear hardcodes affine and never validates the bias plane #1028
Description
Problem / Background
mlxcel_core::layers::QuantizedMultiLinear reaches quantized_matmul with the mode string hardcoded to "affine" and a biases pointer that may be null:
src/lib/mlxcel-core/src/layers.rs:2858-2878(forward)src/lib/mlxcel-core/src/layers.rs:2880-2899(forward_no_transpose)
Both build biases_ptr as match &self.biases { Some(b) => ..., None => std::ptr::null() } and then pass the literal "affine". QuantizedMultiLinear::from_weights (src/lib/mlxcel-core/src/layers.rs:2825) takes .biases as an Option and never calls validate_quantization_biases. It does call validate_quantization_params (added for #958), so the group_size/bits pair is bounded, but nothing checks the bias plane against the mode.
MultiLinear::from_weights (src/lib/mlxcel-core/src/layers.rs:3279) selects the quantized arm on .scales alone, which is the correct gate. The defect is that the type it hands off to can only express affine. A block-float embed_q or unembed_out plane (mxfp4 / nvfp4 / mxfp8 ship .scales and no .biases by design, which is what infer_quantization_mode keys on) is therefore loaded without complaint and then hands MLX a null biases pointer under a declared affine mode. MLX throws "Biases must be provided for affine quantization", and that throw crosses the cxx bridge as an uncatchable std::terminate at the first forward pass rather than a load error.
Reached by four families through MultiLinear::from_weights: src/models/deepseek_v3.rs:428, src/models/deepseek_v32.rs:1090, src/models/glm4_moe_lite.rs:327, src/models/longcat_flash_ngram.rs:348. The type's own doc comment already lists exactly these four.
Currently latent, and the reachability is narrower than it looks
The three families that ship a kv_b_proj sanitizer (deepseek_v3, deepseek_v32, longcat_flash_ngram) dequantize kv_b_proj and insert dense embed_q.weight / unembed_out.weight, so .scales is normally absent by the time MultiLinear::from_weights runs. Those three reach this type only via a checkpoint that ships pre-decomposed quantized embed_q planes, because each sanitizer skips decomposition when embed_q.weight already exists.
GLM4 MoE Lite has no sanitizer at all, but it cannot currently load the canonical layout for a different reason, tracked separately in #1029. So this is latent today in all four families.
Verified against real checkpoints by reading safetensors headers, no weights downloaded: mlx-community/LongCat-Flash-Lite-8bit ships kv_b_proj with weight, scales and biases and no embed_q; mlx-community/Kimi-Linear-48B-A3B-Instruct-4bit and -8bit likewise.
Proposed Solution
Follow the pattern that just landed for kimi_linear's private MultiLinear in PR #1027, and the older precedent at src/models/gpt_oss.rs:463:
- Give
QuantizedMultiLinearamode: Stringfield, derived at load withinfer_quantization_mode(biases.is_some(), group_size, bits), and passself.modetoquantized_matmulin BOTHforwardandforward_no_transposeinstead of the hardcoded literal. Missing one of the two is the obvious way to get this half right. - Call
validate_quantization_biases(mode, biases.is_some())infrom_weights, prefixing the error with the tensor prefix the way the sibling loaders do. QuantizedMultiLinear::new(src/lib/mlxcel-core/src/layers.rs:2795) has no callers in tree and also cannot express a mode. Decide deliberately whether to give it the same treatment or remove it, and record the reasoning. Note that the siblingQuantizedEmbedding::newwas removed for exactly this reason in PR fix(models): mamba and mamba2 treat a scales-without-biases embedding as non-quantized #1025.
Do not attempt to make this loader reconcile against tensor shapes. Its doc comment already records that it stores the declared pair verbatim and never calls reconcile_quantization_layout; that is a separate, deliberately scoped gap.
Acceptance Criteria
- An
embed_qorunembed_outplane carrying.scalesand no.biasesloads with the inferred block-float mode and that mode reachesquantized_matmulon bothforwardandforward_no_transpose. - A plane whose declared mode contradicts the bias plane it ships fails at load with a message naming the prefix, rather than aborting at the first forward pass.
- An affine plane (both planes present) continues to load and produce identical output for all four families.
-
QuantizedMultiLinear::newis either mode-aware or removed, with the choice justified in the PR body. - Unit tests cover the block-float load and the contradiction rejection, asserting on the load
Resultwithout running a forward pass. -
cargo clippy --workspace --all-targets -- -D warningsandcargo fmt --all -- --checkpass.
Refs
Follow-up to #1026 (PR #1027), which corrected the same defect in kimi_linear's private MultiLinear, and to #976 (PR #1025), which corrected the same defect class in the mamba embedding loaders. Surfaced by the review sweep on PR #1027. Related: #958, #973.
Activity
fix(core): QuantizedMultiLinear infers its quantization mode and vali…
fix(core): QuantizedMultiLinear infers its quantization mode and vali…