Fix LFM2 MoE routing to match the HF reference (sigmoid gating) (ml-explore#1354)

Lfm2Moe is a sigmoid-gated MoE (DeepSeek-V3 style): the router scores are
sigmoid(gate(x)); expert_bias is an aux-loss-free load-balancing bias used only
to select the top-k experts; the combination weights are gathered from the
unbiased sigmoid scores and then scaled by routed_scaling_factor.

The current implementation used softmax, folded expert_bias into the weights, and
omitted routed_scaling_factor, so MLX routing diverges from the HF Transformers
reference (and llama.cpp). This aligns it with the reference; routed_scaling_factor
is added to ModelArgs and defaults to 1.0.

Co-authored-by: Yi Zhong <207368749+vincentzed@users.noreply.github.com>

Author: ykhrustalev

mlx_lm/models/lfm2_moe.py

@@ -36,6 +36,7 @@ class ModelArgs(BaseModelArgs):
     conv_bias: bool
     conv_L_cache: int
     rope_theta: float = 1000000.0
+    routed_scaling_factor: float = 1.0
     rope_parameters: Optional[dict] = None
     full_attn_idxs: Optional[List[int]] = None
     layer_types: Optional[List[str]] = None
@@ -196,6 +197,7 @@ def __init__(self, args: ModelArgs):
         self.top_k = args.num_experts_per_tok
         self.norm_topk_prob = args.norm_topk_prob
         self.use_expert_bias = args.use_expert_bias
+        self.routed_scaling_factor = args.routed_scaling_factor
 
         self.gate = nn.Linear(dim, num_experts, bias=False)
         self.switch_mlp = SwitchGLU(dim, intermediate_size, num_experts)
@@ -206,18 +208,23 @@ def __call__(
         self,
         x: mx.array,
     ):
-        gates = self.gate(x).astype(mx.float32)
-        gates = mx.softmax(gates, axis=-1)
-
-        if self.use_expert_bias:
-            gates += self.expert_bias
+        # Sigmoid-gated routing (matches the HF Transformers Lfm2Moe reference):
+        # the expert_bias is used only to select the top-k experts; the routing
+        # weights are gathered from the unbiased sigmoid scores and then scaled by
+        # routed_scaling_factor.
+        routing_weights = mx.sigmoid(self.gate(x))
 
         k = self.top_k
-        inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:]
+        if self.use_expert_bias:
+            scores_for_routing = routing_weights.astype(mx.float32) + self.expert_bias
+            inds = mx.argpartition(scores_for_routing, kth=-k, axis=-1)[..., -k:]
+        else:
+            inds = mx.argpartition(routing_weights, kth=-k, axis=-1)[..., -k:]
 
-        scores = mx.take_along_axis(gates, inds, axis=-1)
+        scores = mx.take_along_axis(routing_weights, inds, axis=-1)
         if self.norm_topk_prob:
-            scores /= mx.sum(scores, axis=-1, keepdims=True) + 1e-20
+            scores = scores / (mx.sum(scores, axis=-1, keepdims=True) + 1e-6)
+        scores = scores * self.routed_scaling_factor
         scores = scores.astype(x.dtype)
 
         y = self.switch_mlp(x, inds)