perf(moe): wire lfm2 to the fused decode-MoE kernel

docs/benchmark_results/fused-moe-decode-kernel-design.md

@@ -207,14 +207,14 @@ the expected numerical consequence of the fusion.
 
 ### Models covered
 
-The fused single-token decode dispatch is wired into five model paths: qwen3_moe
+The fused single-token decode dispatch is wired into six model paths: qwen3_moe
 (Qwen3 MoE), qwen3_next (qwen3.5/3.6), dots.llm1 (mixed 4/6-bit), gemma4 (GeGLU),
-and qwen2_moe (qwen1.5-moe / Qwen2-MoE; migrated from its local `SwitchGLU` to the
+qwen2_moe (qwen1.5-moe / Qwen2-MoE; migrated from its local `SwitchGLU` to the
 shared one, which gained a per-expert stacking loader for the `experts.{idx}`
-checkpoint layout). Other MoE families reuse the shared `SwitchGLU` for the expert
+checkpoint layout), and lfm2 (LFM2-MoE; sigmoid-routed, optional expert_bias and
+norm_topk_prob). Other MoE families reuse the shared `SwitchGLU` for the expert
 matmul but were not wired with the fused decode dispatch, so they stay on
 `gather_qmm` regardless of `MLXCEL_FUSED_MOE`: mixtral, olmoe, minimax, phimoe,
-lfm2, and qwen3_vl_moe (it imports `SwitchGLU` from qwen3_moe but keeps its own
-non-fused MoE forward). Wiring those is a separate follow-up. nemotron-h's MoE runs
-through the separate C++ `fused_moe_forward` and is wired behind
-`MLXCEL_FUSED_MOE_RELU2` only.
+and qwen3_vl_moe (it imports `SwitchGLU` from qwen3_moe but keeps its own
+non-fused MoE forward). nemotron-h's MoE runs through the separate C++
+`fused_moe_forward` and is wired behind `MLXCEL_FUSED_MOE_RELU2` only.

src/models/lfm2.rs

@@ -516,8 +516,24 @@ impl Lfm2MoeSparseMoeBlock {
             scores = mlxcel_core::multiply_scalar(&scores, self.routed_scaling_factor);
         }
 
-        let expert_out = self.switch_mlp.forward(&x_flat, &topk_indices);
-        let result = moe_weighted_sum(&expert_out, &scores, mlxcel_core::array_dtype(&x_flat));
+        let result = {
+            let fused = if mlxcel_core::array_shape(&x_flat)[0] == 1
+                && crate::models::switch_layers::fused_moe_enabled()
+            {
+                self.switch_mlp
+                    .forward_fused_kernel(&x_flat, &topk_indices, &scores)
+                    .map(|out| mlxcel_core::reshape(&out, &[1, hidden_dim]))
+            } else {
+                None
+            };
+            match fused {
+                Some(out) => out,
+                None => {
+                    let expert_out = self.switch_mlp.forward(&x_flat, &topk_indices);
+                    moe_weighted_sum(&expert_out, &scores, mlxcel_core::array_dtype(&x_flat))
+                }
+            }
+        };
 
         if orig_shape.len() > 2 {
             mlxcel_core::reshape(&result, &orig_shape)