Skip to content

fix(core): QuantizedMultiLinear infers its quantization mode at load - #1030

Merged
inureyes merged 1 commit into
mainfrom
fix/issue-1028-multilinear-quant-mode
Aug 5, 2026
Merged

inureyes merged 1 commit into
mainfrom
fix/issue-1028-multilinear-quant-mode

Conversation

@inureyes

@inureyes inureyes commented Aug 5, 2026

Copy link
Copy Markdown
Member

Summary

mlxcel_core::layers::QuantizedMultiLinear stored no quantization mode and hardcoded "affine" at all three of its kernel calls, while from_weights took .biases as an Option and never checked it. A block-float embed_q / unembed_out plane (.scales, no .biases, which is what mxfp4 / nvfp4 / mxfp8 ship by design and exactly what infer_quantization_mode keys on) therefore satisfied MultiLinear::from_weights's .scales-only gate, loaded without complaint, and then handed quantized_matmul a null bias pointer under a declared affine mode. MLX's validate_mode_with_type throws "Biases must be provided for affine quantization" on precisely that, and the throw crosses the cxx bridge as an uncatchable abort at the first MLA forward rather than a load error.

This is the sibling PR #1027 flagged and deliberately left out of scope. It follows that PR's treatment of kimi_linear's private MultiLinear (mode inferred from bias presence, validate_quantization_biases at load) and the older precedent at src/models/gpt_oss.rs.

What changed

  • src/lib/mlxcel-core/src/layers.rs, QuantizedMultiLinear: new mode: &'static str field, derived at load with infer_quantization_mode(biases.is_some(), group_size, bits) in both constructors, plus a validate_quantization_biases call in from_weights that prefixes the error with the tensor prefix the way the sibling loaders do.
  • forward and forward_no_transpose were two copies of the same body differing only in the transpose bool, each with its own hardcoded "affine". They now share one private quantized_matmul helper, so the mode reaches both by construction rather than by two edits staying in sync. The issue called out missing one of the two as the obvious way to get this half right; collapsing them is how this PR makes that impossible rather than merely careful.
  • dequantize reads the same field and the same new biases_ptr() helper. It is callerless in tree, but MLX applies validate_mode_with_type inside dequantize too, so leaving one of the three hardcoded would have parked the same abort one call away from where this removes it. That is a small deliberate widening of the issue's stated scope, which named only the two forwards.
  • Doc updates: the Used by: lists on infer_quantization_mode and validate_quantization_biases, a note on MultiLinear::from_weights recording why gating on .scales alone is correct and what it needs from the quantized arm, and a paragraph in docs/adding-models.md explaining why QuantizedMultiLinear::new survives on the "already carries the bound" list when the embedding constructor next to it in the same section does not.

Decision on QuantizedMultiLinear::new: kept, made mode-aware

Kept, with the signature unchanged and the mode inferred internally.

The removal of QuantizedEmbedding::new in PR #1025 was about a signature that forced the defect: it required a biases argument and hardcoded mode: "affine", so a block-float caller could not describe its own checkpoint at all and fell to the dense branch over a packed uint32 table. QuantizedMultiLinear::new takes Option<biases>, and once the mode is derived from that Option with the same helper from_weights uses, it can describe every plane layout the loader can and cannot store a mode that contradicts them. The removal criterion does not apply to it.

