fix(models): decompose kv_b_proj so GLM4 MoE Lite loads its own layout

[inureyes]

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 fix(models): bound quantization params in the family-local MoE expert loaders #958) and refuses a scales-without-biases plane in the wording standardized across all five sanitizers in PR fix(models): MLA kv_b_proj sanitizers reject scales-without-biases #1027 (issue fix(models): the MLA kv_b_proj sanitizers panic on a scales-without-biases plane #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 fix(models): MLA kv_b_proj sanitizers reject scales-without-biases #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:

• cargo test --lib models::glm4_moe_lite (6 passed, 0 failed)
• cargo test --lib models::longcat_flash_ngram (2 passed, 0 failed)
• cargo test --lib models::youtu_vl (4 passed, 0 failed)
• cargo clippy --lib --tests -- -D warnings (exit 0)
• cargo fmt --all -- --check (exit 0)
• 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

[inureyes]

Implementation Review Summary

Intent

Give glm4_moe_lite the MLA kv_b_proj decomposition every other MLA family in the tree already has, so the family can load the canonical checkpoint layout instead of failing with Weight not found: model.layers.0.self_attn.embed_q.weight.

The decomposition is numerically correct

Checked four independent ways, since this PR adds real numerical behavior rather than converting a panic into an error:

1. Against the reference. transformers Glm4MoeLiteAttention declares kv_b_proj = nn.Linear(kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim), bias=False), so the stored weight is [num_heads * head_dim, kv_lora_rank], row-major and head-major, and expand_kv views (..., -1, qk_nope + v_head) then torch.split([qk_nope, v_head], dim=-1), which fixes the nope rows before the v rows inside each head. That is exactly reshape(&[num_heads, head_dim, -1]) plus a slice on axis 1.
2. Against two siblings. Identical to deepseek_v3.rs and youtu_vl_lm_sanitize.rs line for line, including the [0, 2, 1] transpose on the nope half only.
3. Against the runtime consumer. MultiLinear::forward is x @ w.swapaxes(-1, -2) on [heads, out, in]. Prefill: embed_q.forward_no_transpose(kv_latent) needs embed_q as [heads, kv_lora, qk_nope] to give k_nope = kv_latent @ Wk^T; unembed_out.forward(kv_latent) needs [heads, v_head, kv_lora] to give v = kv_latent @ Wv^T. Decode: embed_q.forward(q_nope) transposes back to Wk and computes the absorbed q_nope @ Wk, and unembed_out.forward(attn) maps the latent output to v_head. Both stored shapes are what the sanitizer emits.
4. Element level, now pinned by a test. See below.

Geometry confirmed against a real config

tiny-random/glm-4-moe-lite config.json: num_attention_heads=4, qk_nope_head_dim=64, v_head_dim=64, kv_lora_rank=384. That gives head_dim=128 and rows = 4 * 128 = 512, matching the shipped kv_b_proj.weight of [512, 384] exactly, with kv_lora_rank on the last axis. The fixture geometry in the tests is a scaled-down version of the same relation.

The other three review points hold as the PR describes them: the slices are copied contiguous (matching the four, not kimi); the shape cross-check cannot reject anything that would otherwise have worked, because a row count that disagrees with num_heads * head_dim reshapes to a latent width the forward path cannot consume; and slice_axis does treat end == -1 as "to the end of the axis" (utils.rs:55-61), so the kimi slice difference really is incidental.

Findings Addressed

• The pipeline stage-executor route never sanitized (HIGH). Glm4MoeLiteStageExecutor::load is a registered load path for this family (StageFamily::Glm4MoeLite), and it delegates to load_glm4_family_stage_model, which went from load_text_weights straight to filter_weight_map with no sanitize step. TransformerBlock::from_weights reads embed_q unconditionally, so the canonical layout still failed there with the same error this PR fixes for the dir route. DeepSeekV3StageExecutor::load and GlmMoeDsaStageExecutor::load both sanitize the full weight map before the filter; this one now does too, via a Sanitize hook on the shared generic with a no_sanitize passthrough for the two non-MLA GLM4 families.
• The tests could not tell the two halves apart (HIGH). The fixture had qk_nope_head_dim == v_head_dim == 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. v_head_dim is now 6, and a new test runs the decomposition over a 0, 1, 2, ... ramp and asserts the exact element identities (embed_q[h][r][c] == kv_b[(h * head_dim + c) * kv_lora_rank + r], unembed_out[h][d][r] == kv_b[(h * head_dim + qk_nope + d) * kv_lora_rank + r]) against the reference layout. It passes, which is direct element-level confirmation of the decomposition rather than shape-level only.

