[inureyes]

Closes #846.

Adds the Klear sparse MoE family, validated on Kwai-Klear/Klear-46B-A2.5B-Instruct converted locally to 4-bit (24 GB, 32 layers, hidden 2048, 256 experts top-8 plus one shared, 1351 tensors). mlx-community/Klear-46B-A2.5B-Instruct-4bit is an empty repo — only the 3-bit conversion has weights — so the checkpoint was produced locally from the 92.5 GB bf16 source.

The backbone is what the issue described; three things are not

1. The model_type is capital-K "Klear". Upstream mlx-lm has to ship Klear.py and a byte-identical klear.py to cover both spellings, because it matches the config value as written. This tree needs only the lowercase detection arm, because get_model_type lowercases model_type before matching — without that normalization the arm would miss every published checkpoint. the_capitalized_model_type_reaches_the_detection_arm drives the real detection entry point rather than a helper, so the normalization is part of what is asserted, across "Klear", "klear" and "KLEAR".

2. The shared expert is BLENDED, not added. This is the one that makes the obvious reuse silently wrong. Every other shared-expert family in this tree adds the shared MLP's output to the routed mixture at a fixed weight of 1. KlearSparseMoeBlock learns a per-token 2-way softmax instead:

coef = mx.softmax(self.coefficient(x), axis=-1, precise=True)
y = y_experts * coef[..., :1] + shared * coef[..., 1:]

So there is an extra mlp.coefficient head per MoE layer — a [2, hidden] projection with a [2] bias, since upstream's nn.Linear(hidden, 2) defaults to bias=True — and a plain add misweights every token's output while leaving it finite and the text fluent.

Pinned by two tests rather than a shape check: the_shared_expert_is_blended_not_added zeroes the routed branch and injects a known constant into the shared one, so the blended reading (shared * coef[1]) and the added reading (shared) are numerically distinguishable, and asserts which is computed; the_blend_weights_come_from_the_coefficient_head_per_token varies the coefficient logits to confirm the weights are read from the checkpoint and that index 1, not 0, weights the shared branch.

3. Routing is sigmoid, and nothing scales the result. routing_weights = mx.sigmoid(self.gate(x).astype(mx.float32)), not softmax, with the expert_bias added to a copy used for selection only while the returned scores are gathered from the unbiased weights — the same silent-misweighting contract bailing_moe documents. norm_topk_prob then renormalizes.

That last point deserves naming. The published checkpoint declares routed_scaling_factor: 2.5, which looks exactly like the knob the neighbouring DeepSeek-derived families apply. Upstream's ModelArgs does not declare the field at all, so it is dropped on parse and KlearSparseMoeBlock scales nothing. This port mirrors upstream rather than guessing, and prints a diagnostic naming the field at load so the choice is visible instead of invisible. Applying it would multiply every routed contribution by 2.5 against the reference.

Validation

• A 5-token prefill reproduces the mlx-lm reference's top-5 logits exactly, values included:

  mlxcel  top5: 12095=19.0000 537=15.3750 3070=14.1250 264=13.8125 279=13.6250
  mlx-lm  top5: 12095=19.0000 537=15.3750 3070=14.1250 264=13.8125 279=13.6250
  

• Greedy decode is token-exact over 20 tokens:

  $ MLXCEL_FUSED_MOE=0 mlxcel generate -m models/klear-46b-a2.5b-instruct-4bit \
      -p "The capital of France is" -n 20 --temp 0 --no-chat-template
  The capital of France is Paris.
  Paris is known for its beautiful architecture, including the Eiffel Tower, the Lou
  

  identical to the reference.
• A reference-free causality check (position 0 of a longer prefill equals a one-token prefill).

One caveat worth knowing, and it is not about this port. That exactness holds with MLXCEL_FUSED_MOE=0. The default-on fused MoE decode kernel (issue #268) engages only at single-token decode, and it is not byte-identical here: with it enabled, a 1-token probe gives 6741=7.7812 387=7.6250 3550=7.4688 ... against the reference's 6741=7.5625 387=7.4062 8160=7.3125 ... (same top-1, reshuffled near-ties at a low-confidence position where the top logit is only 7.5), and greedy decode picks a different but equally coherent continuation. Disabling the kernel makes the 1-token logits match exactly too, which is what isolates it. The 5-token prefill, where the kernel does not engage, is exact either way. That kernel was validated byte-identical on qwen3-30b-a3b when it landed; this is a second data point saying that does not generalize, and it affects every MoE family in the tree rather than Klear specifically.

Other guards

The sparse schedule is layer_idx not in mlp_only_layers and num_experts > 0 and (layer_idx + 1) % decoder_sparse_step == 0, which on the published checkpoint (decoder_sparse_step: 1, empty mlp_only_layers) makes every one of the 32 layers sparse. A zero step is rejected because upstream would divide by it; an mlp_only_layers entry past the end of the stack is rejected too. Loading also rejects an architecture scalar that could size an allocation or divide (zero checks precede the divisibility check, since 0.is_multiple_of(0) is true in Rust), a num_experts_per_tok outside 1..=num_experts, an odd head width, a non-positive or non-finite rope_theta or rms_norm_eps, a non-default rope_scaling block (upstream does not declare the field and so drops one silently; rejecting is the safer reading), and any tensor whose real shape disagrees with the config — including an expert stack with fewer planes than num_experts and a missing mlp.coefficient.

attention_bias is false on the published checkpoint, so the bias path is never exercised by a real-model run; it is covered by a unit test that loads the block with biases and asserts that dropping them changes the output.

Tests

22 new unit tests, all checkpoint-free. cargo test --workspace --profile test-fast --features metal,accelerate green; cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check clean.