The closer in-tree precedent is QuantizedWeight::new_with_mode (issue #973): also callerless, also pre-emptive, kept and made fallible so the bound lands on the path the next caller that hand-builds a quantized layer will take. docs/adding-models.md also lists QuantizedMultiLinear::{new, from_weights} in the table of loaders that already carry the bound, i.e. it is documented API model authors are pointed at, unlike the embedding constructor the same section now explicitly steers people away from. Removing it would additionally have stranded hand_built_quantized_layers_bound_their_declared_params, whose embedding half PR #1025 already retargeted at from_weights; the remaining half would have collapsed into the existing multi_linear_loader_bounds_its_declared_params and the test name would have stopped meaning anything.

Worth recording precisely, because it is a correction rather than an agreement: PR #1025's stated reason for keeping it ("it takes Option<biases> and so cannot express this bug") was true of the embedding gate defect but not of this one. With biases: None and a hardcoded affine forward, new could express issue #1028 exactly. Inferring the mode is what makes that sentence true for both defect classes, and the doc comment on new now says so instead of repeating the incomplete claim.

One thing in the issue that is not satisfiable as written

Acceptance criterion 2 asks that "a plane whose declared mode contradicts the bias plane it ships fails at load with a message naming the prefix". Nothing declares a mode on this path: QuantizedMultiLinear::from_weights takes only group_size and bits, its four callers thread the top-level quantization block into exactly those two, and this PR infers the mode rather than adding a declared-mode parameter, for the same reason PR #1027 recorded for kimi_linear (every other layer in these models infers through the shared loaders, so reading a declaration here alone would make one projection disagree with the rest of the model). With the mode inferred from biases.is_some(), the validate_quantization_biases call one line later is consistent by construction and cannot fire on any input.

It is not decorative, and this is stronger than the equivalent note in PR #1027. Regressing the inference back to a literal "affine" makes the guard fire and turns the block-float load into a load error naming the prefix, which is what the new tests actually observed when the regression was applied. So the check converts a future break of the coupling, including a declared-mode or per-prefix-override caller, into a load error at the point the mode is stored instead of an abort inside quantized_matmul. The tests assert that pairing rather than pretending a contradiction is reachable through the public loader today.

Out of scope, deliberately

Test plan

New tests in the layers.rs inline mod tests:

  • multi_linear_loader_infers_its_quantization_mode_from_the_biases_plane: pins the stored mode across affine, mxfp4, nvfp4 and mxfp8 planes, each with the packed width and group count its group_size / bits pair actually implies, and pins it on the MultiLinear enum path the four families call as well as on the inner loader. Keeps a dense plane as a negative control so a change that treats everything as quantized cannot pass.
  • multi_linear_loader_cannot_store_a_mode_that_contradicts_its_planes: loads six real export pairs with and without a bias plane and asserts the stored (mode, biases) pair always satisfies MLX's own precondition, then asserts the guard rejects both directions of the contradiction, then asserts the hand-built constructor derives the mode identically so it is not the one remaining way to build a contradicting layer.
  • The existing hand_built_quantized_layers_bound_their_declared_params gains a block-float positive control on new and pins the mode on both halves. The existing multi_linear_loader_bounds_its_declared_params is unchanged in behavior; its inline fixture was hoisted into the shared mla_plane / mla_quantized_weights helpers the new tests use, so all three drive the same geometry.

Every assertion is on the load Result or on a stored field. No test runs a forward pass, because a forward on a bad load aborts the test binary instead of failing the test.

Failure was confirmed rather than assumed: temporarily replacing the infer_quantization_mode call in from_weights with a literal "affine" made both new tests fail (the mode assertion in one, the load Result in the other, the latter because the validate_quantization_biases guard fired). The temporary change was reverted before committing and the restored line verified by grep.

Commands actually run on this branch, all green:

  • cargo fmt --all -- --check
  • cargo check -p mlxcel-core --lib --tests
  • cargo check --lib --tests
  • cargo clippy -p mlxcel-core --lib --tests -- -D warnings, exit 0
  • cargo clippy --lib --tests -- -D warnings, exit 0
  • cargo test -p mlxcel-core --lib layers::tests (74 passed, up from 72)
  • cargo test --lib models::deepseek_v3:: (14 passed)
  • cargo test --lib models::deepseek_v32 (14 passed)
  • cargo test --lib models::longcat_flash_ngram (2 passed)
  • cargo test --lib models::glm4_moe_lite (1 passed)
  • cargo doc -p mlxcel-core --no-deps --lib, checked only to confirm the new intra-doc links resolve. The single unresolved link in layers.rs is at line 4244, outside every hunk in this diff, and the crate's 102 doc warnings are pre-existing.

Not run here: cargo clippy --workspace --all-targets -- -D warnings and the full test suite, both left to the orchestrator. No MLA checkpoint is present on this machine, so the affine positive path was confirmed only against synthetic fixtures, not against real weights, and no block-float MLA export exists to confirm the new path end to end.

Closes #1028

`QuantizedMultiLinear` stored no mode and hardcoded `"affine"` at all three of its kernel calls, while `from_weights` took `.biases` as an `Option` and never checked it. A block-float `embed_q` / `unembed_out` plane (`.scales`, no `.biases`, which is what mxfp4 / nvfp4 / mxfp8 ship by design and what `infer_quantization_mode` keys on) therefore satisfied `MultiLinear::from_weights`'s `.scales`-only gate, loaded without complaint, and then handed `quantized_matmul` a null bias pointer under a declared affine mode. MLX's `validate_mode_with_type` throws "Biases must be provided for affine quantization" on exactly that, and the throw crosses the cxx bridge as an uncatchable abort at the first MLA forward rather than a load error.

The layer now carries a `mode: &'static str` derived at load with `infer_quantization_mode(biases.is_some(), group_size, bits)`, and `from_weights` calls `validate_quantization_biases` with the tensor prefix on the error. `forward` and `forward_no_transpose` were two copies of the same body differing only in `transpose`; they now share one private `quantized_matmul` helper, so there is a single mode argument rather than two that can drift. `dequantize` reads the same field for the same reason: MLX applies that precondition there too.

`QuantizedMultiLinear::new` is kept and made mode-aware rather than removed the way `QuantizedEmbedding::new` was in issue #976. That removal was about a signature that forced the defect, since requiring a `biases` argument meant a block-float caller could not describe its own checkpoint. This one takes `Option<biases>` and now derives the mode from it with the same helper the loader uses, so it can describe every plane layout the loader can. The closer precedent is `QuantizedWeight::new_with_mode` from issue #973: also callerless, also pre-emptive, kept so the bound lands on the path the next hand-builder takes.

Reached by DeepSeek V3, DeepSeek V3.2, GLM4 MoE Lite and LongCat Flash NGram through `MultiLinear::from_weights`. Latent today in all four, since the three with a `kv_b_proj` sanitizer normally insert dense planes and GLM4 MoE Lite cannot load the canonical layout for an unrelated reason (issue #1029).

Two new tests in the `layers.rs` inline module pin the inferred mode across affine, mxfp4, nvfp4 and mxfp8 planes on both the inner loader and the `MultiLinear` enum path, and pin that the stored mode never contradicts the bias plane for any declared pair the loader accepts. Both assert on the load `Result` with no forward pass, because a forward on a bad load aborts the test binary instead of failing the test. Both were confirmed to fail when the inference is regressed to a literal `"affine"`.

Closes #1028
@inureyes inureyes added status:review Under review type:bug Bug fixes, error corrections, or issue resolutions priority:low Low priority area:models Model architectures, weights, loading, metadata area:core mlxcel-core: MLX FFI, primitives, KV cache, layers labels Aug 5, 2026
@inureyes

inureyes commented Aug 5, 2026

Copy link
Copy Markdown
Member Author

Implementation Review Summary

Intent

QuantizedMultiLinear must stop hardcoding "affine" next to a possibly-null bias pointer, derive its mode from the planes the checkpoint actually ships, and validate the pair at load, so a block-float embed_q / unembed_out becomes a load error rather than an uncatchable abort inside quantized_matmul.

Findings Addressed

None. No CRITICAL or HIGH finding was found, so nothing was auto-fixed and the branch is unchanged at a3b63d49.

Verification of the four points that carried real risk

Mode threading is complete. The impl block contains exactly two kernel calls, and both take the pointer and the mode from the same object:

ffi::quantized_matmul(... self.biases_ptr(), transpose, self.group_size, self.bits, self.mode)
ffi::dequantize(...     self.biases_ptr(), self.group_size, self.bits, self.mode)

Every remaining "affine" literal inside QuantizedMultiLinear (lines 2789, 2817, 2931) is in a doc comment. self.mode is written in exactly two places, both infer_quantization_mode(biases.is_some(), ...) keyed on the same biases that becomes self.biases, so the pointer and the mode cannot disagree short of a pub-field struct literal, which the doc already calls out as the unenforceable escape hatch.

Transpose polarity is preserved exactly. Compared against main argument by argument: forward was true, // transpose and is now self.quantized_matmul(x, true); forward_no_transpose was false, // no transpose and is now self.quantized_matmul(x, false). Argument order into ffi::quantized_matmul is unchanged. The public method names and signatures are unchanged, so all 16 call sites across the four families (deepseek_v3.rs:331,345,348,349, deepseek_v32.rs:436,465,480,493,494, glm4_moe_lite.rs:238,252,255,256, longcat_flash_ngram.rs:274,288,291,292) are untouched and unaffected.

dequantize is genuinely callerless, and including it was right. No .dequantize() call site exists anywhere in src/, MultiLinear exposes no forwarder for it, and git log -S shows no commit in the repository's history ever contained embed_q.dequantize or unembed_out.dequantize. This is not scope creep: on main all three kernel calls were consistently affine, so fixing only the two the issue named would have left one method of the same type reading a different mode from the other two, which is worse than either endpoint. MLX applies validate_mode_with_type inside dequantize too.

mode: &'static str is the right choice and does not diverge awkwardly. infer_quantization_mode returns &'static str, ffi::quantized_matmul / ffi::dequantize take mode: &str, and nothing on this type needs an owned string. QuantizedWeight and QuantizedEmbedding carry String because each has a *_with_mode constructor that stores a caller-supplied mode; QuantizedMultiLinear has none, so the type now states the narrower contract. It also matches PR #1027's kimi_linear::MultiLinear, which uses mode: &'static str for the identical reason. Consistent, not divergent.

Judgment on the two decisions the PR argues for

Keeping QuantizedMultiLinear::new: the distinction holds. QuantizedEmbedding::new was removed because its signature forced the defect (required biases, hardcoded mode: "affine"), so a block-float caller could not describe its checkpoint at all. This one takes Option<biases> and now derives the mode from it with the same helper the loader uses, so the class of layer it can build is exactly the class the loader can build. QuantizedWeight::new_with_mode (#973) is the accurate precedent: callerless, pre-emptive, kept and made fallible. docs/adding-models.md already lists QuantizedMultiLinear::{new, from_weights} in the "already carries the bound" table, so it is documented API model authors are pointed at, and the new paragraph correctly explains why the embedding constructor two paragraphs above is not. The PR body's self-correction (that #1025's stated reason for keeping it was true of the gate defect but not of this one, since biases: None plus a hardcoded affine forward could express #1028) is right, and the doc comment now says the accurate thing rather than repeating the incomplete claim.

Acceptance criterion 2 is genuinely unsatisfiable as written, and keeping the guard is correct. Confirmed there is no declared-mode path into this type: from_weights takes only (prefix, group_size, bits), new takes only (weight, scales, biases, group_size, bits), and none of the four families calls any *_with_mode variant anywhere (grep for from_weights_with_mode across deepseek_v3.rs, deepseek_v32.rs, glm4_moe_lite.rs, longcat_flash_ngram.rs returns nothing). Nothing in the tree parses a quantization.mode and threads it here. The guard is a provable tautology: infer_quantization_mode returns "affine" on exactly the has_biases == true arm and one of three non-affine literals otherwise, all four members of SUPPORTED_QUANTIZATION_MODES, so both validate_quantization_mode and the (mode == "affine", has_biases) match are satisfied unconditionally. The only mode-declaring route is a pub-field struct literal, which bypasses every constructor by construction and is not a load path.

Keeping the call is right rather than decorative. It is the same pairing QuantizedEmbedding::from_weights_with_mode carries with the same comment, and the same one PR #1027 put on kimi_linear's copy, so removing it here would be the odd one out. The regression argument checks out analytically as well: replacing the inference with a literal "affine" makes validate_quantization_biases("affine", false) return Err for every with_biases == false fixture, which fails multi_linear_loader_cannot_store_a_mode_that_contradicts_its_planes on the load Result and the block-float rows of multi_linear_loader_infers_its_quantization_mode_from_the_biases_plane on the stated mode. No reachable declared-mode path was missed.

Remaining Items

  • FusedQKVLinear::from_weights_separate hardcodes "affine" (MEDIUM, pre-existing, not this PR's scope). layers.rs:2260-2264 passes the literal "affine" to from_weights_separate_with_mode, whose quantized branch builds QuantizedWeight as a struct literal at layers.rs:2399-2409 with mode: mode.to_string() and biases: qkv_biases, where qkv_biases is None unless all three .biases planes are present. That literal bypasses QuantizedWeight::new_with_mode, so no validate_quantization_biases runs on this path. A block-float separate-q/k/v export therefore reaches quantized_matmul with a null bias pointer under a declared affine mode: the identical uncatchable abort this PR removes. Sixteen families reach it through from_weights_separate (llama3, gemma, gemma2, gemma3, gemma4, qwen3, qwen3_moe, qwen3_vl, qwen3_vl_moe, cohere2, cohere2_moe, starcoder2, internlm3, jamba, helium, and the helium_tests fixture), and nothing in tree calls from_weights_separate_with_mode with a non-affine mode. Worth filing as the next issue in the fix(models): mamba and mamba2 treat a scales-without-biases embedding as non-quantized #976 / fix(models): the MLA kv_b_proj sanitizers panic on a scales-without-biases plane #1026 / fix(core): QuantizedMultiLinear hardcodes affine and never validates the bias plane #1028 series rather than widening this PR.
  • The mode is inferred from the declared group_size / bits (LOW, pre-existing, explicitly out of scope). This loader still never calls reconcile_quantization_layout, so a mixed-precision export whose embed_q is 8-bit under a 4-bit top-level default infers mxfp4 rather than mxfp8. The issue explicitly forbids fixing this here and the type's doc already records the gap.
  • No test pins the transpose polarity of the two forward methods (LOW, introduced by this PR in the sense that the refactor is the change that could have swapped them). Verified correct by direct comparison against main instead. A test would need a real ffi::quantize round-trip, which no existing layers.rs test does for this type, and every new assertion is deliberately on the load Result or a stored field because a forward on a bad load aborts the test binary.

One behavioral note that is inherent to this defect class rather than a finding: replacing the hardcoded affine with inference trades a loud abort for possibly-silent garbage in the case of a genuinely affine export whose .biases plane is missing. That case is indistinguishable from a legitimate block-float export at the plane level, and the rule is already what UnifiedLinear::from_weights, QuantizedEmbedding::from_weights and kimi_linear::MultiLinear::from_weights use tree-wide. Nothing to change.

Commands run on this branch

  • cargo fmt --all -- --check, exit 0
  • cargo clippy -p mlxcel-core --lib --tests -- -D warnings, exit 0
  • cargo clippy -p mlxcel-core --all-targets -- -D warnings, exit 0 (widens the author's run to benches for the changed crate)
  • cargo test -p mlxcel-core --lib layers::tests: 74 passed, matching the claimed count
  • cargo test -p mlxcel-core --lib multi_linear: 3 passed, both new tests present and green
  • cargo test -p mlxcel-core --lib quantiz: 46 passed
  • cargo test --lib quantiz: 136 passed, 1 ignored
  • cargo test --lib models::deepseek_v3:: 14, models::deepseek_v32 14, models::longcat_flash_ngram 2, models::glm4_moe_lite 1, models::kimi_linear 4, all passed
  • cargo test --lib mla 5 passed, cargo test --lib sanitize 126 passed, 5 ignored
  • cargo doc -p mlxcel-core --no-deps --lib: the new intra-doc links resolve. The three warnings in layers.rs are at 4244, 4399 and 4401, all outside every hunk in this diff, and only 4244 is an unresolved link (the other two are the pre-existing "links to private item" pair on select_pooled_paged_dispatch), so the PR body's claim is accurate.

Test-fixture geometry was checked by hand: every (packed_in, num_groups, group_size, bits) tuple in both new tests satisfies packed_in * 32 == bits * num_groups * group_size, and every expected mode matches infer_quantization_mode's branch order.

Verification

  • All stated requirements implemented (criterion 2 excepted, and its unsatisfiability is established above rather than asserted)
  • No placeholder/mock code remaining
  • Integrated into project code flow: the four families reach the fix through MultiLinear::from_weights with no call-site change required, since the public surface of the type is unchanged
  • Project conventions followed: matches PR fix(models): MLA kv_b_proj sanitizers reject scales-without-biases #1027's kimi_linear treatment field for field, guard for guard, comment for comment
  • Existing modules reused: infer_quantization_mode, validate_quantization_biases, validate_quantization_params, and the Used by: lists on the first two were updated
  • No unintended structural changes: two files, one of them docs, no API removed, no rename or move
  • Tests pass

Verdict

The PR correctly and completely implements issue #1028. Approve as-is.

@inureyes

inureyes commented Aug 5, 2026

Copy link
Copy Markdown
Member Author

Review addendum: one further finding, adjudicated

A second review pass raised a P1 against the infer_quantization_mode call in from_weights (layers.rs:2891), which I investigated. Summarising it so the reasoning is on the record: MLX fixes the block-float layouts to mxfp4 at 32/4, nvfp4 at 16/4 and mxfp8 at 32/8, but infer_quantization_mode keys only on bits == 8 and group_size == 16, so a no-bias plane declared (group_size = 64, bits = 4) is stored as mxfp4 at group size 64, and (64, 6) is stored as mxfp4 at 6 bits. Neither pair describes a layout MLX can honor, and QuantizedMultiLinear hands the stored pair to the kernel without reconciling it. The suggested remedy was to reconcile from tensor shapes or reject mode-incompatible params before constructing the layer.

Adjudication: real in substance, MEDIUM, not this PR's to fix, no change made. Four reasons, each checked against the tree rather than asserted.

It is a property of the shared helper, not of this PR. infer_quantization_mode returns mxfp4 or nvfp4 for any bits in 1..=32 other than 8, and it is the shared rule for UnifiedLinear::from_weights, QuantizedEmbedding::from_weights, kimi_linear::MultiLinear::from_weights and the BailingMoe validators. The loaders that do reconcile are not actually protected from the bits half either: reconcile_quantization_layout's block-float branch gates on let can_infer = ... && 32 % bits == 0, and at bits == 6 that is false, so it returns the declared (64, 6) unchanged with reconciled: false, silently. UnifiedLinear::from_weights(weights, prefix, 64, 6) on a no-bias plane therefore stores mxfp4 at 6 bits today, with no warning. This PR extends an existing tree-wide contract to one more loader; it does not create the contract.

The primary remedy is explicitly forbidden by the issue. #1028 says: "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." The PR body lists the same thing under "Out of scope, deliberately". Adding the reconciler here would be the out-of-scope architectural change a review is supposed to flag, not commit.

The secondary remedy is a shared-contract change, so it belongs in its own issue. Rejecting mode-incompatible params is shape-free and genuinely non-tautological on the bits axis ((64, 6) and (16, 6) would both be refused), so it is a real check rather than another restatement of the inference. But it is not safe on the group-size axis: reconcile_block_float_group_size_override_flags_reconciled pins the minicpm-v case where a config declares group size 64 for an mxfp4 plane that is really 32, and the correct handling there is re-derivation, not rejection. So the check would have to be bits-only, would have to live in the shared helper to be worth anything, and would change behavior for five loaders at once. That is a separate issue on infer_quantization_mode / reconcile_quantization_layout, not a fix to this PR.

The pre-PR behavior was strictly worse, so this is not a regression. On main every block-float MLA plane, canonical or not, was stored as "affine" over a null bias pointer and aborted the process at the first forward. The narrowing this PR performs is unambiguously in the right direction; what it leaves is that QuantizedMultiLinear remains weaker than UnifiedLinear, because it does not even get the reconciler's group-size re-derivation for the minicpm-v shape. That delta is the pre-existing scoped-out gap, present on main too and merely masked there by the affine hardcode aborting everything block-float first.

One correction to my earlier comment while this is in view. I described the (12, 1, 64, 6) fixture in multi_linear_loader_cannot_store_a_mode_that_contradicts_its_planes as "honest geometry". It is arithmetically honest (12 * 32 == 6 * 1 * 64) but it is not an honest mxfp4 layout, since block-float bits are fixed by the mode. The test asserting that it loads is consistent with the shared contract as it stands today, and it is a fair pin of current behavior, but it is worth knowing that the row encodes the gap above rather than a layout any exporter produces.

Updated remaining items

The verdict is unchanged: the PR correctly and completely implements issue #1028. Two follow-ups worth filing, in priority order:

  1. FusedQKVLinear::from_weights_separate hardcodes "affine" and builds QuantizedWeight by struct literal, bypassing validate_quantization_biases, on a path 16 families reach. Same defect class as fix(models): mamba and mamba2 treat a scales-without-biases embedding as non-quantized #976 / fix(models): the MLA kv_b_proj sanitizers panic on a scales-without-biases plane #1026 / fix(core): QuantizedMultiLinear hardcodes affine and never validates the bias plane #1028.
  2. infer_quantization_mode can return a block-float mode with mode-incompatible bits, and reconcile_quantization_layout's 32 % bits == 0 gate lets that pair through silently. Affects UnifiedLinear, QuantizedEmbedding, kimi_linear::MultiLinear, QuantizedMultiLinear and the BailingMoe validators alike.

@inureyes

inureyes commented Aug 5, 2026

Copy link
Copy Markdown
Member Author

Security and performance review

Threat model: a hostile or malformed checkpoint reaching MLX C++ across the cxx bridge, where a C++ throw is an uncatchable std::terminate and a Rust panic in the server is a denial of service.

No CRITICAL or HIGH finding. Nothing auto-fixed, branch unchanged at a3b63d49, working tree clean.

1. The mode and the bias pointer cannot diverge

self.mode is written in exactly two places, new (layers.rs:2834) and from_weights (layers.rs:2891), both infer_quantization_mode(biases.is_some(), ...) keyed on the same binding that is moved into self.biases in the same struct expression, with no statement in between that can change either. Both kernel call sites take self.biases_ptr() and self.mode from the same &self within one expression (layers.rs:2933-2946, 2968-2979). A repo-wide grep for .biases = , .mode = , .biases.take() and mem::take/mem::replace on these fields returns nothing for any quantized layer type, so no post-construction mutation exists either.

The pub-field struct literal remains the one theoretical divergence, and the doc already names it. It is unreachable in practice: QuantizedMultiLinear is referenced in zero files outside layers.rs repo-wide, so no literal exists to go stale.

For the shipping path the change is provably a no-op: infer_quantization_mode returns "affine" on its first branch for any plane with .biases, regardless of group_size / bits, so every affine MLA export gets the identical literal that was hardcoded before, and the C++ is_affine memcmp fast path (mlx_cxx_bridge.cpp:2942) is taken identically. Acceptance criterion 3 is satisfied by construction rather than by measurement.

2. No new panic reachable from a malformed weight map

Nothing was added to the load path that can panic: no unwrap, no expect, no slice indexing, no arithmetic. validate_quantization_params still runs first, before any tensor is copied, so group_size and bits are bounded before anything multiplies them.

biases_ptr()'s b.as_ref().unwrap() (layers.rs:2922) is a verbatim hoist of the three identical inline copies it replaces, and a mirror of QuantizedWeight::biases_ptr (layers.rs:116) and QuantizedEmbedding::biases_ptr (layers.rs:340). Reaching the None arm of UniquePtr::as_ref requires self.biases == Some(<null UniquePtr>). The only producer on the loader path is weights.get(&biases_name).map(|w| ffi::copy(w)), and ffi::copy is std::make_unique<MlxArray>(...) (mlx_cxx_bridge.cpp:825), which never returns null. Not reachable from a weight map. A hand-built new(..., Some(UniquePtr::null()), ...) could reach it, which is pre-existing across all three sibling types and not this PR's to change.

3. Lifetime and pointer safety of the refactor

&'static str cannot dangle. cxx maps &str to rust::Str, a pointer and length pair, and the C++ side never assumes NUL termination: quantized_matmul does std::memcmp(mode.data(), "affine", 6) on the fast path and std::string(mode.data(), mode.size()) otherwise (mlx_cxx_bridge.cpp:2942, 2957), and dequantize does the same (:2969). Nothing retains the Rust pointer past the call.

biases_ptr() returns a pointer into the MlxArray owned by self.biases and it is consumed as an argument in the same expression under the &self borrow, so no reallocation, move or drop can intervene, and the pointer is never stored. Moving the computation from a let binding into the argument list does not change that.

4. The guard covers every route for the four families

MultiLinear::from_weights (layers.rs:3349) is the only route, confirmed rather than assumed. QuantizedMultiLinear appears in exactly one file repo-wide, all four families name only mlxcel_core::layers::MultiLinear (deepseek_v3.rs:428, deepseek_v32.rs:1090, glm4_moe_lite.rs:327, longcat_flash_ngram.rs:348), and no code outside layers.rs even pattern-matches MultiLinear::Quantized, so there is no route to the inner type's fields either. kimi_linear keeps its own private MultiLinear and is unaffected, as PR #1027 already handled it.

5. Performance: load-time only

Per call this is a two-word field read. mode: &'static str adds 16 bytes per layer instance and allocates nowhere; the String the siblings carry would have cost one to_string() per layer at load, not per call, since &self.mode on a String also derefs to &str for free. No allocation was added or removed from the forward path on the Rust side.

The C++ side adds none either. "mxfp4", "nvfp4" and "mxfp8" are 5 bytes and fit small-string optimization, so the std::string(mode.data(), mode.size()) on the block-float branch is stack-only, and affine layers skip it entirely through the memcmp fast path. Collapsing the two forwards into one private helper adds one statically dispatched call that inlines; it is unmeasurable next to a per-head matmul.

Findings

MEDIUM, pre-existing, agreed out of scope. infer_quantization_mode can pair a block-float mode with mode-incompatible group_size / bits. I agree with the earlier adjudication and will not restate it. One corroborating fact worth having on the record: the residual is abort to abort, not abort to wrong output. MLX's CUDA quantized dispatch throws no kernel for ... bits, group size, mode (mlx/backend/cuda/quantized/quantized.cpp:252 and :484) for an unsupported tuple, so a mode-incompatible pair cannot silently produce garbage numbers.

LOW, correction to the earlier review's closing note, no code change. The note said the inference "trades a loud abort for possibly-silent garbage in the case of a genuinely affine export whose .biases plane is missing." That case does not go silent. Affine .scales are float16 or bfloat16, and validate_mode_with_type rejects any non-affine mode whose scales.dtype() != uint8 (mlx/ops.cpp:4614-4620) before any kernel runs, so such a checkpoint still fails loudly, with Scale type must be uint8 but received type float16 in place of Biases must be provided for affine quantization. There is no wrong-output path anywhere in this change, which is a better result than the PR body claims for itself.

LOW, introduced by this PR, no action recommended. Adding a pub field to a pub all-public-fields struct breaks downstream struct-literal construction. mlxcel-core is 0.4.3 with no publish = false, and there are zero such literals in the workspace, so the practical blast radius is nil.

No new attack surface. The set of checkpoints that load successfully is unchanged, since a no-biases quantized plane already loaded before this PR and only its forward differed. MLX still validates w.dtype() == uint32, the scales and biases shape agreement, the batch shape and the packing invariant w.shape(-1) * 32 / bits == scales.shape(-1) * group_size (mlx/ops.cpp:112-152) before any kernel dispatch, so a hostile checkpoint gets a throw rather than a memory-safety violation, exactly as before.

Verdict

Approve. The fix is correct, the coupling it relies on is airtight on every construction path in the tree, it introduces no panic, no dangling pointer and no allocation in the forward path, and it leaves affine deployments bit-identical.

@inureyes inureyes added status:done Completed and removed status:review Under review labels Aug 5, 2026
@inureyes
inureyes merged commit 3c38f28 into main Aug 5, 2026
5 checks passed
@inureyes
inureyes deleted the fix/issue-1028-multilinear-quant-mode branch August 5, 2026 03:58
@inureyes inureyes self-assigned this Aug 31, 2026
be-free-snowboy pushed a commit to freestylejs/snowbrain-mlxcel that referenced this pull request Oct 27, 2026
…dates the bias plane (lablup#1030)

## Summary

`mlxcel_core::layers::QuantizedMultiLinear` stored no quantization mode and hardcoded `"affine"` at all three of its kernel calls, while `from_weights` took `.biases` as an `Option` and never checked it. A block-float `embed_q` / `unembed_out` plane (`.scales`, no `.biases`, which is what mxfp4 / nvfp4 / mxfp8 ship by design and exactly what `infer_quantization_mode` keys on) therefore satisfied `MultiLinear::from_weights`'s `.scales`-only gate, loaded without complaint, and then handed `quantized_matmul` a null bias pointer under a declared affine mode. MLX's `validate_mode_with_type` throws "Biases must be provided for affine quantization" on precisely that, and the throw crosses the cxx bridge as an uncatchable abort at the first MLA forward rather than a load error.

This is the sibling PR lablup#1027 flagged and deliberately left out of scope. It follows that PR's treatment of `kimi_linear`'s private `MultiLinear` (mode inferred from bias presence, `validate_quantization_biases` at load) and the older precedent at `src/models/gpt_oss.rs`.

## What changed

- `src/lib/mlxcel-core/src/layers.rs`, `QuantizedMultiLinear`: new `mode: &'static str` field, derived at load with `infer_quantization_mode(biases.is_some(), group_size, bits)` in both constructors, plus a `validate_quantization_biases` call in `from_weights` that prefixes the error with the tensor prefix the way the sibling loaders do.
- `forward` and `forward_no_transpose` were two copies of the same body differing only in the `transpose` bool, each with its own hardcoded `"affine"`. They now share one private `quantized_matmul` helper, so the mode reaches both by construction rather than by two edits staying in sync. The issue called out missing one of the two as the obvious way to get this half right; collapsing them is how this PR makes that impossible rather than merely careful.
- `dequantize` reads the same field and the same new `biases_ptr()` helper. It is callerless in tree, but MLX applies `validate_mode_with_type` inside `dequantize` too, so leaving one of the three hardcoded would have parked the same abort one call away from where this removes it. That is a small deliberate widening of the issue's stated scope, which named only the two forwards.
- Doc updates: the `Used by:` lists on `infer_quantization_mode` and `validate_quantization_biases`, a note on `MultiLinear::from_weights` recording why gating on `.scales` alone is correct and what it needs from the quantized arm, and a paragraph in `docs/adding-models.md` explaining why `QuantizedMultiLinear::new` survives on the "already carries the bound" list when the embedding constructor next to it in the same section does not.

## Decision on `QuantizedMultiLinear::new`: kept, made mode-aware

Kept, with the signature unchanged and the mode inferred internally.

The removal of `QuantizedEmbedding::new` in PR lablup#1025 was about a signature that forced the defect: it required a `biases` argument and hardcoded `mode: "affine"`, so a block-float caller could not describe its own checkpoint at all and fell to the dense branch over a packed uint32 table. `QuantizedMultiLinear::new` takes `Option<biases>`, and once the mode is derived from that `Option` with the same helper `from_weights` uses, it can describe every plane layout the loader can and cannot store a mode that contradicts them. The removal criterion does not apply to it.

The closer in-tree precedent is `QuantizedWeight::new_with_mode` (issue lablup#973): also callerless, also pre-emptive, kept and made fallible so the bound lands on the path the next caller that hand-builds a quantized layer will take. `docs/adding-models.md` also lists `QuantizedMultiLinear::{new, from_weights}` in the table of loaders that already carry the bound, i.e. it is documented API model authors are pointed at, unlike the embedding constructor the same section now explicitly steers people away from. Removing it would additionally have stranded `hand_built_quantized_layers_bound_their_declared_params`, whose embedding half PR lablup#1025 already retargeted at `from_weights`; the remaining half would have collapsed into the existing `multi_linear_loader_bounds_its_declared_params` and the test name would have stopped meaning anything.

Worth recording precisely, because it is a correction rather than an agreement: PR lablup#1025's stated reason for keeping it ("it takes `Option<biases>` and so cannot express this bug") was true of the embedding gate defect but not of this one. With `biases: None` and a hardcoded affine forward, `new` could express issue lablup#1028 exactly. Inferring the mode is what makes that sentence true for both defect classes, and the doc comment on `new` now says so instead of repeating the incomplete claim.

## One thing in the issue that is not satisfiable as written

Acceptance criterion 2 asks that "a plane whose declared mode contradicts the bias plane it ships fails at load with a message naming the prefix". Nothing declares a mode on this path: `QuantizedMultiLinear::from_weights` takes only `group_size` and `bits`, its four callers thread the top-level `quantization` block into exactly those two, and this PR infers the mode rather than adding a declared-mode parameter, for the same reason PR lablup#1027 recorded for `kimi_linear` (every other layer in these models infers through the shared loaders, so reading a declaration here alone would make one projection disagree with the rest of the model). With the mode inferred from `biases.is_some()`, the `validate_quantization_biases` call one line later is consistent by construction and cannot fire on any input.

It is not decorative, and this is stronger than the equivalent note in PR lablup#1027. Regressing the inference back to a literal `"affine"` makes the guard fire and turns the block-float load into a load error naming the prefix, which is what the new tests actually observed when the regression was applied. So the check converts a future break of the coupling, including a declared-mode or per-prefix-override caller, into a load error at the point the mode is stored instead of an abort inside `quantized_matmul`. The tests assert that pairing rather than pretending a contradiction is reachable through the public loader today.

## Out of scope, deliberately

- No reconciliation against tensor shapes. The type's doc already records that it stores the declared pair verbatim and never calls `reconcile_quantization_layout`; that is a separate scoped gap and this PR does not touch it.
- `src/models/glm4_moe_lite.rs` is untouched. Its missing MLA sanitizer is issue lablup#1029.

## Test plan

New tests in the `layers.rs` inline `mod tests`:

- `multi_linear_loader_infers_its_quantization_mode_from_the_biases_plane`: pins the stored mode across affine, mxfp4, nvfp4 and mxfp8 planes, each with the packed width and group count its `group_size` / `bits` pair actually implies, and pins it on the `MultiLinear` enum path the four families call as well as on the inner loader. Keeps a dense plane as a negative control so a change that treats everything as quantized cannot pass.
- `multi_linear_loader_cannot_store_a_mode_that_contradicts_its_planes`: loads six real export pairs with and without a bias plane and asserts the stored `(mode, biases)` pair always satisfies MLX's own precondition, then asserts the guard rejects both directions of the contradiction, then asserts the hand-built constructor derives the mode identically so it is not the one remaining way to build a contradicting layer.
- The existing `hand_built_quantized_layers_bound_their_declared_params` gains a block-float positive control on `new` and pins the mode on both halves. The existing `multi_linear_loader_bounds_its_declared_params` is unchanged in behavior; its inline fixture was hoisted into the shared `mla_plane` / `mla_quantized_weights` helpers the new tests use, so all three drive the same geometry.

Every assertion is on the load `Result` or on a stored field. No test runs a forward pass, because a forward on a bad load aborts the test binary instead of failing the test.

Failure was confirmed rather than assumed: temporarily replacing the `infer_quantization_mode` call in `from_weights` with a literal `"affine"` made both new tests fail (the mode assertion in one, the load `Result` in the other, the latter because the `validate_quantization_biases` guard fired). The temporary change was reverted before committing and the restored line verified by grep.

Commands actually run on this branch, all green:

- [x] `cargo fmt --all -- --check`
- [x] `cargo check -p mlxcel-core --lib --tests`
- [x] `cargo check --lib --tests`
- [x] `cargo clippy -p mlxcel-core --lib --tests -- -D warnings`, exit 0
- [x] `cargo clippy --lib --tests -- -D warnings`, exit 0
- [x] `cargo test -p mlxcel-core --lib layers::tests` (74 passed, up from 72)
- [x] `cargo test --lib models::deepseek_v3::` (14 passed)
- [x] `cargo test --lib models::deepseek_v32` (14 passed)
- [x] `cargo test --lib models::longcat_flash_ngram` (2 passed)
- [x] `cargo test --lib models::glm4_moe_lite` (1 passed)
- [x] `cargo doc -p mlxcel-core --no-deps --lib`, checked only to confirm the new intra-doc links resolve. The single unresolved link in `layers.rs` is at line 4244, outside every hunk in this diff, and the crate's 102 doc warnings are pre-existing.

Not run here: `cargo clippy --workspace --all-targets -- -D warnings` and the full test suite, both left to the orchestrator. No MLA checkpoint is present on this machine, so the affine positive path was confirmed only against synthetic fixtures, not against real weights, and no block-float MLA export exists to confirm the new path end to end.

Closes lablup#1028
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:core mlxcel-core: MLX FFI, primitives, KV cache, layers area:models Model architectures, weights, loading, metadata priority:low Low priority status:done Completed type:bug Bug fixes, error corrections, or issue resolutions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant