fix(models): read the declared quantization block in deepseek::ModelArgs

[inureyes]

Summary

deepseek::ModelArgs declared only flat top-level group_size / bits keys and had no nested quantization field, so serde silently discarded the MLX-style "quantization": {"group_size": N, "bits": M, "mode": ...} block that every conversion tool emits. group_size() / bits() returned the hardcoded 64 / 4 regardless of what the checkpoint declared. The three DeepSeek-OCR loaders in src/loading/vlm_deepseekocr.rs copy the root block into language_config before deserializing into that struct, so the inheritance ran, succeeded, and accomplished nothing on all three; DeepSeekModel::load dropped the block the same way for a plain deepseek config.json.

The declared pair feeds four places in this family, and they fail differently. The dense projections, the embedding and the OCR projector go through reconcile_quantization_layout, which trusts group_size and re-derives bits from the tensor shapes. The MoE expert planes never reach the reconciler: SwitchLinear::from_stacked_parts stores the pair verbatim and hands it to gather_qmm, which crosses the cxx bridge as UniquePtr<MlxArray> rather than Result, so a mismatch is an uncatchable std::terminate at the first routed forward pass. Reading the declared block is what closes both at the root.

What changed

• src/models/deepseek.rs gains pub quantization: Option<QuantizationArgs>. group_size() and bits() resolve nested block, then flat key, then family default, independently per key. A new quantization_mode() returns the declared mode (defaulting to "affine" when no block names one) through QuantizationArgs::get_mode, which validates against SUPPORTED_QUANTIZATION_MODES. DeepSeekModel::from_weights calls it and refuses anything but "affine": every plane this family builds is affine, and from_stacked_parts takes its quantized arm only when a .biases plane is present, so a block-float checkpoint would land on the Regular arm with a packed uint32 weight and be reinterpreted rather than refused. This mirrors the guard jamba.rs carries for the same reason.
• src/models/switch_layers.rs gains validate_expert_quantization_shapes, called from deepseek::SwitchLinear::from_stacked_parts. A declared pair now actually reaches gather_qmm for the first time on this family, and validate_expert_quantization_params deliberately stays a range rather than an allowlist (so mixed-precision exports keep loading), so an in-range pair can still describe a different tensor than the one on disk. This restates MLX's own packed_in * 32 == bits * num_groups * group_size precondition where the shapes are still in hand, so it can only refuse a load that was already going to abort. The other eighteen families on validate_expert_quantization_params's Used by: list are left alone to keep this one logical change; they are unaffected either way, since none of them was discarding a declared pair.
• src/models/config.rs: doc only. The shared QuantizationArgs is adopted rather than a sixtieth private Quantization struct being added (grep -c '^pub struct Quantization\b' over src/models/ gives 59, and the divergence between the deepseek.rs and deepseek_v2.rs copies is what produced this bug). get_mode's "no production callers yet" note is now false and is replaced with a Used by: line, and the struct doc records why it is the type to reach for.
• src/loading/vlm_deepseekocr.rs: the triplicated language_config fixup collapses into one deepseekocr_language_args(full_config, family) helper used by all three loaders, and the Unlimited-OCR ring-cache window moves into unlimited_ocr_sliding_window so it still reads the untouched full_config. One funnel makes the inheritance testable without a checkpoint on disk and keeps the three from drifting apart again.
• Tests: new src/loading/vlm_deepseekocr_tests.rs drives the helper once per loader label (inherited block honoured, a language_config-declared block still winning over the root, the in-tree checkpoints' declaration resolving to the previous value, no declaration keeping the defaults, a partial block resolving per key, an unparseable inherited mode refused, model_type supplied, sliding-window fallback and clamp). src/models/deepseek_tests.rs adds a nested-spelling run of the expert-plane guard and a shape-mismatch case, so the MoE half is covered at the loader that stores the pair, plus the plain-config.json nested read, the flat-spelling fallback, precedence, per-key fallback and the mode allowlist. 21 tests in these two modules.

Validation

Real checkpoints, --temp 0, one 1400x460 text image, main (39e2ceb60) versus this branch.

Checkpoint	Declared block	Matches the old hardcoded 64 / 4?	Output before vs after
models/deepseek-ocr-4bit	{"group_size": 64, "bits": 4, "mode": "affine"}	yes	byte-identical, OCR correct both sides
models/deepseek-ocr-2-4bit	{"group_size": 64, "bits": 4, "mode": "affine"}	yes	byte-identical, OCR correct both sides
models/deepseek-vl2-small-4bit	{"group_size": 64, "bits": 4}	n/a (loads through deepseek_v2, which already had the field)	byte-identical

Both DeepSeek-OCR checkpoints declare exactly the fallback pair, so the bug was latent rather than live for them and no output can change. To show the declared block now arrives, three symlinked copies of deepseek-ocr-4bit with only config.json edited (weights symlinked, nothing under models/ written):

Edited declaration	Before	After
"mode": "mxfp4"	loads, OCR identical to the honest checkpoint	clean load error naming the mode
{"group_size": 64, "bits": 8}	loads, OCR identical	clean load error naming the expert plane and the arithmetic (weight [64, 896, 160] with scales [64, 896, 20], 5120 against 10240)
{"group_size": 32, "bits": 8}	loads, OCR identical	quantized_matmul rejects it on the dense path

The middle row is the one the new shape guard converts: without it, that config reached gather_qmm and killed the process with SIGABRT. The last row still aborts, and honestly so: the pair is self-consistent under reconcile_quantization_layout's own arithmetic, and what rejects it is MLX comparing the expanded width against the activation width, which has no load-time equivalent because the activation is not in hand at load. That is the pre-existing #958 class, not something reading the block can close, and a candidate follow-up.

Notes

• The issue's "silently wrong output" prediction did not reproduce. It expected a real group_size-32 4-bit tensor read as group_size 64 to land on a self-consistent (64, 2) pair that MLX accepts and compute wrong numbers. It is self-consistent by the reconciler's arithmetic, but the expanded input width it implies no longer matches the activation width, and MLX's quantized_matmul throws on exactly that. On these checkpoints the observed consequence of a wrong pair is an abort, not silent corruption, on both halves. That does not change the fix (the values still have to be read) but it is worth recording, because it means the defect's practical blast radius was "aborts on anything but a 64 / 4 export", not "quietly degrades".
• fix(models): deepseek_v2 generates degenerate output on a real checkpoint while mlx-lm is correct on the same weights #991 does not interact. It was closed by fix(models): make the deepseek_v2 prefill causal #998, which made the deepseek_v2 prefill causal: a masking defect in a different module, not config plumbing. Nothing here touches deepseek_v2.rs, and models/deepseek-vl2-small-4bit is re-run above as the cross-family control.
• Two citation groups in the issue body have drifted as main moved. deepseek_v2.rs's nested field is at :117 (struct :121, accessors :184 / :191), not :114 / :118 / :181 / :188. In layers.rs, validate_quantization_params is at :1322, infer_quantization_bits at :1183, reconcile_quantization_layout at :1623 and the divergence tracing::warn! at :1751. Everything else checked out, including the 59 private Quantization copies, QuantizationArgs having zero production callers, the lib.rs:1092 / :1108 bridge citations, and both OCR checkpoints declaring exactly the fallback values.
• Supersedes fix(models): read the nested quantization block in deepseek::ModelArgs #989, which diagnosed the same defect correctly and independently. That PR reuses deepseek_v2::Quantization, whose fields are required and which carries no mode, so a partially declared block would become a hard deserialization failure and the declared mode would stay invisible; it also predates fix(core): bound the declared quantization mode string, not just group_size and bits #973 / fix(core): bound the declared quantization mode string at load #1006, which added the mode allowlist used here. It is a fork PR, so completing it in place would mean force-pushing over a contributor's branch.

Test plan

• cargo fmt --all -- --check
• cargo clippy --release --lib --tests --features metal,accelerate -- -D warnings
• cargo clippy -p mlxcel-core --release --all-targets --features metal,accelerate -- -D warnings
• cargo test --release --features metal,accelerate (77 targets, 5284 tests, 0 failures)
• cargo test --release -p mlxcel-core --features metal,accelerate --no-fail-fast -- --test-threads=1 (1411 tests, 1 failure: fused_moe_parity_tests::fused_moe_geglu_kernel_matches_references_gemma4_shape, the known deterministic failure tracked as fix(tests): fused MoE GeGLU parity threshold is tighter than the gather_qmm reference's own jitter #964; this PR changes no file under src/lib/mlxcel-core/)
• Real-checkpoint before/after on deepseek-ocr-4bit, deepseek-ocr-2-4bit and deepseek-vl2-small-4bit, plus the three edited-config demonstrations above

Closes #975