Fix mixed-bit quantized load for sanitize()-derived MLA projections

[rajanshxrma]

Summary

Loading a mixed-bit quantized checkpoint fails with a shape mismatch when the source tensor a model's sanitize() derives another tensor from carries a per-tensor quantization override.

Concretely: deepseek_v3.py, deepseek_v32.py, glm4_moe_lite.py, kimi_linear.py, and longcat_flash.py all derive the absorbed MLA projections embed_q/unembed_out from kv_b_proj inside sanitize(), dequantizing and re-quantizing them at whatever bits/group_size kv_b_proj itself used. But embed_q/unembed_out never appear in config["quantization"] — the converter never saw those paths, since they only come into existence inside sanitize(). So load_model's class_predicate (mlx_lm/utils.py) falls through to f"{p}.scales" in weights, which is True (sanitize just created those scales), and nn.quantize allocates the module at the global default bits/group_size instead of the bits it was actually packed at. On any layer where the source tensor's override differs from the default, load_weights then raises a shape mismatch.

Repro (from #1451):

from mlx_lm import load
load("kwonhan/GLM-4.7-Flash-mixed-4-6-g32-mlx")
# ValueError: Expected shape (20, 512, 24) but received shape (20, 512, 36)
# for parameter model.layers.0.self_attn.embed_q.weight

Fixes #1451.

Fix

When a path is missing from config["quantization"] but its weight already carries .scales, infer the actual group_size/bits from the packed weight and scales shapes (in_dims = m.weight.shape[-1], same formula every affected sanitize() already uses on the source tensor) instead of assuming the top-level default. This is generic — it isn't specific to any one model file — and a no-op for the common uniform-bit case, since the inferred bits equal the default that was used to produce the checkpoint in the first place.

Testing

• Added test_load_model_with_mixed_bit_derived_mla_projection in tests/test_utils.py: builds a tiny deepseek_v3 model, quantizes it uniformly, then re-packs embed_q/unembed_out at a different bit width (without adding their paths to the quantization map) to simulate a kv_b_proj override — confirms the old code raises the shape-mismatch error from Mixed-bit MLA checkpoints fail to load — sanitize() re-quantizes derived embed_q/unembed_out at kv_b_proj's bits, but the modules are built at the global default #1451, and the fix loads cleanly with the modules correctly reporting the overridden bits.
• python -m unittest tests.test_models — 76 passed.
• python -m unittest tests.test_utils tests.test_prompt_cache — all pass except two pre-existing failures in test_convert/test_load_model_with_custom_get_classes caused by an incomplete local HF cache with network disabled, unrelated to this change.
• pre-commit run --files mlx_lm/utils.py tests/test_utils.py (black, isort) — clean.

[michalk8]

Thanks for spotting this, I've modified your code a bit, happy to merge now.