fix(models): MLA kv_b_proj sanitizers reject scales-without-biases - #1027
Conversation
Four MLA `kv_b_proj` sanitizers decided the plane was quantized by probing for `.scales`, then unconditionally `.unwrap()`ed the `.biases` removal: `kimi_linear`, `deepseek_v3`, `deepseek_v32` and `longcat_flash_ngram`. Affine quantization stores zero-point `.biases`; the block-float modes (mxfp4, nvfp4, mxfp8) carry none by design, which is exactly what `infer_quantization_mode` keys on. A block-float `kv_b_proj` therefore satisfies the `.scales` gate and then panics on the `.unwrap()`. That is a Rust panic rather than the uncatchable C++ `std::terminate` class, because it fires before `dequantize` is reached, but in the server it still takes the process down instead of rejecting one model load. All four now use the `ok_or_else` form already in tree at `src/models/youtu_vl_lm_sanitize.rs`, which is the corrected fifth sibling of this exact code, with the same message wording so the five converge rather than diverge further. Each enclosing `sanitize_weights` already returned `Result<WeightMap, String>` and already used `?` for `infer_mla_quantization_params`, so no signature changed. The `dequantize` call stays hardcoded `"affine"` in all five: a genuine block-float `kv_b_proj` is still unsupported, and the change is that it is refused by name instead of unwinding. The `youtu_vl_lm_sanitize` comment is rewritten to the same wording. It carried a terse "M1:" review-artifact label; the five sites are now identical in behavior, message and rationale. Second, lower-severity item in the same file: `kimi_linear`'s private `MultiLinear` passed a null biases pointer to `quantized_matmul` under a hardcoded `"affine"` mode and never validated the bias plane at load. It now derives the mode with `infer_quantization_mode(biases.is_some(), group_size, bits)`, stores it, passes it to `quantized_matmul`, and calls `validate_quantization_biases` at load, following `src/models/gpt_oss.rs` `ExpertLinear::from_weights`. This branch is unreachable today, not a live defect: the sanitizer always dequantizes `kv_b_proj` and inserts `embed_q.weight` / `unembed_out.weight` as dense tensors with no `.scales`, and those two tensors exist in no shipped checkpoint, so `is_quantized` is always false there. It is a bound on what a future change to that sanitizer can reintroduce. Tests: a scales-without-biases rejection per family, all four rather than the two the issue asked for, each asserting on the load `Result` and its message and running no forward pass, because a forward pass on a bad load aborts the test binary instead of failing the test. Each keeps a both-planes-present positive control so a check that rejected every quantized `kv_b_proj` could not pass. The longcat case pins the full key because it is the one family whose prefix carries a sub-attention index. A new `MultiLinear` test pins the stored mode across affine, mxfp4, nvfp4 and mxfp8 planes. The per-family `build` fixture closures are hoisted into shared `affine_kv_b_proj_weights` / `block_float_kv_b_proj_weights` helpers so both tests in each family drive the same geometry. Verified by temporarily restoring the `.unwrap()` in `kimi_linear`: the new test failed with `called Option::unwrap() on a None value` at that line, confirming both that it is a panic and that the test catches it. The temporary change was reverted before committing. Refs #958, #973, #976
inureyes
added
type:bug
Implementation Review SummaryIntentStop four MLA Findings AddressedNone. No CRITICAL or HIGH finding was raised, so nothing was auto-fixed and the branch is unchanged at VerifiedPartial mutation on the new error path is not observable. Each site Interpolated index is the layer index at every site. dsv3 Affine behavior is preserved. The only non-comment change in the four sanitizers is The Tests. All assert on the load Ran locally on this branch, all green: Scope disciplineIssue item 3 asked the implementer to evaluate a shared helper and report, not build one. No helper was built: the diff is 7 files, all under
Remaining items (reported, not fixed)
Verification
|
Security and performance reviewThreat model applied: a hostile or malformed checkpoint reaching model-loading code that feeds checkpoint tensors and config-declared integers across the cxx bridge, where a C++ throw is an uncatchable Verdict: approve. No CRITICAL and no HIGH findings, and none introduced by this PR. Four MEDIUM/LOW items below, all reported rather than fixed; three are pre-existing and one is documentation-only. 1. The panic is removed, not movedConfirmed at all five sites. Between the 2. Resource handling on the new early return: cleanAll four 3. Error message: clean, and safer than the brief assumedThe interpolated 4. The guard's completeness, and the prior reviewer's bypass claimConfirmed, and it is broader than stated. Pre-existing, not introduced here. Two routes, not one:
Failure mode is MLX's 5. Performance: clean, and it is load-timeAll five sites are inside Remaining findingsMEDIUM, pre-existing. The MEDIUM, pre-existing. Four of the five sites hand LOW, introduced here, documentation only. The doc comment at LOW, not a regression. The inferred mode is not cross-checked against MLX's mode-fixed geometry. MLX requires LOW, acknowledged in the PR body. For a genuine block-float export, "the checkpoint may be corrupted or only partially converted" mis-describes a well-formed but unsupported checkpoint. Wording consistency across the five sites is the better property here, so the trade the PR made is the right one. NotesNo code changes were made: nothing reached the CRITICAL or HIGH bar that this pass auto-fixes, and every MEDIUM is pre-existing and outside the changed files. Working tree left clean. |
…laim
The doc comment on kimi_linear's private MultiLinear::from_weights and the matching test doc comment both claimed the quantized branch was "not reachable today" and that the mode threading was only a bound against a future sanitizer change. That claim is wrong: sanitize_weights only inserts a dense embed_q.weight / unembed_out.weight pair inside the if weights.contains_key(&kv_b_key) guard, so a checkpoint that ships no kv_b_proj.weight at all skips that block outright and leaves its own embed_q / unembed_out planes untouched. MlaAttention::from_weights calls this loader unconditionally, so a checkpoint shipping pre-decomposed, quantized embed_q.weight plus embed_q.scales with no kv_b_proj.weight lands here with is_quantized == true today, not hypothetically.
Both comments now state the actual reachability path and explicitly flag that the branch must not be deleted as dead code, since the prior wording invited exactly that. The validate_quantization_biases call-site comment was checked and does not claim the branch is unreachable (it only notes that specific tautological check cannot fire, which remains true regardless of branch reachability), so it was left as is. Also pinned the stored mode ("affine") on the dense leg of kimi_linear_multi_linear_infers_its_quantization_mode_from_the_biases_plane, alongside the existing !is_quantized assertion.
Validation:
- cargo check --lib --tests: clean
- cargo test --lib models::kimi_linear: 4 passed
- cargo clippy --lib --tests -- -D warnings: clean
- cargo fmt --all -- --check: clean
Refs #1026
…dates the bias plane (#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 #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 - 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 #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 #1028
#1032) ## Summary `src/models/glm4_moe_lite.rs` had no `sanitize_weights` and no MLA `kv_b_proj` decomposition anywhere, while `MlaAttention::from_weights` loads `embed_q` and `unembed_out` unconditionally (`src/models/glm4_moe_lite.rs:327-333`). Every published `glm4_moe_lite` checkpoint stores the MLA up-projection as a single `self_attn.kv_b_proj.weight` and ships no `embed_q` tensor at all, so the family failed the canonical layout of its own architecture with `Weight not found: model.layers.0.self_attn.embed_q.weight`. This is a live failure, not a latent one, and every other MLA family in the tree already carries this decomposition. ## What changed - `src/models/glm4_moe_lite_sanitize.rs` (new): `sanitize_weights` decomposes `kv_b_proj` into the per-head `embed_q` / `unembed_out` pair. Adapted from `src/models/youtu_vl_lm_sanitize.rs`, the most complete of the five existing implementations and the only one that cross-checks the tensor against `[num_heads * head_dim, kv_lora_rank]` before reshaping. Reuses `mlxcel_core::layers::infer_mla_quantization_params` for the packed pair (issue #958) and refuses a scales-without-biases plane in the wording standardized across all five sanitizers in PR #1027 (issue #1026), verbatim. - `src/models/glm4_moe_lite.rs`: declares the sanitize module, re-exports `sanitize_weights`, and calls it from `Glm4MoeLiteModel::load` between `load_text_weights` and `from_weights`. - `src/models/glm4_moe_lite_sanitize_tests.rs` (new): five tests. The code lives in its own module rather than being appended to `glm4_moe_lite.rs`, the same split `youtu_vl_lm_sanitize.rs` uses, because the runtime module is already well past the repo's 500-line target. ### Why the shape cross-check is carried over Of the five sibling sanitizers only Youtu-VL cross-checks the tensor shape, and it is worth having. The reshape below the split is where a `config.json` that disagrees with the stored tensor stops being recoverable: MLX reports a bad reshape by throwing, and that throw crosses the cxx bridge as `UniquePtr<MlxArray>` rather than `Result`, so it aborts the process during weight sanitization instead of failing the load. The check cannot reject anything that would otherwise have worked, because a row count that disagrees with `num_heads * head_dim` produces a latent width the forward path cannot consume even when the reshape happens to succeed. ### Why the quantized plane is dequantized first The split runs per head along the row axis. Slicing the packed plane would leave each half carrying group scales and biases that describe the whole row, so the halves have to be produced in dense space. `dequantize` is hardcoded `"affine"` here, the same as in all five siblings, which is why a block-float `kv_b_proj` is refused by name rather than silently mis-decomposed. ## The `load` versus `from_weights` wiring decision `src/model_metadata.rs:185` registers two entry points for this family, and the `WeightLoadRoute::ConfigBacked` arm at `src/loading/mod.rs:673` reaches `from_weights` without going through `load`, so a sanitizer placed only in `load` does not cover it. I checked what the siblings actually do for their own weight route before choosing, and the answer is not uniform: - `KimiLinear` and `LongcatFlash` / `LongcatFlashNgram` are `WeightLoadRoute::Special`, and that route **does** sanitize: `src/loading/special.rs:241-262` takes an owned copy through `copy_weight_map` and runs the family sanitizer before `from_weights`. - `Youtu-VL` is a VLM with `adapter: Some(...)`, so `load_model_with_adapter` rejects it before any weight route runs. Not applicable. - `DeepSeekV3` and `DeepSeekV32` are `WeightLoadRoute::ConfigBacked`, the same route as this family, and both **share the gap**: `try_load_config_backed_model_from_weights` is macro-generated (`src/loading/config_backed.rs:50-68`) and calls `$weight_builder(weights, &args)` with no hook where a sanitizer could go. So the gap is not a property of MLA families, it is a property of the `ConfigBacked` route: the hand-written `Special` route has a place to put the sanitizer and the macro-generated one does not. It affects exactly the three `ConfigBacked` MLA families, and `load_model_from_weights` is called from one place only (`src/loading/mod.rs:636`), so the blast radius is LoRA adapter loading, not ordinary model loading. **Decision: sanitize in `load` only, matching `DeepSeekV3` and `DeepSeekV32` exactly.** Closing the gap unilaterally for this one family means either moving it to the `Special` route or adding a per-family hook to a macro that drives around sixty registrations, and the second of those fixes the siblings, which this PR was explicitly scoped not to do. Sanitizing inside `from_weights` was rejected on a separate ground: it would need an owned `copy_weight_map`, and the `load` path already owns its map, so the dir route would copy the whole weight map a second time. That is a real peak-memory regression on a large MoE in exchange for a route this family is not yet used on. **The sibling gap deserves its own issue.** It is a genuine pre-existing bug, not an accepted design: LoRA adapter loading is already broken for `DeepSeekV3` and `DeepSeekV32` on the canonical checkpoint layout today, failing with the same `Weight not found: ...embed_q.weight` this PR fixes for the dir route. The asymmetry with the `Special` route is the evidence that it was an oversight. The fix wants to be one hook on the `ConfigBacked` registration, applied to all three families at once with its own tests, not three copies. ## On extracting a shared helper for the six decomposition blocks PR #1027 declined this for three recorded reasons. Checked against the tree, all three hold but two are weaker than they look: - **Only Youtu cross-checks the `w_full` shape.** True, and this PR makes it two of six. This is the one real behavioral difference, and it is a difference the other four should lose, not one worth preserving. - **The slice and transpose helpers differ.** Kimi uses raw `mlxcel_core::slice` with explicit stops plus `swap_axes`; the rest use `slice_axis` plus `transpose_axes`. This is incidental, not semantic: `slice_axis` treats `end == -1` as "to the end of the axis" (`src/lib/mlxcel-core/src/utils.rs:52-61`), which is exactly what Kimi's explicit stop computes, and its in-code note about `stop=-1` is about the raw `slice`, not about `slice_axis`. - **Kimi does not copy its slices contiguous.** True, and it is the odd one out rather than a considered choice: the other five all copy so `MultiLinear`'s matmul sees well-formed strides regardless of backend. What genuinely varies per family is the prefix and the layer predicate: four use `model.layers.{l}.self_attn`, LongCat carries a sub-attention index (`self_attn.0` / `self_attn.1`), and Kimi skips its linear-attention layers. Both are naturally parameters, not obstacles. A `decompose_mla_kv_b_proj(&mut weights, prefix, layer_label, geometry) -> Result<(), String>` next to `infer_mla_quantization_params` in `mlxcel_core::layers` would be roughly sixty lines and would delete something like two hundred and fifty across five files. **Conclusion: the extraction is now justified, and it should not ride in this PR.** It changes behavior for four shipped families in three ways at once (they gain the shape cross-check, Kimi gains contiguity, Kimi changes slice helper), which needs a per-family test matrix and its own bisect point. Bundling it here would also make a fix that restores a family's basic loadability impossible to revert on its own. PR #1027 declined once on the record, so overturning that call belongs in an issue where the three reasons can be answered, not in a paragraph attached to an unrelated fix. Recommend filing it as a follow-up. ## Test plan Five tests in `src/models/glm4_moe_lite_sanitize_tests.rs`, all driving the real `sanitize_weights` and the real `Glm4MoeLiteModel::from_weights` with synthetic weight maps built from `mlxcel_core::from_slice_f32`, all asserting on the returned `Result`. None runs a forward pass: a load that produced a malformed `embed_q` reaches MLX through the cxx bridge as `UniquePtr<MlxArray>` rather than `Result`, so a throw there aborts the whole test binary instead of failing one test. - `sanitize_synthesizes_the_mla_pair_no_checkpoint_ships`: asserts the unsanitized checkpoint fails with exactly `Weight not found: model.layers.0.self_attn.embed_q.weight` (the failure the issue reports, taken from the real loader rather than restated), then that the sanitized one loads and carries `embed_q` at `[heads, kv_lora_rank, qk_nope]` and `unembed_out` at `[heads, v_head, kv_lora_rank]`. - `sanitize_dequantizes_kv_b_proj_before_the_per_head_split`: an honest affine 4-bit plane (`packed_in * 32 == bits * num_groups * group_size`, UINT32 packed so `dequantize` does not throw on the positive control) decomposes to the same shapes, consumes all three source planes, and leaves both halves dense. - `sanitize_rejects_scales_with_no_biases`: positive control first, then a block-float plane is refused naming `model.layers.0.self_attn.kv_b_proj.biases` in PR #1027's wording. - `sanitize_leaves_an_already_decomposed_checkpoint_alone`: a checkpoint shipping the pair plus a stale `kv_b_proj` keeps its `kv_b_proj.weight`, which is the observable proof the `continue` fired rather than the split running. - `sanitize_rejects_a_kv_b_proj_that_disagrees_with_the_config`: a tensor carrying an extra head is refused before the reshape, naming the shape the config describes. Commands actually run, all from a clean run on this branch: - [x] `cargo test --lib models::glm4_moe_lite` (6 passed, 0 failed) - [x] `cargo test --lib models::longcat_flash_ngram` (2 passed, 0 failed) - [x] `cargo test --lib models::youtu_vl` (4 passed, 0 failed) - [x] `cargo clippy --lib --tests -- -D warnings` (exit 0) - [x] `cargo fmt --all -- --check` (exit 0) - [x] `cargo check --lib --tests` (clean) Not run here: the full-workspace clippy and test sweep, and any load against a real `glm4_moe_lite` checkpoint. No `glm4_moe_lite` weights are present on this machine, and the three public repos are tiny random fixtures rather than a released model, so the fix is verified against the published tensor layout and not against a production checkpoint. ## Review round 2 Two findings from implementation review, fixed in `0069c6c7`. **The pipeline stage-executor route never sanitized.** `Glm4MoeLiteStageExecutor::load` is a second registered load path for this family (`StageFamily::Glm4MoeLite`), and it delegates to the shared `load_glm4_family_stage_model`, which went from `load_text_weights` straight to `filter_weight_map` with no sanitize step. `models::glm4_moe_lite::TransformerBlock::from_weights` reads `embed_q` unconditionally, so the canonical layout still failed there with the same `Weight not found: model.layers.N.self_attn.embed_q.weight`. `DeepSeekV3StageExecutor::load` (`src/distributed/pipeline/stage_executor/deepseek_v3.rs:84`) and `GlmMoeDsaStageExecutor::load` (`src/distributed/pipeline/stage_executor/glm_moe_dsa.rs:42`) both sanitize the full weight map before the filter. The generic now takes a `Sanitize: Fn(WeightMap, &A) -> Result<WeightMap, String>` hook called in that position; `Glm4MoeLiteStageExecutor` passes `sanitize_weights` and the two non-MLA GLM4 families pass a `no_sanitize` passthrough. This is distinct from the `ConfigBacked` gap discussed above, which stays open and cross-family. **The tests could not tell the two halves apart.** The fixture set `qk_nope_head_dim` and `v_head_dim` both to 4, which leaves `[HEADS, KV_LORA_RANK, QK_NOPE]` and `[HEADS, V_HEAD, KV_LORA_RANK]` both valid with the nope and v halves swapped, and every fixture tensor was zero-filled, so no assertion could observe element order at all. A wrong half, a wrong reshape axis order or a missing transpose would all have passed. `v_head_dim` is now 6, and a sixth test decomposes a `0, 1, 2, ...` ramp and asserts the exact element identities against the reference layout: `embed_q[h][r][c] == kv_b[(h * head_dim + c) * kv_lora_rank + r]` for the transposed nope half, and `unembed_out[h][d][r] == kv_b[(h * head_dim + qk_nope + d) * kv_lora_rank + r]` for the v half as stored. It passes, so the decomposition is now pinned at the element level rather than by shape alone. Re-verified after both: `cargo test --lib models::glm4_moe_lite` (7 passed, 0 failed), `cargo test --lib distributed::pipeline` (317 passed, 0 failed), `cargo clippy --lib --tests -- -D warnings` (exit 0), `cargo fmt --all -- --check` (exit 0). Closes #1029
1 participant
Summary
Four MLA
kv_b_projsanitizers decided the plane was quantized by probing for.scales, then unconditionally.unwrap()ed the.biasesremoval. Affine quantization stores zero-point.biases; the block-float modes (mxfp4, nvfp4, mxfp8) carry none by design, which is exactly whatinfer_quantization_modekeys on. A block-floatkv_b_projtherefore satisfies the.scalesgate and then panics. All four now use theok_or_elseform already in tree atsrc/models/youtu_vl_lm_sanitize.rs:66-73, the corrected fifth sibling of this exact code, with the same message wording so the five converge instead of diverging further.This is latent today. No shipped checkpoint takes the path:
mlx-community/Kimi-Linear-48B-A3B-Instruct-4bitand-8bitboth shipkv_b_projwithweight,scalesandbiases, andmoonshotai/Kimi-Linear-48B-A3B-Instructshipsweightonly. It goes live the moment a block-float MLA export appears, and it presents as a process kill rather than a load error.What changed
src/models/kimi_linear.rs,src/models/deepseek_v3.rs,src/models/deepseek_v32.rs,src/models/longcat_flash_ngram.rs: the.biasesremoval is nowok_or_else(...)?with the wording fromyoutu_vl_lm_sanitize.rs. No signature changed: every enclosingsanitize_weightsalready returnsResult<WeightMap, String>and already used?forinfer_mla_quantization_params, so the?had a home at all four sites. This was flagged as the interesting part of the task, and it turned out to be a non-event.dequantizecall stays hardcoded"affine"at all five sites. A genuine block-floatkv_b_projis still unsupported; what changes is that it is refused by name instead of unwinding. See "What this deliberately does not do" below.src/models/youtu_vl_lm_sanitize.rs: comment only. It carried a terseM1:review-artifact label; it now states the same rationale as the other four, so the five are identical in behavior, message and comment. The code there was already correct.src/models/kimi_linear.rsprivateMultiLinear: gained amode: &'static strfield, derived at load withmlxcel_core::layers::infer_quantization_mode(biases.is_some(), group_size, bits)and passed toquantized_matmulin place of the hardcoded"affine", plus amlxcel_core::layers::validate_quantization_biasescall at load.src/models/gpt_oss.rs:463is the precedent.The
MultiLinearchange is reachable todayCorrecting an earlier claim in the other direction, per the security review: the branch is reachable today, not merely a bound against a future sanitizer change.
sanitize_weights(src/models/kimi_linear.rs:1323-1325) only inserts a denseembed_q.weight/unembed_out.weightpair inside itsif weights.contains_key(&kv_b_key)guard, wherekv_b_keyis{attn_prefix}.kv_b_proj.weight; when a checkpoint carries nokv_b_proj.weightat all, that whole decomposition block is skipped, so anyembed_q/unembed_outplanes the checkpoint shipped itself pass through sanitization untouched.MlaAttention::from_weights(src/models/kimi_linear.rs:576-577) then calls this loader unconditionally for every non-linear-attention layer, so a checkpoint that ships pre-decomposed, quantizedembed_q.weightplusembed_q.scales(and the matching pair forunembed_out) with nokv_b_proj.weightto decompose from lands here withis_quantized == truetoday. This branch must not be deleted as dead code.One honest caveat on the
validate_quantization_biasescall: with the mode inferred frombiases.is_some()on the line above, the check is consistent by construction and cannot fire as written. It is kept because it is the assertion that couples the two helpers, so the day the mode comes from anywhere else (a declaredquantization.mode, a per-prefix override) the contradiction is caught where the mode is stored rather than insidequantized_matmul, where it is an uncatchable abort. The comment at the call site says exactly this rather than implying it is load-bearing.What this deliberately does not do
Supporting a genuine block-float
kv_b_projwould mean threading the inferred mode into thedequantizecall at all five sites and verifying the decomposition against a real block-float MLA export. That is a feature, not this bug fix, and the acceptance criteria here ask for a load error naming the missing key. Worth recording: for a genuine block-float export the shared message ("the checkpoint may be corrupted or only partially converted") understates the situation slightly, since such a checkpoint is well-formed and merely unsupported here. Wording consistency across the five sites was judged the more valuable property, per the issue.On issue item 3, the shared helper
Evaluated, and the recommendation is yes, but not in this PR. Detail so the next person does not have to re-derive it.
It is more feasible than the issue's caveat suggests. The differing config accessors are a non-issue: all five resolve to
i32at the call site already, so a helper taking(weights, prefix, num_heads, qk_nope, v_head, kv_lora_rank, layer_label)and returningResult<(wk, wv), String>absorbs the whole block, and the per-family loop (kimi skipping linear-attention layers, longcat iterating two sub-attentions per layer) stays at the call site where it belongs. That is roughly 225 lines across five files collapsing to about 110. These five sites have now been touched twice for the same defect class, once for #958 and once here, which is the strongest argument for consolidating them.What makes it wrong to do here is that the five are not actually identical, and unifying them forces three behavioral decisions that this PR is not equipped to make:
youtu_vl_lm_sanitize.rscross-checksw_fullagainst[num_heads * head_dim, kv_lora_rank]before reshaping; the other four do not. A shared helper either adds a new load-time rejection to four families or drops an existing guard from one.kimi_linear.rsusesmlxcel_core::slice+swap_axeswhere the other four useslice_axis+transpose_axes. Equivalent, so this one is free.kimi_linear.rsdoes notcopy()thewk/wvslices contiguous, where the other four do, with a comment in youtu saying it exists soMultiLinear's matmul has well-formed strides regardless of the upstream backend. Unifying silently changes kimi's contiguity. That deserves to be the change under review, not a rider.A refactor that adds a new rejection to four families and changes the contiguity of a fifth belongs in a PR where that is the headline. Happy to file it as a follow-up issue on request.
Separately, a sibling that this PR does not touch
mlxcel_core::layers::QuantizedMultiLinear(src/lib/mlxcel-core/src/layers.rs:2775) has the identical defect inforwardandforward_no_transpose: a null biases pointer under a hardcoded"affine", withfrom_weightsgating on.scalesalone and takingbiasesas anOptionwith no validation. It is reached by DeepSeek V3, DeepSeek V3.2, GLM4 MoE Lite and LongCat Flash NGram throughMultiLinear::from_weights, which also gates on.scalesalone.Its reachability story is arguably better than the
kimi_linearone fixed here, because all four sanitizers skip decomposition whenembed_q.weightalready exists, so a checkpoint shipping a pre-decomposed quantizedembed_qreaches it directly. Left out of scope on purpose: it is a different crate, the reachability claim needs verifying against whatmlx_lm.convertactually emits for a pre-decomposed MLA model, and adding it would blur this PR's story. Flagging it rather than silently fixing or silently skipping it.Test plan
Commands actually run on this branch, all green:
cargo fmt --all -- --checkcargo clippy --lib --tests -- -D warnings, exit 0cargo check --lib --testscargo test --lib models::kimi_linear(4 passed)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::youtu_vl(4 passed), since its comment was touched and it shares the message wordingNew tests, one rejection case per family rather than the two families the issue required:
kimi_linear_sanitize_rejects_a_kv_b_proj_with_scales_and_no_biasesquantized_kv_b_proj_rejects_scales_with_no_biases(deepseek_v3)quantized_kv_b_proj_rejects_scales_with_no_biases(deepseek_v32)sanitize_rejects_scales_with_no_biases(longcat_flash_ngram), which pins the full key because this is the one family whose prefix carries a sub-attention index (self_attn.0), the site where a copy of the fix could most easily blame the wrong tensorkimi_linear_multi_linear_infers_its_quantization_mode_from_the_biases_plane, pinning the stored mode across affine, mxfp4, nvfp4 and mxfp8 planes and confirming a dense projection stays unquantizedEach assertion is on the load
Resultand its message, and no test runs a forward pass, because a forward pass on a bad load aborts the test binary instead of failing the test. Each rejection test keeps a both-planes-present positive control first, so a check that rejected every quantizedkv_b_projcould not pass it. The per-familybuildfixture closures were hoisted into sharedaffine_kv_b_proj_weights/block_float_kv_b_proj_weightshelpers so both tests in each family drive the same geometry.Failure was confirmed rather than assumed: temporarily restoring the
.unwrap()inkimi_linear.rsmade the new test fail withcalled Option::unwrap() on a None valueat that exact line, which confirms both that the failure mode is a Rust panic (not thestd::terminateclass, as the issue states) and that the test catches it. The temporary change was reverted before committing, and a grep confirms no.unwrap()remains on anykv_b_projbiases removal.Not run here:
cargo clippy --workspace --all-targets -- -D warningsand the full test suite, both left to the orchestrator.mlxcel-corewas not modified, so no-p mlxcel-corerun was made. No real MLA checkpoint is present on this machine, so the affine positive path was confirmed only against the synthetic UINT32 fixtures, not against real weights.Closes #1026