Skip to content

fix(core): QuantizedMultiLinear hardcodes affine and never validates the bias plane #1028

Closed

Description

@inureyes

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:

  1. Give QuantizedMultiLinear a mode: String field, derived at load with infer_quantization_mode(biases.is_some(), group_size, bits), and pass self.mode to quantized_matmul in BOTH forward and forward_no_transpose instead of the hardcoded literal. Missing one of the two is the obvious way to get this half right.
  2. Call validate_quantization_biases(mode, biases.is_some()) in from_weights, prefixing the error with the tensor prefix the way the sibling loaders do.
  3. 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 sibling QuantizedEmbedding::new was 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_q or unembed_out plane carrying .scales and no .biases loads with the inferred block-float mode and that mode reaches quantized_matmul on both forward and forward_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::new is 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 Result without running a forward pass.
  • cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check pass.

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

  1. added
    type:bugBug fixes, error corrections, or issue resolutions
    area:coremlxcel-core: MLX FFI, primitives, KV cache, layers
    area:modelsModel architectures, weights, loading, metadata
    and removed on Aug 5, 2026
  2. self-assigned this
    on Aug 31, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

area:coremlxcel-core: MLX FFI, primitives, KV cache, layersarea:modelsModel architectures, weights, loading, metadatapriority:lowLow prioritystatus:doneCompletedtype:bugBug fixes, error corrections, or issue resolutions

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions