[inureyes]

Summary

MambaModel::from_weights and Mamba2Model::from_weights decided whether the embedding table was quantized by requiring both .scales and .biases. Affine quantization stores zero-point biases; the block-float modes (mxfp4, nvfp4, mxfp8) carry none by design, which is exactly what infer_quantization_mode keys on. A block-float embedding export therefore fell to the else branch and built a plain Embedding over the raw packed uint32 table. Embedding::forward is a bare mlx::core::take with no dequantization, so the lookup returned a uint32 array whose last axis was the packed width (hidden_size * bits / 32) rather than hidden_size. Nothing rejected it at load: the model reported as loaded and the failure surfaced later as garbage hidden states, or as an MLX C++ throw crossing the cxx bridge as an uncatchable std::terminate at the first forward pass.

The defect is latent today because no mamba or mamba2 checkpoint in the recommended set ships a block-float embedding. It goes live the moment one appears, and it would present as wrong output rather than as a load error, which is the expensive way to find it.

Both files now resolve the embedding prefix (backbone.embeddings or model.embed_tokens, one contains_key probe) and hand the whole map to UnifiedEmbedding::from_weights, which gates on .scales alone and infers the mode from bias presence.

This supersedes PR #972

#972, merged for issue #958, declined to route mamba through the shared loader and instead made the hand-built QuantizedEmbedding::new fallible. Two of its four stated reasons do not hold:

• "takes the table out of the map by remove, so it cannot address a single-prefix map loader": the prefix is resolvable at the call site with one contains_key, and the remove maintains no invariant. The same function already loads in_proj / x_proj / dt_proj / out_proj / lm_head through UnifiedLinear::from_weights(&weights, ...), which copies and leaves the keys in the map, and there is no leftover-weight check anywhere in the load path.
• "treats .biases as optional where mamba requires both": that requirement is the bug this PR fixes, so the reason is circular.

The other two reasons are real, and this PR accepts them as costs:

• Routing costs one lazy ffi::copy (mlx::core::copy) of the embedding table instead of moving it out of the map. That is a lazy-array node share rather than an element copy, and it is what every other model already pays.
• The loader stores the reconciled pair rather than the declared one. For a uniform-quant checkpoint reconciliation is a no-op and emits no warning; for a mixed-precision one it is a strict improvement, since the declared top-level pair is the wrong one for an embedding quantized at a different bit width.

Routing through the shared loader also picks up reconcile_quantization_layout and therefore validate_quantization_params, which is the bound QuantizedEmbedding::new was carrying on its own since #958.

New load-time rejection for these two families

Both files now call validate_embedding_table (src/models/gpt2.rs) right after the embedding loads. This is a NEW rejection for mamba and mamba2: it requires a 2-D table with non-zero rows, vocab_size <= rows, and, for a non-quantized table, a width equal to hidden_size (for a quantized table it reconstructs the dequantized width as scales.shape(-1) * group_size through validate_quantized_packing). It is the same guard BailingMoe, Gpt2, GptBigCode and GptNeoX already run, so this is not a new policy, only a new caller.

It is also what gives the regression tests an observable. MambaModel::embeddings is private and the test module is a sibling rather than a descendant, so a test cannot inspect the variant, and a forward pass on a regressed load aborts the test binary instead of failing it. With the guard in place the assertion is simply that the block-float load succeeds, because the regressed path (a Regular embedding over the packed table) now fails at load.

QuantizedEmbedding::new removed

After this change it has zero production callers across all four workspace members (verified by grep). It is removed, so the next family cannot rebuild the same defect by hand-building an affine-only layer that requires a biases argument. QuantizedMultiLinear::new stays: it is pre-emptive and also callerless, but it takes Option<biases> and so cannot express this bug.

What changed

• src/models/mamba.rs, src/models/mamba2.rs: prefix resolution plus UnifiedEmbedding::from_weights, then validate_embedding_table. The missing-embedding error still names both spellings rather than regressing to the loader's single-prefix "Weight not found". The stale comment describing the remove chain is replaced with one that says why the gate is .scales-only.
• src/lib/mlxcel-core/src/layers.rs: QuantizedEmbedding::new removed, with the reason recorded on from_weights; the three doc references to it retargeted; the QuantizedEmbedding half of hand_built_quantized_layers_bound_their_declared_params retargeted at QuantizedEmbedding::from_weights over a synthetic WeightMap, keeping the QuantizedMultiLinear::new half intact.
• src/models/mamba_tests.rs, src/models/mamba2_tests.rs: the existing fixture closures hoisted into shared helpers, the existing test's doc comment corrected, and three new tests per family (see below).
• src/models/gpt2.rs, src/lib/mlxcel-core/src/layers.rs: "Used by" lists updated to include Mamba and Mamba2.
• docs/adding-models.md: QuantizedEmbedding::new removed from the table of loaders that already carry the quant bound, with a note explaining why hand-building an embedding is no longer the answer for a multi-spelling prefix.

Tests

Per family:

• The block-float regression (*_loads_a_block_float_embedding_as_quantized): mxfp4 at packed_in 1 and mxfp8 at packed_in 2, both dequantizing to hidden_size 8, .scales present and .biases absent, load must succeed. nvfp4 is not covered because group_size 16 cannot describe a width of 8 with a non-zero number of groups, and inventing a second hidden_size for it would have been dishonest geometry.
• Both prefix spellings (*_resolves_both_embedding_prefix_spellings): affine and block-float both load under model.embed_tokens.*, and a checkpoint carrying neither spelling fails with a message naming both.
• The guard itself (*_rejects_an_embedding_table_narrower_than_hidden_size): a float table whose width is not hidden_size errors at load with a message naming hidden_size.
• The existing float-table case in *_rejects_quantization_params_that_would_abort_the_embedding_lookup now inserts a genuine [8, 8] float table instead of reusing the packed [8, 1] one with its scales stripped, because the new guard correctly rejects the latter. The guard was not weakened to accommodate the old fixture.
• The hostile-pair loop over switch_layers::HOSTILE_QUANT_PARAMS is unchanged and still passes: with .biases present the inferred mode is affine, and reconcile_quantization_layout calls validate_quantization_params before it looks at any shape.

The block-float tests were confirmed to fail on regression by temporarily restoring the bias-requiring gate in mamba.rs: both new cases failed with backbone.embeddings.weight is [8, 1] but hidden_size is 8, which is a load error rather than an abort. The temporary change was reverted before committing.

Test plan

Commands actually run, all green:

• cargo fmt --check
• cargo clippy --lib --tests -- -D warnings
• cargo clippy -p mlxcel-core --lib --tests -- -D warnings
• cargo check --lib --tests
• cargo check -p mlxcel-core --lib --tests
• cargo test --lib models::mamba (14 passed, 2 ignored; the ignored two are the pre-existing requires serial MLX execution conv-state tests)
• cargo test -p mlxcel-core --lib layers::tests (72 passed), including hand_built_quantized_layers_bound_their_declared_params
• cargo test --lib models::gpt2::tests (27 passed), since the shared guard gained callers

Orchestrator-side, after the above:

• cargo clippy --workspace --all-targets -- -D warnings (all four workspace members, every target), exit 0
• cargo fmt --all -- --check, the exact CI gate
• Real-checkpoint shape verification of the new guard against falcon-mamba-7b-4bit, mamba2-1.3b-4bit and mamba2-130m, see the comment below

Pre-existing and unrelated: a bare cargo test --workspace --lib on a Linux/CUDA box built without --features cuda aborts the root lib test binary with terminate called ... [cuda_kernel] No CUDA back-end (SIGABRT) after roughly 1926 tests. This was confirmed to reproduce identically on main at e1ea5c2 with the same warm target directory (1927 tests, same abort), so it is a property of this build configuration rather than of this branch. Every candidate test in the in-flight window passes when run individually. CI does not use this configuration.

Not run: the acceptance criterion comparing models/falcon-mamba-7b-4bit output before and after. No mamba or mamba2 checkpoint is present under models/ on this machine, so the affine path was not confirmed against real weights. The affine path is covered by the synthetic positive control in the existing guard test, which is weaker evidence. cargo clippy --all-targets -- -D warnings and the full workspace test suite were also not run here; only --lib --tests was, plus the narrow test filters above.

Closes #976

[inureyes]

Real-checkpoint verification of the new validate_embedding_table guard

The PR body notes that no mamba or mamba2 checkpoint is present locally, so the acceptance criterion naming models/falcon-mamba-7b-4bit could not be executed. That gap is now closed for the part that actually matters, which is whether the new load-time rejection would refuse a checkpoint that loads today.

Tensor shapes come from each repo's safetensors header (two HTTP Range requests: bytes 0-7 give the little-endian u64 header length, then the header JSON maps tensor name to dtype and shape), and the config values from raw/main/config.json. No weights were downloaded.

checkpoint	config	embeddings.weight	embeddings.scales	biases
mlx-community/falcon-mamba-7b-4bit	falcon_mamba, vocab 65024, hidden 4096, gs 64, bits 4	U32 [65024, 512]	F16 [65024, 64]	present
mlx-community/mamba2-1.3b-4bit	mamba2, vocab 50288, hidden 2048, gs 64, bits 4	U32 [50288, 256]	F16 [50288, 32]	present
mlx-community/mamba2-130m	mamba2, vocab 50288, hidden 768, no quantization	F16 [50288, 768]	absent	absent

Against the four checks the guard applies:

checkpoint	2-D	vocab_size <= rows	width	packing invariant	verdict
falcon-mamba-7b-4bit	yes	65024 <= 65024	scales.shape(-1) * gs = 64 * 64 = 4096 == hidden_size	512 * 32 == 4 * 64 * 64 (16384)	loads
mamba2-1.3b-4bit	yes	50288 <= 50288	32 * 64 = 2048 == hidden_size	256 * 32 == 4 * 32 * 64 (8192)	loads
mamba2-130m	yes	50288 <= 50288	cols 768 == hidden_size (non-quantized branch)	n/a, float table	loads

All three pass, so the guard adds no rejection for the checkpoints this family is validated against today. It exercises both branches of the guard as well: two affine quantized tables and one float table.

The same arithmetic covers the "loads and generates identically" criterion for the affine path. Both quantized checkpoints satisfy packed_in * 32 == bits * groups * group_size exactly at their declared group_size 64 / bits 4, so reconcile_quantization_layout returns the declared pair unchanged with reconciled = false and emits no warning. The reconciled pair this PR now stores is byte-identical to the declared pair the old hand-built path stored, which is why routing through the shared loader cannot change what these two checkpoints produce.

Still not executed: an actual generation run on real weights. The above is shape-level ground truth, not an output comparison.

[inureyes]

Implementation Review Summary

Intent

Stop mamba / mamba2 from gating the quantized embedding branch on .biases presence, so a block-float export (.scales, no .biases) loads as UnifiedEmbedding::Quantized instead of as a plain Embedding over the raw packed uint32 table.

Verdict: the PR implements issue #976 correctly and completely. No CRITICAL or HIGH finding, nothing auto-fixed, no commit added.

1. Prefix resolution: dropping the per-tensor .or_else() chain is not a regression

The old chain did not support mixed spellings, it silently paired tensors from different prefixes. Both pathological orderings get strictly better:

• weight under backbone.embeddings, .scales under model.embed_tokens: the old code built a QuantizedEmbedding from mismatched planes, which MLX rejects in validate_quantized_input as an uncatchable abort at the first forward pass. The new code resolves backbone.embeddings, finds no .scales there, builds a Regular embedding over the packed table, and validate_embedding_table refuses it at load naming hidden_size.
• The mirrored ordering behaves the same way.

No real checkpoint can carry mixed spellings here, because both files hardcode backbone.layers.{i}.* and backbone.norm_f.weight for everything else (src/models/mamba.rs:571, src/models/mamba2.rs:807). A genuinely transformers-spelled checkpoint would fail at backbone.layers.0.mixer.conv1d.weight regardless of how the embedding is spelled, so the model.embed_tokens branch is only reachable at num_hidden_layers: 0, which is what the new *_resolves_both_embedding_prefix_spellings fixture uses. Theoretical only, and the change is the safer direction.

2. validate_embedding_table cannot refuse a checkpoint that works on main

Checked all four conditions against the failure modes they replace:

• 2-D and non-zero rows: mamba embedding tables are always [rows, packed_or_width].
• vocab_size <= rows: padded-vocab checkpoints go the safe direction (config vocab_size at or below the padded row count). The state-spaces lineage (vocab_size 50277 against a 50280-row table) passes; all three validated checkpoints have vocab_size == rows.
• Width: for a float table cols == hidden_size; for a quantized one scales.shape(-1) * reconciled_group_size == hidden_size plus MLX's own validate_quantized_input predicate, mirrored in validate_quantized_packing. Anything failing either throws inside MLX at the first forward pass on main, so the guard converts an abort into a load error rather than refusing a working load.
• falcon_mamba specifically: model_type == "falcon_mamba" only flips use_bcdt_rms in compute_time_step_rank (src/models/mamba.rs:117). Nothing about the variant changes the embedding geometry, and 65024 rows / 64 * 64 == 4096 == hidden_size passes.
• Tied vs untied lm_head is not reachable by the guard: it inspects only the embedding table, and lm_head still loads through UnifiedLinear::from_weights.

The one class that "works on main and now fails" is a config whose vocab_size exceeds the table rows and whose prompts happen never to emit a high token id. That is an out-of-bounds gather on main, so refusing it is the point of the guard.

3. Affine behavior is preserved, and no new warning fires

• infer_quantization_mode(true, _, _) returns "affine", identical to the removed constructor's hardcoded mode.
• The affine branch of reconcile_quantization_layout trusts group_size and re-derives bits through infer_quantization_bits, which returns caller_bits unchanged when packed_in * 32 == num_groups * group_size * bits. That holds exactly for falcon-mamba-7b-4bit (512 * 32 == 64 * 64 * 4) and mamba2-1.3b-4bit (256 * 32 == 32 * 64 * 4), so reconciled == false and reconcile_quantization_layout_logged emits no tracing::warn!.
• The group_size == 16 && bits == 4 NVFP4-repack special case is unreachable for both, since both declare group_size 64.

So the reconciled pair the loader now stores is the declared pair for every validated checkpoint. The only genuinely new work on the affine path is validate_quantization_params, which is what issue #958 asked for.

4. QuantizedEmbedding::new removal is complete

git ls-files | xargs grep -l 'QuantizedEmbedding::new' returns one file, src/lib/mlxcel-core/src/layers.rs, and the single hit is the doc comment at line 5310 recording the removal. That covers all four workspace members, tests/, examples/, benchmarks/, python/, spike/ and the docs. src/lib/mlxcel-surgery/src/ops/scale.rs:264 mentions QuantizedEmbedding::from_weights in prose only. mlxcel-core is not published, and the struct fields stay pub, so this is not an external break.

5. The new tests do fail on a reverted fix, and they fail rather than abort

Every new assertion is on the Result of MambaModel::from_weights / Mamba2Model::from_weights. There is no forward pass anywhere in the new tests, so nothing can abort the test binary in place of failing it.

The revert claim is provable from evidence already on the board, without re-running it: *_rejects_an_embedding_table_narrower_than_hidden_size is green, which establishes that a Regular embedding over an [8, 1] table produces a load error containing hidden_size. Restoring the bias-requiring gate routes block_float_embedding_weights(prefix, 1) to exactly that Regular path, so *_loads_a_block_float_embedding_as_quantized panics on its .expect. The mxfp8 case lands the same way at [8, 2]. The error text the author reports matches validate_embedding_table's non-quantized branch verbatim.

One coupling worth recording: the block-float tests get their observable from the guard, so if validate_embedding_table were ever dropped from these two loaders they would go green again even with the bias gate restored. The test doc comments say so explicitly, so it is a stated tradeoff rather than an oversight.

