feat(core): Support mixed-bit-width q/k/v projections in the fused QKV loader

[inureyes]

Problem / Background

FusedQKVLinear::from_weights_separate (src/lib/mlxcel-core/src/layers.rs:2252, delegating to from_weights_separate_with_mode at :2281) concatenates the three packed q/k/v planes along axis 0 and infers a single quantization width from q_proj. When a checkpoint quantizes the three projections at different bit widths, the packed planes have incompatible shapes and the concatenation is a hard MLX shape error, for example [2048, 256] against [256, 512].

This is not hypothetical. mlx-community/LocateAnything-3B-4bit is a mixed_4_8 conversion: 18 of its 36 layers hold v_proj at 8 bits while q_proj and k_proj stay at 4 bits. The model could not load at all until PR #1070 worked around it.

Scope: this is a shared-loader problem, not a one-model problem

FusedQKVLinear::from_weights_separate is called by 16 model files in this tree: src/models/llama3.rs, src/models/gemma.rs, src/models/gemma2.rs, src/models/gemma3.rs, src/models/gemma4.rs, src/models/cohere2.rs, src/models/cohere2_moe.rs, src/models/qwen3.rs, src/models/qwen3_moe.rs, src/models/qwen3_vl.rs, src/models/qwen3_vl_moe.rs, src/models/jamba.rs, src/models/helium.rs, src/models/internlm3.rs, src/models/starcoder2.rs, plus src/loading/vlm_locateanything_quant.rs. The doc comment at layers.rs:2273 also names Llama3 and Mistral as users.

Any of these families will fail to load a mixed-bit-width conversion. mixed_4_8 is a standard mlx_lm quantization predicate, so such checkpoints will keep appearing on the Hub.

What PR #1070 did, and why it is a stopgap rather than the fix

The LocateAnything loader detects affected layers and dequantizes the three planes for those layers only (src/loading/vlm_locateanything_quant.rs). Properties of that workaround worth carrying into the real fix:

• It is exact, asserted at max_abs_diff == 0.0.
• Widening 4-bit to 8-bit instead was tried and measured a 3.7e-3 error, because MLX's affine quantizer snaps the group scale onto the larger-magnitude edge rather than using (max-min)/(2^bits-1), so a 4-bit group does not land on the 8-bit grid. Do not attempt requantization as the fix.
• It only engages for genuinely mixed layers; a layer whose three planes agree on both bits and group_size is left alone, with positive-control tests for uniform 4-bit and uniform 8-bit.
• The cost on the 3B checkpoint is roughly 190 MB, load-time only, and is logged.

The limitation is that this lives in one model's loader. The other 15 families get an unhelpful MLX shape error instead.

Proposed Solution

Move mixed-width handling into FusedQKVLinear itself so every caller benefits, either by keeping the three planes separate when their widths disagree (preferred, since it avoids the dequantization memory cost entirely) or by performing the same exact dequantization the LocateAnything loader does. If the fused path genuinely requires uniform width, then at minimum detect the mismatch and fail with an actionable message naming the layer and the three widths, rather than surfacing a raw MLX shape error.

Acceptance Criteria

• A mixed_4_8 checkpoint loads and generates correctly on at least one family other than LocateAnything.
• Uniform 4-bit and uniform 8-bit checkpoints are byte-identical to their current behavior, with tests pinning that.
• Numerical exactness is preserved; no requantization across bit widths.
• The per-family workaround in src/loading/vlm_locateanything_quant.rs is either removed in favor of the shared path or explicitly documented as still needed and why.
• Memory cost of any dequantization is bounded, load-time only, and logged.

Technical Considerations

The mismatch can be in bits, in group_size, or in both, so the detection predicate must compare the full (bits, group_size) pair per projection rather than bits alone. Since the fused layer today derives one width from q_proj, whichever direction is chosen has to thread per-plane quantization metadata (or a resolved "planes stay separate" decision) through to the forward path, not just through loading.