Both in 0069c6c7.

Remaining Items (reported, not fixed here)

• WeightLoadRoute::ConfigBacked never sanitizes (MEDIUM, pre-existing). Confirmed independently: special.rs:241-262 runs copy_weight_map then the family sanitizer then from_weights, while the config_backed.rs macro calls $weight_builder(weights, &args) with no hook. LoRA adapter loading is therefore still broken on the canonical layout for glm4_moe_lite, deepseek_v3 and deepseek_v32 alike. The PR's reasoning for not closing it here is sound, including the rejection of sanitizing inside from_weights (it takes &WeightMap, so it would need an owned copy_weight_map and would copy the whole map twice on the dir route). Wants one hook on the macro covering all three families, with its own tests.
• No MoE expert stacking (MEDIUM, pre-existing, out of scope for fix(models): GLM4 MoE Lite never decomposes kv_b_proj, so it cannot load the canonical checkpoint layout #1029). The public checkpoints ship mlp.experts.{N}.{gate,up,down}_proj.weight, but Glm4MoeLiteMoE::from_weights reads only pre-stacked mlp.switch_mlp.*. So a raw HF glm4_moe_lite directory still cannot load end to end, for a reason unrelated to kv_b_proj. This is consistent with glm4_moe, which also expects the stacked form, so it looks like a deliberate "MLX-converted checkpoints only" convention for the GLM families rather than a regression. Worth its own issue if raw HF directories are meant to load. Note that an MLX conversion does not pre-decompose kv_b_proj (mlx-lm sanitizes at load, not at convert), so this PR's fix is still required on that path.
• No MTP trailer strip (MEDIUM, pre-existing). glm4_moe_lite configs carry num_nextn_predict_layers, and the published checkpoints do ship the trailer at model.layers.{num_hidden_layers}, including eh_proj, enorm, hnorm, a full embed_tokens and a full shared_head.head. sanitize_weights iterates 0..num_hidden_layers and from_weights builds the same range, so nothing breaks, but the trailer stays resident and unused. deepseek_v3.rs strips exactly this at the end of its sanitizer. On a real MoE that is a whole extra decoder layer plus two vocab-sized planes.
• The ordering rationale is slightly overstated (LOW, introduced in 0069c6c7). The comment in glm4.rs and the commit body say sanitizing before filter_weight_map is required because the filter would drop out-of-range kv_b_proj keys. Out-of-range layers are never built, and sanitize_weights already continues when kv_b_proj is absent, so post-filter sanitation would still be correct for this family's sanitizer. The ordering is right and matches the siblings, it is just defensive consistency rather than a hard requirement here.
• The w_shape.len() != 2 arm is marginally stronger than claimed (LOW, inherited from youtu). A hypothetical kv_b_proj stored as exactly [num_heads, head_dim, kv_lora_rank] would reshape to a no-op and decompose correctly without the check, and is now refused. nn.Linear is always 2-D and no exporter ships that, so the "cannot reject anything that would otherwise have worked" claim holds in practice.

On the shared-helper extraction

The reasoning checks out. The slice and transpose difference really is incidental (slice_axis end == -1 computes what kimi's explicit stop computes), so the only substantive differences across the six blocks are kimi's missing contiguity copy and the shape cross-check that now exists in two of six. Both are differences the other implementations should lose rather than keep. Agreed that it does not belong in this PR: it would change behavior for four shipped families at once and would make a loadability fix impossible to revert on its own.

Verification

• All stated requirements implemented (kv_b_proj decomposition, dequantize before the split, PR fix(models): MLA kv_b_proj sanitizers reject scales-without-biases #1027's scales-without-biases wording, continue on an existing embed_q)
• No placeholder or mock code
• Integrated into the project code flow (dir route and, as of 0069c6c7, the pipeline stage-executor route; the ConfigBacked adapter route stays open as a tracked cross-family gap)
• Project conventions followed (module split matching youtu_vl_lm_sanitize.rs, error wording matching the five siblings)
• Existing modules reused (infer_mla_quantization_params, slice_axis, transpose_axes)
• No unintended structural changes
• Tests pass: 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 clean, cargo fmt --all -- --check clean

[inureyes]

Security and Performance Review

Scope: 5e3c5c30..0069c6c7 against 3c38f28c. Threat model: hostile or malformed checkpoint reaching model-loading code that feeds checkpoint tensors and config-declared integers into MLX across a 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 findings. Nothing auto-fixed; working tree untouched. The two HIGH items from the implementation review were already closed in 0069c6c7.

1. Abort and panic surface in the decomposition

src/models/glm4_moe_lite_sanitize.rs is a faithful copy of youtu_vl_lm_sanitize.rs, so every arithmetic hazard below is pre-existing in kind across all six MLA sanitizers and newly reachable for this family. None is introduced by this PR as a novel pattern.

• Guarded, correctly. infer_mla_quantization_params (layers.rs:1456) checks each divisor before dividing, evaluates packed_in * 32 in i64, bounds the solved pair, and re-checks described == num_groups * group_size, which is exactly what affine_dequantize compares. The .biases ok_or_else matches PR fix(models): MLA kv_b_proj sanitizers reject scales-without-biases #1027's wording verbatim. The &*b as *const _ raw pointer is sound: mlx_cxx_bridge.cpp:2968 copies biases->inner into a std::optional<array> before the call returns, and b outlives the call. The // SAFETY: comment at :108 is accurate and is an improvement over the youtu original, which has none.
• MEDIUM (pre-existing in kind, newly reachable): dequantize runs before the shape cross-check and can abort. glm4_moe_lite_sanitize.rs:111-120 is reached before the :132 check. infer_mla_quantization_params validates only the trailing axes, but mlx::core::dequantize (ops.cpp:5346-5355) additionally throws eagerly on w.dtype() != uint32 and w.ndim() < 2, and affine_dequantize (ops.cpp:5115-5124) throws on any rank or leading-axis disagreement between .weight, .scales and .biases. A crafted quantized kv_b_proj with, say, [512, 48] weight and [4, 12] scales solves cleanly through the helper and then aborts the process. Not fixable in this file alone without diverging from the five siblings; belongs in the shared-helper extraction the PR proposes.
• MEDIUM (pre-existing in kind, newly reachable): the shape cross-check does not stop a zero dimension. num_attention_heads: 0, or qk_nope_head_dim + v_head_dim == 0, makes expected_rows == 0; a kv_b_proj.weight of shape [0, kv_lora_rank] passes :134 and then hits reshape(&w_full, &[num_heads, head_dim, -1]) at :143. Reshape::output_shape (primitives.cpp:3967-3970) throws [reshape] Cannot infer the shape of an empty array when the known product is zero and an axis is inferred. Uncatchable abort. A positive-dimension precondition on num_heads, head_dim and kv_lora_rank before the reshape would close it for all six.
• LOW (pre-existing in kind, newly reachable): unchecked usize -> i32 casts and one unchecked i32 product. :51-54 truncate four config fields silently, and :133 computes num_heads * head_dim as i32 * i32. The shipped [profile.release] sets no overflow-checks, so this wraps rather than panicking in production; a wrapped-to-zero or wrapped-to-negative product is caught by the :134 comparison against a real tensor dimension in every case except the zero-dimension one above. In debug and test builds it is an overflow panic instead. Low practical reach, and identical in the five siblings.
• LOW (pre-existing in kind, newly reachable): for layer_idx in 0..args.num_hidden_layers is unbounded. :56 iterates the declared layer count doing three format!s and two lookups per pass, and continues silently when keys are absent. On main this family failed fast at layer 0 inside from_weights; the sanitizer now runs the full declared range first, so a config declaring 1e11 layers is a multi-hour hang before the fast failure. Same shape in all five siblings.

2. The shape cross-check

Confirmed it cannot reject a checkpoint that would otherwise have loaded and run correctly. The contains_key(embed_q_key) || !contains_key(kv_b_key) guard at :66 short-circuits before the check, so a pre-decomposed checkpoint (including a quantized embed_q, which still ships .weight) and a checkpoint carrying a stale kv_b_proj alongside the pair are both left untouched. A row count disagreeing with num_heads * head_dim produces a latent width the forward path cannot consume even where the reshape happens to succeed. It fires before the reshape, the slices and the transpose. It does not fire before dequantize, per the finding above, and that is the one gap in the "before any operation that could abort" claim.

3. The Sanitize hook in 0069c6c7

• load_glm4_family_stage_model has exactly three callers, all in glm4.rs:56/75/94, all explicitly updated. GlmMoeDsaStageExecutor and DeepSeekV3StageExecutor live in their own files and were not touched. No other StageFamily variant changed behavior.
• no_sanitize (glm4.rs:205) is Ok(weights) with no other statement, and map_err(anyhow::Error::msg) on an Ok is a no-op. The WeightMap moves in and out by value, which is a struct-header move with no reallocation and no iteration-order change. Glm4 and Glm4Moe are genuinely unaffected.
• Hook position relative to filter_weight_map is correct and matches the two sibling stage executors. Verified that decomposed keys are filtered properly: classify_weight_key (partial_loading.rs:145-154) parses the layer index out of the model.layers. prefix generically, so out-of-range self_attn.embed_q.weight and unembed_out.weight classify as Layer(N) and are dropped. No retention of other stages' planes.
• LOW (introduced in 0069c6c7): sanitizing before the filter widens the abort surface from the stage's own layers to all layers. A quantized kv_b_proj inconsistency in a layer this stage does not own now reaches dequantize where previously the filter would have discarded the key. Correct ordering regardless (it matches the siblings, and out-of-range keys must survive until the split), but worth naming.
• LOW (introduced in 0069c6c7, already noted in the implementation review): the ordering rationale in the comment is stronger than the requirement. sanitize_weights already continues on an absent kv_b_proj, so post-filter sanitation would also be correct for this family. Defensive consistency, not a hard constraint.

4. Resource handling

Clean. WeightMap is HashMap<String, UniquePtr<MlxArray>>; remove transfers ownership to a local that drops at the end of the iteration, and the .unwrap()s at :73 and :76 are each guarded by a contains_key on the same key in the same single-threaded scope. No double free: the MLX graph node holds its own mlx::core::array copies of w, s and b, so dropping the Rust handles is not a use-after-free for the retained result. No leak: keys the loop inserts are either consumed by from_weights or dropped with the map.

A partially decomposed map cannot be observed after an error. sanitize_weights takes WeightMap by value and both call sites move their binding into it (glm4_moe_lite.rs:869, glm4.rs:245), so on Err the map is dropped with the error and no caller holds a handle to it.

5. Performance and memory

• Nothing added to the decode path. The diff touches load functions only.
• Every MLX op used here is lazy, so the per-layer dequantized plane and the two copies are graph nodes at sanitize time, not allocations. The removed kv_b_proj planes stay reachable through the graph until the first eval, not because the remove failed but because w_full retains them as inputs; they are released as the graph is consumed. This is the same lifetime the five siblings have.
• Peak at eval is roughly the packed plane plus the dense plane plus the two halves, transiently, per layer, and the halves together are about the size of the dense plane. On the largest plausible geometry for this family that is tens of MB per layer, sequential, not cumulative.
• The stage-executor path builds decomposition graphs for all declared layers before the filter trims them. Because the ops are lazy, out-of-range layers cost graph nodes only and are freed by filter_weight_map. Negligible.

Verdict

Approve on security and performance grounds. The decomposition adds no novel unsafe construct, the raw-pointer use across the cxx boundary is sound and now carries a // SAFETY: comment the original lacked, ownership is clean on every path including the error paths, and the Sanitize hook is inert for the two non-MLA families. The remaining abort surfaces are the shared MLA-sanitizer pattern, not this PR's work, and every one of them argues for the decompose_mla_kv_b_proj extraction the PR already recommends filing. The two MEDIUM items above are worth carrying into that issue as hardening requirements: validate w's dtype and rank plus the leading-axis agreement between the three planes before dequantize, and require positive num_heads, head_dim and kv_lora_rank before the reshape.