6. The same defect class elsewhere

A full sweep of the tracked tree found no second instance of the exact defect, a both-planes-required predicate choosing Quantized over Regular for an embedding. Nearby findings, all out of scope for #976:

• src/models/deepseek.rs:971 (SwitchLinear::from_stacked_parts) is the only surviving (Some(scales), Some(biases)) => Quantized, _ => Regular for MoE experts. It is already masked: DeepSeekModel::from_weights rejects mode != "affine" up front at src/models/deepseek.rs:690 with a message that names this exact failure. The residual hole is that ModelArgs::quantization_mode() returns "affine" when the quantization block is absent entirely. Fixing it properly means giving the local SwitchLinear::Quantized an Option biases and a mode field, not flipping the match. (MEDIUM, separate issue)
• src/models/kimi_linear.rs:169-190 (MultiLinear::forward) gates correctly on .scales alone and keeps Option biases, then passes a null biases pointer with a hardcoded "affine" mode, with no validate_quantization_biases at from_weights. That is an uncatchable abort rather than a load error, the same shape gpt_oss.rs:463 already fixed. Cheapest high-value follow-up. (MEDIUM, separate issue)
• Roughly 18 family-local MoE expert loaders require .biases and return a clean Err without it. They refuse a block-float export rather than mis-loading it, which is the safe direction, but each carries a non-Option biases and a hardcoded "affine" at gather_qmm. (LOW, one sweep)
• Four MLA kv_b_proj sanitizers (deepseek_v3.rs:1138, deepseek_v32.rs:920, kimi_linear.rs:1280, longcat_flash_ngram.rs:1096) gate on .scales then .unwrap() the biases, so a block-float plane panics instead of erroring. youtu_vl_lm_sanitize.rs:71 is the already-corrected sibling. (LOW)

Confirmed clean: switch_layers.rs:461, gpt_oss.rs:445, bailing_moe.rs:879, nemotron_h.rs:2187, the fused-QKV bias triple in layers.rs:2358, mlxcel-surgery, and the shape-only bias checks in gemma3n.rs / helium.rs.

Verification

