Repository navigation
feat(models): add Kuaishou Klear (klear) text model support #846
Description
Corrected against
references/mlx-lm(commit15b522f, 2026-07-11) on 2026-08-05.
Checkpoint confirmed and prepared
The original body lists the checkpoint as "unconfirmed". It exists: Kwai-Klear/Klear-46B-A2.5B-Instruct (32 layers, hidden 2048, 256 experts, 1 shared expert).
mlx-community/Klear-46B-A2.5B-Instruct-4bit is an empty repo (only .gitattributes); only the 3-bit conversion has weights. A 4-bit checkpoint was produced locally from the 92.5 GB bf16 source and is available at models/klear-46b-a2.5b-instruct-4bit (24 GB, 4.501 bits/weight).
Detection string is capital-K
The config declares "model_type": "Klear", not "klear". The detection arm must match the capitalized spelling (mlx-lm ships both Klear.py and klear.py, byte-identical, to paper over the same issue). The original acceptance criterion says "keyed on the klear model_type" and would miss every published checkpoint.
What the original body missed: the coefficient blend
The original body says the shared expert output is added to the routed mixture. Upstream does not add it. KlearSparseMoeBlock learns a 2-way blend:
self.coefficient = nn.Linear(args.hidden_size, 2)
...
coef = mx.softmax(self.coefficient(x), axis=-1, precise=True)
shared = self.shared_experts(x)
y = y_experts * coef[..., :1] + shared * coef[..., 1:]So there is an extra coefficient weight per MoE layer, and a plain add would misweight every token's output while leaving it finite and fluent.
Two further routing details, both load-bearing:
- Routing activation is sigmoid, not softmax:
routing_weights = mx.sigmoid(self.gate(x).astype(mx.float32)). expert_biasis added to a copy used for selection only; the returned scores are gathered from the unbiasedrouting_weights. Gathering from the biased copy is the classic silent-misweighting bug.norm_topk_probthen renormalizes the selected scores.
Confirmed as originally written
QK-RMSNorm on q/k, RMSNorm layer norms, rope_theta, and attention_bias plumbing are all accurate. Note the published Instruct checkpoint sets attention_bias: false, so that flag is not exercised by the validation run.
Original issue body (retained for history)
Summary
Add support for the Kuaishou Klear mixture-of-experts family to the mlxcel Rust inference runtime, ported from mlx-lm. Klear is a MoE decoder with attention_bias, routed top-k experts plus a shared expert, QK-RMSNorm, RMSNorm, and RoPE (rope_theta). It maps onto the existing qwen3_moe / deepseek MoE path plus the qwen3 QK-norm, with the attention-bias plumbing from a TeleChat3-style dense port.
Upstream reference
Public checkpoint
- Unconfirmed. Verify a Kuaishou/Kwai Klear checkpoint (and an MLX conversion) is publicly available before implementation. If none exists yet, keep this issue in the backlog and note the blocker; a real-checkpoint validation is required to close.
Architecture notes
- MoE decoder with routed top-k experts (
num_experts,num_experts_per_tok) plus a shared expert (n_shared_experts). attention_bias: bias terms on the attention projections when set.- QK-RMSNorm: RMSNorm on query/key (as in Qwen3).
- RMSNorm for the layer norms; RoPE with
rope_theta.
Implementation plan
- Reuse the MoE path from
src/models/qwen3_moe.rsand the deepseek MoE decoders (src/models/deepseek_v2.rs,src/models/deepseek_v3.rs) oversrc/models/switch_layers.rsfor routed top-k plus shared-expert combination. - Reuse the QK-RMSNorm construction from
src/models/qwen3.rs. - Plumb the
attention_biasflag into the attention projection construction so q/k/v/o pick up bias tensors at load time (same mechanism as the TeleChat3 port). - Confirm router normalization and shared-expert scaling match upstream before trusting parity.
Touchpoints & acceptance criteria
Follow docs/adding-models.md. Done only when the model loads and generates from a real checkpoint.
- Config struct + serde parse (
num_experts,num_experts_per_tok,n_shared_experts,attention_bias, QK-norm flag,rope_theta). -
from_weightsconstructs routed experts (viaswitch_layers.rs) plus the shared expert, QK-norm, and attention bias. -
sanitizehandles the checkpoint weight layout (stacked vs per-expert experts; q/k/v/o bias; standard bf16 to f16; key remapping). - Detection arm added in
src/models/detection.rskeyed on theklearmodel_type. - Registration in
src/model_metadata.rsviafor_each_model_registration!. - TP/distributed arch-string wired if Klear is a target for tensor/pipeline parallel.
- Unit tests in a
_tests.rsfile beside the implementation. -
docs/supported-models.mdupdated with the new family. - Real-checkpoint validation:
./target/release/mlxcel generateproduces coherent output andmlxcel listreports the architecture.
Effort
MEDIUM. A qwen3_moe/deepseek MoE derivative with QK-norm plus attention-bias; contingent on a public checkpoint being available for validation.
Activity
docs(models): document the modelsize / modeltype / arch label taxonomy (
feat(models): add Kuaishou Klear (Klear) text model support (lablup#1044
docs(models): document the modelsize / modeltype / arch label taxonomy (
feat(models): add Kuaishou Klear (Klear) text model support (lablup#1044