Repository navigation
fix(core): QuantizedMultiLinear infers its quantization mode at load - #1030
Conversation
`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
Implementation Review SummaryIntent
Findings AddressedNone. No CRITICAL or HIGH finding was found, so nothing was auto-fixed and the branch is unchanged at Verification of the four points that carried real riskMode threading is complete. The impl block contains exactly two kernel calls, and both take the pointer and the mode from the same object: Every remaining Transpose polarity is preserved exactly. Compared against
Judgment on the two decisions the PR argues forKeeping Acceptance criterion 2 is genuinely unsatisfiable as written, and keeping the guard is correct. Confirmed there is no declared-mode path into this type: Keeping the call is right rather than decorative. It is the same pairing Remaining Items
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 Commands run on this branch
Test-fixture geometry was checked by hand: every Verification
VerdictThe PR correctly and completely implements issue #1028. Approve as-is. |
Review addendum: one further finding, adjudicatedA second review pass raised a P1 against the 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. 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 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 ( The pre-PR behavior was strictly worse, so this is not a regression. On One correction to my earlier comment while this is in view. I described the Updated remaining itemsThe verdict is unchanged: the PR correctly and completely implements issue #1028. Two follow-ups worth filing, in priority order:
|
Security and performance reviewThreat model: a hostile or malformed checkpoint reaching MLX C++ across the cxx bridge, where a C++ throw is an uncatchable No CRITICAL or HIGH finding. Nothing auto-fixed, branch unchanged at 1. The mode and the bias pointer cannot diverge
The For the shipping path the change is provably a no-op: 2. No new panic reachable from a malformed weight mapNothing was added to the load path that can panic: no
3. Lifetime and pointer safety of the refactor
4. The guard covers every route for the four families
5. Performance: load-time onlyPer call this is a two-word field read. The C++ side adds none either. FindingsMEDIUM, pre-existing, agreed out of scope. 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 LOW, introduced by this PR, no action recommended. Adding a 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 VerdictApprove. 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. |
…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
Summary
mlxcel_core::layers::QuantizedMultiLinearstored no quantization mode and hardcoded"affine"at all three of its kernel calls, whilefrom_weightstook.biasesas anOptionand never checked it. A block-floatembed_q/unembed_outplane (.scales, no.biases, which is what mxfp4 / nvfp4 / mxfp8 ship by design and exactly whatinfer_quantization_modekeys on) therefore satisfiedMultiLinear::from_weights's.scales-only gate, loaded without complaint, and then handedquantized_matmula null bias pointer under a declared affine mode. MLX'svalidate_mode_with_typethrows "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 privateMultiLinear(mode inferred from bias presence,validate_quantization_biasesat load) and the older precedent atsrc/models/gpt_oss.rs.What changed
src/lib/mlxcel-core/src/layers.rs,QuantizedMultiLinear: newmode: &'static strfield, derived at load withinfer_quantization_mode(biases.is_some(), group_size, bits)in both constructors, plus avalidate_quantization_biasescall infrom_weightsthat prefixes the error with the tensor prefix the way the sibling loaders do.forwardandforward_no_transposewere two copies of the same body differing only in thetransposebool, each with its own hardcoded"affine". They now share one privatequantized_matmulhelper, 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.dequantizereads the same field and the same newbiases_ptr()helper. It is callerless in tree, but MLX appliesvalidate_mode_with_typeinsidedequantizetoo, 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.Used by:lists oninfer_quantization_modeandvalidate_quantization_biases, a note onMultiLinear::from_weightsrecording why gating on.scalesalone is correct and what it needs from the quantized arm, and a paragraph indocs/adding-models.mdexplaining whyQuantizedMultiLinear::newsurvives 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-awareKept, with the signature unchanged and the mode inferred internally.
The removal of
QuantizedEmbedding::newin PR #1025 was about a signature that forced the defect: it required abiasesargument and hardcodedmode: "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::newtakesOption<biases>, and once the mode is derived from thatOptionwith the same helperfrom_weightsuses, 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.mdalso listsQuantizedMultiLinear::{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 strandedhand_built_quantized_layers_bound_their_declared_params, whose embedding half PR #1025 already retargeted atfrom_weights; the remaining half would have collapsed into the existingmulti_linear_loader_bounds_its_declared_paramsand 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. Withbiases: Noneand a hardcoded affine forward,newcould express issue #1028 exactly. Inferring the mode is what makes that sentence true for both defect classes, and the doc comment onnewnow 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_weightstakes onlygroup_sizeandbits, its four callers thread the top-levelquantizationblock into exactly those two, and this PR infers the mode rather than adding a declared-mode parameter, for the same reason PR #1027 recorded forkimi_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 frombiases.is_some(), thevalidate_quantization_biasescall 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 insidequantized_matmul. The tests assert that pairing rather than pretending a contradiction is reachable through the public loader today.Out of scope, deliberately
reconcile_quantization_layout; that is a separate scoped gap and this PR does not touch it.src/models/glm4_moe_lite.rsis untouched. Its missing MLA sanitizer is issue fix(models): GLM4 MoE Lite never decomposes kv_b_proj, so it cannot load the canonical checkpoint layout #1029.Test plan
New tests in the
layers.rsinlinemod 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 itsgroup_size/bitspair actually implies, and pins it on theMultiLinearenum 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.hand_built_quantized_layers_bound_their_declared_paramsgains a block-float positive control onnewand pins the mode on both halves. The existingmulti_linear_loader_bounds_its_declared_paramsis unchanged in behavior; its inline fixture was hoisted into the sharedmla_plane/mla_quantized_weightshelpers the new tests use, so all three drive the same geometry.Every assertion is on the load
Resultor 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_modecall infrom_weightswith a literal"affine"made both new tests fail (the mode assertion in one, the loadResultin the other, the latter because thevalidate_quantization_biasesguard 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 -- --checkcargo check -p mlxcel-core --lib --testscargo check --lib --testscargo clippy -p mlxcel-core --lib --tests -- -D warnings, exit 0cargo clippy --lib --tests -- -D warnings, exit 0cargo 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 inlayers.rsis 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 warningsand 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