fix(models): mamba and mamba2 treat a scales-without-biases embedding as non-quantized #976

Problem / Background

MambaModel::from_weights and Mamba2Model::from_weights decide whether the embedding table is quantized by requiring both .scales and .biases:

let embed_scales = weights
    .remove("backbone.embeddings.scales")
    .or_else(|| weights.remove("model.embed_tokens.scales"));
let embed_biases = weights
    .remove("backbone.embeddings.biases")
    .or_else(|| weights.remove("model.embed_tokens.biases"));

let embeddings = if let (Some(scales), Some(biases)) = (embed_scales, embed_biases) {
    UnifiedEmbedding::Quantized(mlxcel_core::layers::QuantizedEmbedding::new(
        embed_weight, scales, biases, group_size, bits,
    ))
} else {
    UnifiedEmbedding::Regular(mlxcel_core::layers::Embedding::new(embed_weight))
};

src/models/mamba.rs:534 through src/models/mamba.rs:551, and the identical block at src/models/mamba2.rs:763 through src/models/mamba2.rs:780.

Affine quantization stores zero-point biases; the block-float schemes (mxfp4, nvfp4, mxfp8) do not, by design. That is exactly what infer_quantization_mode (src/lib/mlxcel-core/src/layers.rs:1030) encodes, and why QuantizedEmbedding::from_weights (src/lib/mlxcel-core/src/layers.rs:253) keys the quantized branch on .scales alone and infers the mode from bias presence:

let has_biases = weights.get(&format!("{}.biases", prefix)).is_some();
let mode = infer_quantization_mode(has_biases, group_size, bits);
Self::from_weights_with_mode(weights, prefix, group_size, bits, mode)

So for a block-float export, the two mamba files take the else branch and build Embedding::new(embed_weight) over the raw packed tensor. Embedding::forward (src/lib/mlxcel-core/src/layers.rs:193) calls ffi::embedding, which is a bare mlx::core::take(weight, indices, 0) (src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:2976). No dequantization happens: the lookup returns a uint32 array whose last dimension is the packed width (hidden_size * bits / 32), not hidden_size. Nothing rejects it at load; the model reports as loaded and the failure surfaces downstream, either as a shape or dtype rejection inside a later MLX op (which crosses the cxx bridge as an uncatchable abort at the first forward pass) or as nonsense hidden states. Either way the load succeeds when it should not.

These two files also bypass the shared reconciler entirely, which is the point #958 raised for them: QuantizedEmbedding::new (src/lib/mlxcel-core/src/layers.rs:235) is a plain field-filling constructor that hardcodes mode: "affine" and calls neither reconcile_quantization_layout_logged nor validate_quantization_params. The group_size and bits come from MambaConfig::group_size() / bits() (src/models/mamba.rs:126 / :134, src/models/mamba2.rs:76 / :84), which read a private Quantization { group_size, bits } struct (src/models/mamba.rs:44, src/models/mamba2.rs:31) with no mode field at all, so even an affine mamba export cannot express a block-float mode.

Currently latent

No mamba or mamba2 checkpoint in the recommended set ships a block-float embedding, so this does not reproduce today. It becomes live the moment an mxfp4 / nvfp4 / mxfp8 mamba export appears, and it will fail as garbage output rather than as a load error, which is the expensive way to find it.

Proposed Solution

Route both files through the shared loader instead of hand-building the layer:

1. Replace the if let (Some(scales), Some(biases)) gate at src/models/mamba.rs:541 and src/models/mamba2.rs:770 with a .scales-only check, and construct via QuantizedEmbedding::from_weights (or from_weights_with_mode with a mode from mlxcel_core::layers::infer_quantization_mode(has_biases, group_size, bits)), so biases stays optional the way the struct already allows (src/lib/mlxcel-core/src/layers.rs:227).
2. Going through from_weights also picks up reconcile_quantization_layout_logged (src/lib/mlxcel-core/src/layers.rs:1175) and therefore validate_quantization_params, which closes the mamba half of fix(models): bound quantization params in the family-local MoE expert loaders #958 as a side effect rather than as a separate bounds check bolted on.
3. The two loaders currently remove() the tensors from the WeightMap while QuantizedEmbedding::from_weights takes (&WeightMap, prefix). The prefix differs by checkpoint (backbone.embeddings vs model.embed_tokens), so resolve the prefix first and pass it, rather than keeping the by-hand extraction.
4. Consider making QuantizedEmbedding::new private or #[doc(hidden)], or removing it in favour of the fallible constructors, so the next family that builds one directly cannot reintroduce either problem. fix(models): bound quantization params in the family-local MoE expert loaders #958 raised the same question.

Acceptance Criteria

• A mamba or mamba2 checkpoint whose embedding carries .scales and no .biases loads as UnifiedEmbedding::Quantized with the inferred block-float mode, not as UnifiedEmbedding::Regular over the packed tensor.
• An affine mamba / mamba2 embedding (.scales plus .biases) continues to load and generate identically.
• A non-quantized mamba / mamba2 embedding (no .scales) still takes the Embedding::new path.
• Both files reach QuantizedEmbedding through a path that runs validate_quantization_params, so a hostile "bits": 0 or "group_size": 0 fails at load rather than aborting at the first forward pass.
• models/falcon-mamba-7b-4bit (or another mamba checkpoint present under models/) produces unchanged output, confirming the affine path is untouched.
• cargo clippy --all-targets -- -D warnings and cargo fmt --check pass.

Refs

Follow-up to #958 (PR #972). #958 listed mamba.rs and mamba2.rs under Scope for the direct QuantizedEmbedding::new bypass; this issue covers the second defect on the same two lines, that the branch is gated on bias presence rather than scales presence. Same abort class as #929.