fix(models): the MLA kv_b_proj sanitizers panic on a scales-without-biases plane #1026

Problem / Background

Four MLA kv_b_proj sanitizers decide the plane is quantized by weights.contains_key("<prefix>.kv_b_proj.scales"), then unconditionally .unwrap() the .biases removal:

• src/models/kimi_linear.rs:1281 (let biases = weights.remove(...kv_b_proj.biases).unwrap();)
• src/models/deepseek_v3.rs:1149
• src/models/deepseek_v32.rs:929
• src/models/longcat_flash_ngram.rs:1105

Affine quantization stores zero-point .biases; the block-float modes (mxfp4, nvfp4, mxfp8) carry none by design, which is exactly what mlxcel_core::layers::infer_quantization_mode keys on. A kv_b_proj exported block-float therefore satisfies the .scales gate and then panics on the .unwrap(). A panic is worse than the load error it should be: in the server it takes down the process rather than rejecting one model load. Note this is a Rust panic, not the uncatchable C++ std::terminate class, because it fires before dequantize is reached.

Each of the four then calls dequantize with a hardcoded "affine" mode, so even a checkpoint that carried both planes could not express a block-float kv_b_proj.

Already solved once, in-tree

src/models/youtu_vl_lm_sanitize.rs:66-73 is the corrected fifth sibling of this exact code and already does the right thing:

let b_key = format!("{}.kv_b_proj.biases", prefix);
let b = weights.remove(&b_key).ok_or_else(|| {
    format!(
        "layer {layer_idx}: kv_b_proj has scales but no biases at key `{b_key}`; \
         the checkpoint may be corrupted or only partially converted"
    )
})?;

All five sites already share mlxcel_core::layers::infer_mla_quantization_params for the packed-pair bound (issue #958), so this is the one remaining divergence between them.

Currently latent

Verified against the real checkpoints by reading their safetensors headers, no weights downloaded: mlx-community/Kimi-Linear-48B-A3B-Instruct-4bit and -8bit both ship kv_b_proj with weight, scales and biases (declared "mode": "affine"), and moonshotai/Kimi-Linear-48B-A3B-Instruct ships weight only. So no shipped checkpoint takes the panic path today. It goes live the moment a block-float MLA export appears, and it presents as a process kill rather than a load error.

Second, lower-severity item in the same file

kimi_linear.rs's private MultiLinear::forward (src/models/kimi_linear.rs:171-186) passes a null biases pointer to quantized_matmul under a hardcoded "affine" mode, and MultiLinear::from_weights (src/models/kimi_linear.rs:231) never calls validate_quantization_biases. That combination is what MLX throws "Biases must be provided for affine quantization" on, and it crosses the cxx bridge as an uncatchable abort.

This branch is currently unreachable, not a live defect. The sanitizer at src/models/kimi_linear.rs:1324-1325 always dequantizes kv_b_proj and inserts embed_q.weight / unembed_out.weight as dense tensors with no .scales ("Store as dense MultiLinear weights (no re-quantization)"), and those two tensors exist in no shipped checkpoint, so is_quantized is always false there. It is worth correcting anyway so a future change to the sanitizer cannot silently reintroduce a live abort. src/models/gpt_oss.rs:463 is the in-tree precedent for the fix (store an inferred mode rather than hardcoding "affine", and call mlxcel_core::layers::validate_quantization_biases).

Proposed Solution

1. Replace the four .unwrap()s with the ok_or_else form from youtu_vl_lm_sanitize.rs:66-73, keeping the message wording consistent across all five sites so the five converge rather than diverge further.
2. In kimi_linear.rs's private MultiLinear, derive the mode with mlxcel_core::layers::infer_quantization_mode(biases.is_some(), group_size, bits), store it, pass it to quantized_matmul instead of the hardcoded "affine", and call mlxcel_core::layers::validate_quantization_biases at load, following gpt_oss.rs:463.
3. Consider whether the five near-identical MLA decomposition blocks should share a helper. This is a question for the implementer to evaluate and report on rather than a required change, since the blocks differ in their config accessors.

Acceptance Criteria

• A kv_b_proj carrying .scales and no .biases produces a load error naming the missing key, not a panic, in all four families.
• The error message matches the wording already used in youtu_vl_lm_sanitize.rs.
• An affine kv_b_proj (both planes present) continues to load and dequantize identically in all four families.
• kimi_linear's MultiLinear no longer hardcodes "affine" and validates the bias plane against the mode at load.
• Unit tests cover the scales-without-biases rejection for at least kimi_linear and one deepseek family, asserting on the load Result and not running a forward pass.
• cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check pass.

Refs

Follow-up to #976 (PR #1025), which addressed the same defect class (bias presence used as a proxy for quantized-ness) in the mamba and mamba2 embedding loaders. Found by the review sweep on PR #1025; see that PR's review comments for the full inventory. Related: #958, #929, #973.