• All stated requirements implemented
• No placeholder or mock code remaining
• Integrated into the project code flow (both MambaModel::from_weights and Mamba2Model::from_weights, which are the live entry points for ModelType::Mamba / ModelType::Mamba2 in src/loading/nonstandard.rs and src/loading/special.rs)
• Project conventions followed (use crate::models::gpt2::validate_embedding_table matches bailing_moe.rs:105 and gpt_neox.rs:69 exactly)
• Existing modules reused where applicable (UnifiedEmbedding::from_weights, validate_embedding_table, validate_quantized_packing, HOSTILE_QUANT_PARAMS)
• No unintended structural changes (one public constructor removed, which issue fix(models): mamba and mamba2 treat a scales-without-biases embedding as non-quantized #976 asked for; no renames, moves or API-surface churn beyond it)
• Tests pass

Remaining items

• The models/falcon-mamba-7b-4bit output-comparison acceptance criterion is still not executed, since no mamba checkpoint is present locally. Shape-level ground truth from the safetensors headers stands in for it, and the arithmetic above shows the affine path stores the same pair it stored before, but that is not an output comparison.

[inureyes]

Security and performance review

Scope: the 7-file diff at c68ffdae against main @ e1ea5c2f, read under a hostile-or-malformed-checkpoint threat model (config-declared integers and tensor shapes crossing the cxx bridge, where an MLX C++ throw is an uncatchable std::terminate rather than a Result).

No CRITICAL and no HIGH findings. Nothing auto-fixed; the branch is unchanged.

Config-declared integers reaching arithmetic: clean, and strictly stronger than the path it replaces

group_size / bits are bounded by validate_quantization_params (bits in [1, 32], group_size in [1, 2^20]) at the top of reconcile_quantization_layout, before any shape is read. The removed QuantizedEmbedding::new called that same bound but nothing else; the shared path adds validate_quantization_mode and validate_quantization_biases on top of it, and reconciles against the tensor shapes. Every multiply on the new path is checked_mul (infer_quantization_bits, infer_quantization_group_size, packed_in.checked_mul(32 / bits)) or widened to i64 (the MLX-mirror predicate in validate_quantized_packing). vocab_size and hidden_size are usize and reach only comparisons and one checked_mul. No overflow is reachable in debug or release.

One property worth recording since two more families now exercise it: the block-float branch of the reconciler derives group_size = in_features / num_groups without re-bounding it against MAX_QUANT_GROUP_SIZE. It is safe, because num_groups * derived_gs == in_features == packed_in * (32 / bits) is already checked_mul-bounded to i32, so MLX's scales.shape(-1) * group_size in C++ int cannot overflow. Pre-existing shared-loader behavior, not introduced here.

Panic reachability: clean, at parity with the old path

The new path is contains_key then get().map(ffi::copy).ok_or_else() then a let [rows, cols] = shape.as_slice() else refutation then length-guarded slicing. The *shape.last().unwrap() inside the reconciler is guarded by its own is_empty() early return, and validate_quantized_packing rejects a rank mismatch before it slices. The only expect on the path is UnifiedEmbedding::weight(), which asserts a std::make_unique result, is pre-existing, and is shared with the four families already running this guard. A malformed weight map produces an Err, not an abort.

Out-of-bounds gather: closed for these two families, and this is the main security gain

Confirmed against the vendored MLX source rather than assumed. Both ffi::embedding and ffi::quantized_embedding are mlx::core::take(..., 0) (mlx_cxx_bridge.cpp:2980,2999), and the gather resolves an index with offset_neg_idx(idx, size) = (idx < 0) ? idx + size : idx (mlx/backend/cpu/indexing.cpp:19), which has no upper bound. validate_embedding_table rejects vocab_size > weight.shape[0], and on the quantized branch validate_quantized_packing additionally requires scales.shape[..rank-1] == weight.shape[..rank-1], so the second take(scales, flat_indices, 0) inside quantized_embedding is bounded by the same row count.

No bypass exists. MambaModel::from_weights and Mamba2Model::from_weights are the only construction sites (src/loading/special.rs:149,156), the guard runs unconditionally on both immediately after the embedding loads, and neither model type can be built by struct literal from outside its module.

Pre-existing and global, stated for completeness rather than as a finding against this PR: config.vocab_size bounds sampled tokens only indirectly (sampling clamps to the logits width, which for a tied embedding is the table row count). Prompt token ids come from the tokenizer, so a checkpoint whose tokenizer vocabulary exceeds its table rows is still an out-of-bounds read no matter what config.json declares. That is equally true of Gpt2, GptNeoX, GptBigCode and BailingMoe.

QuantizedEmbedding construction after the removal: clean

Two struct literals exist and both are inside layers.rs: from_weights_with_mode (post-validation) and clone_shared (copies already-validated fields). There is no QuantizedEmbedding { .. } anywhere else in the workspace and no remaining reference to QuantizedEmbedding::new outside one doc comment that explains its removal. The pub fields remain a theoretical bypass, which the doc comment on QuantizedMultiLinear::new now states explicitly.

Performance: the PR body's claim is correct

ffi::copy is mlx::core::copy (mlx_cxx_bridge.cpp:825), which builds a lazy Copy node (mlx/ops.cpp:303). Copy::eval is out.copy_shared_buffer(inputs[0]) (mlx/backend/common/common.cpp:57), reached by both eval_cpu and eval_gpu. It is a buffer alias, not an element copy. For falcon-mamba-7b-4bit that is roughly 133 MB of packed table plus two 16.6 MB planes that are aliased rather than duplicated; the delta is three array handles.

Peak load memory is unchanged for a second reason as well: weights was already retained for the whole of from_weights before this change, since every UnifiedLinear::from_weights(&weights, ...) borrows it. Swapping the embedding's three remove calls for get does not extend the map's lifetime by a single statement.

validate_embedding_table is O(rank) shape arithmetic. array_shape reads arr.inner.shape() metadata (mlx_cxx_bridge.cpp:300), which forces no evaluation and touches no element data. There is no data-dependent scan anywhere in the guard.

MEDIUM

The vocab_size clause of the new guard has no test for either family. The three new tests per family cover the block-float regression, both prefix spellings, and the hidden_size width clause, but nothing exercises claimed_rows > rows. That is the clause that closes the out-of-bounds gather, and it is the one place a caller can silently mis-wire the guard by passing the wrong config field into the claimed_rows slot. Every peer family that adopted this guard shipped exactly such a test: gpt_neox_tests.rs:1191 from_weights_rejects_a_config_that_overstates_the_token_table, gpt2_tests.rs:554, gpt_bigcode_tests.rs:558, helium_tests.rs:544. The shared helper is therefore covered; what is not covered is that Mamba and Mamba2 wire config.vocab_size into it. With the fixtures this PR already hoisted, mirroring the GptNeoX test is about ten lines per family: insert a [3, 8] float table under backbone.embeddings.weight, assert the load fails naming vocab_size, then assert a padded [64, 8] table still loads.

LOW

The prefix-mixing hole is closed too, and the PR body does not claim it. The old code resolved weight, scales and biases through three independent or_else chains, so a checkpoint carrying backbone.embeddings.weight alongside model.embed_tokens.scales would have built one QuantizedEmbedding out of two unrelated tables. Resolving a single prefix and passing it to the shared loader makes all three planes come from the same table by construction. Worth a line in the body, since it is a second latent defect this fixes.

Mamba and Mamba2 still have no config-scalar guard. Unlike Gpt2, GptBigCode, GptNeoX and BailingMoe, which all call args.validate() before touching the checkpoint, these two validate nothing. Vec::with_capacity(config.num_hidden_layers) (mamba.rs:576, mamba2.rs:805) is an allocation sized straight from config.json, and intermediate_size = num_heads * head_dim then conv_dim = intermediate_size + 2 * n_groups * ssm_state_size (mamba2.rs) are unchecked usize products that wrap in release and panic in debug. Pre-existing, untouched by this PR, and narrower than before now that the embedding surface is guarded. Natural follow-up issue rather than a change request here.

Verdict

Approve on security and performance grounds. The change removes a real out-of-bounds-read and garbage-output path for two families, replaces a hand-built constructor that could only express affine with the shared loader that reconciles and bounds, and pays for it with a lazy buffer alias that costs no measurable memory at load. The only thing I would ask for before merge is the vocab_size rows test, and that is a coverage gap rather than a defect.

[inureyes]

PR Finalization Complete

Summary

• Tests: added mamba_rejects_a_config_that_overstates_vocab_size (src/models/mamba_tests.rs) and mamba2_rejects_a_config_that_overstates_vocab_size (src/models/mamba2_tests.rs), closing the MEDIUM coverage gap flagged in the security review. Both drive MambaModel::from_weights / Mamba2Model::from_weights with a [3, 8] table under backbone.embeddings.weight against the existing vocab_size: 8 fixture config, assert the load fails naming vocab_size, then assert a [64, 8] padded table still loads. Both reuse the fixtures this PR already hoisted (embedding_config, embedding_test_base_weights) and assert on the load Result only, no forward pass.
• Documentation: none needed. docs/adding-models.md already covers the QuantizedEmbedding::new removal and does not enumerate per-clause test coverage, so there is nothing for this addition to update.
• Lint/Format: no issues found; nothing to fix.

Validation

• cargo fmt --all -- --check: clean
• cargo check --lib --tests: clean
• cargo clippy --lib --tests -- -D warnings: clean
• cargo test --lib models::mamba: 16 passed, 2 pre-existing ignored (serial-MLX conv-state tests)
• Confirmed both new tests fail cleanly (not abort) when the claimed_rows > rows clause in validate_embedding_table is temporarily disabled, then reverted before committing.

Commit 118ddb9f on top of c68ffdae.