feat(models): add Ant Group Ling / Bailing MoE (bailing_moe) text model support

docs/supported-models.md

@@ -25,6 +25,7 @@ Implemented model families include:
 - Mixtral and other MoE families
 - DeepSeek v1 / v2 / v3 / v3.2
 - dots.llm1 (`dots1`, rednote: a DeepSeek-V3-style Mixture-of-Experts without MLA. Standard multi-head attention with per-head Q/K RMSNorm, a dense first layer (`first_k_dense_replace`), then sigmoid-routed experts that select on `gate.weight` logits plus an `e_score_correction_bias`, with a single always-on shared expert. Validated against `mlx-community/dots.llm1.inst-mixed-4-6bit`, a mixed 4/6-bit export whose `v_proj` and `down_proj` tensors are 6-bit while the rest are 4-bit; the unified loaders detect the per-tensor bit width from shape.)
+- Ling / Bailing MoE (`bailing_moe`, Ant Group: a DeepSeek-shaped sparse decoder, RMSNorm before attention and before the FFN, grouped-query attention with RoPE, `num_experts` routed SwiGLU experts selected per token, and one always-on shared expert, so the expert machinery is the shared `switch_layers` path unchanged. What is not shared with DeepSeek is where the port's work went. The checkpoint uses GPT-NeoX naming rather than Llama naming: the token table is `model.word_embeddings`, the output projection is `attention.dense`, and Q, K and V arrive fused in one `attention.query_key_value` matrix that on `inclusionAI/Ling-lite-1.5` is `[3072, 2048]`, which is `(16 + 2 * 4) * 128` for 16 query heads and 4 KV heads of width 128; the split offsets are therefore `(2048, 2560)`, the uneven GQA form rather than an even three-way split, and MLX's `slice` clamps an out-of-range stop instead of throwing, so the even split would produce tensors of exactly the right shape assembled from the wrong channels. The router is renamed by upstream mlx-lm's `sanitize` from `mlp.gate.weight` to `mlp.gate.gate_proj.weight`, which collides with the `gate_proj` every routed expert also has; the loader resolves the router from two fully anchored keys instead of matching a suffix, so a raw `inclusionAI` export and an mlx-lm conversion both load and neither can read an expert's SwiGLU gate as the router. The raw checkpoint stores each routed expert separately at `mlp.experts.{e}.{gate,up,down}_proj`, 64 experts times 3 projections times 28 layers on Ling-lite; the loader joins those individually-keyed tensors into `SwitchLinear`'s stacked `[num_experts, out, in]` layout at load, checking every expert index rather than only index 0 so a checkpoint missing a late expert cannot register a short stack the router can index past, and an mlx-lm conversion that already ships the pre-stacked `mlp.switch_mlp.{proj}` form loads unchanged on the same path. Three routing details are load-bearing and invisible downstream: the correction bias is added to a copy of the scores used for SELECTION only, and the returned weights are gathered from the unbiased scores (gathering from the biased copy leaves the output finite and plausible while misweighting every routed contribution); `norm_topk_prob` is applied only when `num_experts_per_tok > 1` and divides by `sum + 1e-20` rather than a bare sum; and `score_function` defaults to `"softmax"`, not to DeepSeek-V3's sigmoid, so reusing the DeepSeek gate as-is would change the routing distribution on every token. `moe_router_enable_routed_scaling` is dead code upstream (`BailingMoeGate` stores it and never reads it, and `group_expert_select` ends with an unconditional multiply), and this port mirrors upstream so a checkpoint decodes here as it does under the reference; the loader prints a diagnostic naming both fields for the one combination where the two readings diverge, a non-unit `routed_scaling_factor` with the flag false. The shared expert is ONE wide MLP of `(moe_shared_expert_intermediate_size or moe_intermediate_size) * num_shared_experts`, which is `1408 * 2 = 2816` on Ling-lite, added to the routed mixture at a fixed weight of 1.0 and never packed into the switch tensors; there is no per-shared-expert axis anywhere in the checkpoint. Grouped top-k routing (`n_group` groups scored by the sum of their top two experts, `topk_group` kept) and the router expert bias are implemented and unit-tested, but no published Bailing checkpoint reaches either: Ling-lite declares neither field and both default to off, so a passing real-checkpoint run says nothing about them. `norm_head` is live upstream despite appearing vestigial, and is implemented as an axis-0 L2 normalization of `lm_head.weight` in float32 with a `1e-7` epsilon cast back to the weight's own dtype; `norm_softmax` is declared by upstream's `ModelArgs` and never read, does not appear in the vendored `modeling_bailing_moe.py`, and is not a named parameter of the vendored config, so a config setting it true is rejected at load rather than parsed and ignored. Loading also rejects an architecture scalar that could size an allocation, divide, or truncate through an `as i32` cast (zero checks precede divisibility checks, since `0.is_multiple_of(0)` is true in Rust); routing parameters that would put `argpartition(kth = n_group - topk_group - 1)` out of range or reshape the score row into unequal groups, which upstream checks neither of; a rotary width that is zero, odd or wider than the head, and a non-positive or non-finite `rope_theta` or `rms_norm_eps`; a `rope_scaling` block other than an absent, empty or `"default"` one, since upstream threads it into `initialize_rope` while this loader always builds the plain rotation; a `quantization` block whose `bits` falls outside `1..=32` or whose `group_size` is non-positive; the output head's input width against `hidden_size`, the one axis nothing else in the loader checked, with its row count checked exactly against upstream's own strict `nn.Linear(hidden_size, vocab_size)` load and the check skipped only when `tie_word_embeddings` ships no separate `lm_head` tensor; a quantized token table's own input width, reconstructed the same way, since the shared `validate_embedding_table` checks the row count unconditionally but the width only on the unquantized path; and any tensor whose real shape disagrees with the config, on both axes, per expert index rather than only the first (a checkpoint missing a late expert would otherwise register a short stack the router can index past), and on the quantized path, where packing compresses the input axis only so the input width is recovered from the scales the way MLX recovers it, `scales.shape(-1) * group_size`. Each of those is a load-time rejection because an MLX C++ exception crossing the cxx bridge is an uncatchable `std::terminate` at the first forward pass rather than a Rust error. Validation checkpoint: `inclusionAI/Ling-lite-1.5` (16.8B total / 2.75B activated, raw bf16, 5603 tensors, no mlx-community conversion needed). A token-exact comparison against the mlx-lm reference needs `--no-chat-template` on both sides, because this checkpoint does ship a chat template embedded in `tokenizer_config.json` while the reference harness feeds the raw prompt.)
 - Mistral 4 (`mistral4`, Mistral Small 4: a DeepSeek-V3-style Multi-Latent Attention decoder with compressed query and KV LoRA projections (`q_lora_rank` 1024, `kv_lora_rank` 256), split rope/nope query-key head dims, and a separate value head dim, paired with a softmax-routed Mixture-of-Experts (128 routed experts plus one always-on shared expert, 4 active per token, `norm_topk_prob`) and Llama-4 position-dependent attention scaling. The only public checkpoint is the Mistral Small 4 119B vision model, whose `text_config.model_type` is `mistral4`; mlxcel detects it as a Mistral 3 VLM (Pixtral vision tower) on the Mistral4 text backbone. Validated against `mlx-community/Mistral-Small-4-119B-2603-4bit` for both text and image-plus-text.)
 - Cohere / Cohere2
 - Command MoE (`cohere2_moe`, Cohere: the Cohere2 parallel-residual backbone (interleaved sliding/global attention, conditional RoPE, logit scaling, tied embeddings) with the dense FFN replaced by a sparse Mixture-of-Experts. The router scores on `mlp.gate` logits activated in f32 by sigmoid (default) or softmax, takes the top-k, and gathers the activated values at the selected experts (not a fresh softmax over the k logits), with optional `norm_topk_prob` renormalization clamped by `1e-12`. Always-on shared experts (`moe_num_shared_experts`) run a dense SwiGLU on the same input and combine by average (default `(y + y_s) / 2`) or sum. An optional dense-FFN prefix (`first_k_dense_replace`, width `prefix_dense_intermediate_size`) and an optional RMSNorm mode (`rms_norm_eps`, else LayerNorm) round out the deltas versus dense `cohere2`. The full runtime (detection, metadata, loaded-model dispatch, CLI and server generation) is wired; real-checkpoint validation is pending a public `cohere2_moe` checkpoint.)

src/distributed/tensor_parallel/inference.rs

@@ -202,6 +202,14 @@ fn fallback_architecture(model_type: ModelType) -> &'static str {
         ModelType::HunyuanMoe => "hunyuan_moe",
         ModelType::HunyuanV1Dense => "hunyuan_v1_dense",
         ModelType::MiMo => "mimo",
+        // Bailing MoE is not shardable by the generic transformer plan: its
+        // attention projection is the fused `attention.query_key_value` rather
+        // than separate q/k/v, and its routed experts live under
+        // `mlp.experts.{idx}` / `mlp.switch_mlp`, neither of which the shard
+        // rules name. The refusal comes from `validate_supported_runtime`, whose
+        // `runtime_kind_for` match has no arm for `ModelType::BailingMoe` and so
+        // returns `None`; this arm only keeps the dispatch table total.
+        ModelType::BailingMoe => "bailing_moe",
         ModelType::Apertus => "apertus",
         ModelType::SeedOss => "seed_oss",
         ModelType::Granite => "granite",

src/loaded_model.rs

@@ -149,6 +149,7 @@ pub enum LoadedModel {
     HunyuanMoe(models::HunyuanMoeModel),
     HunyuanV1Dense(models::HunyuanV1DenseModel),
     MiMo(models::MiMoModel),
+    BailingMoe(models::BailingMoeModel),
     Apertus(models::ApertusModel),
     SeedOss(models::SeedOssModel),
     Granite(models::GraniteModel),
@@ -300,6 +301,7 @@ macro_rules! delegate_language_model {
             LoadedModel::HunyuanMoe(inner) => LanguageModel::$method(inner, $($arg),*),
             LoadedModel::HunyuanV1Dense(inner) => LanguageModel::$method(inner, $($arg),*),
             LoadedModel::MiMo(inner) => LanguageModel::$method(inner, $($arg),*),
+            LoadedModel::BailingMoe(inner) => LanguageModel::$method(inner, $($arg),*),
             LoadedModel::Apertus(inner) => LanguageModel::$method(inner, $($arg),*),
             LoadedModel::SeedOss(inner) => LanguageModel::$method(inner, $($arg),*),
             LoadedModel::Granite(inner) => LanguageModel::$method(inner, $($arg),*),

src/main.rs

@@ -1820,6 +1820,7 @@ const FAMILY_ORDER: &[&str] = &[
     "GLM",
     "ERNIE",
     "Hunyuan",
+    "Bailing",
     "Granite",
     "ExaOne",
     "Solar",
@@ -1848,6 +1849,8 @@ const FAMILY_ORDER: &[&str] = &[
     "GLM VLM",
     "Kimi VLM",
     "PaddleOCR VLM",
+    "MiniMax VLM",
+    "Step VLM",
     "Other VLM",
 ];
 

src/model_metadata.rs

@@ -189,6 +189,7 @@ macro_rules! for_each_model_registration {
             HunyuanMoe => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::HunyuanMoeModel::load, args: models::hunyuan_moe::ModelArgs, weight_builder: models::HunyuanMoeModel::from_weights, wrap: LoadedModel::HunyuanMoe } };
             HunyuanV1Dense => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::HunyuanV1DenseModel::load, args: models::hunyuan_v1_dense::ModelArgs, weight_builder: models::HunyuanV1DenseModel::from_weights, wrap: LoadedModel::HunyuanV1Dense } };
             MiMo => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::MiMoModel::load, args: models::mimo::ModelArgs, weight_builder: models::MiMoModel::from_weights, wrap: LoadedModel::MiMo } };
+            BailingMoe => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::BailingMoeModel::load, args: models::bailing_moe::ModelArgs, weight_builder: models::BailingMoeModel::from_weights, wrap: LoadedModel::BailingMoe } };
             Apertus => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::ApertusModel::load, args: models::apertus::ModelArgs, weight_builder: models::ApertusModel::from_weights, wrap: LoadedModel::Apertus } };
             SeedOss => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::SeedOssModel::load, args: models::seed_oss::ModelArgs, weight_builder: models::SeedOssModel::from_weights, wrap: LoadedModel::SeedOss } };
             Granite => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::GraniteModel::load, args: models::granite::ModelArgs, weight_builder: models::GraniteModel::from_weights, wrap: LoadedModel::Granite } };