perf(moe): GeGLU fused decode-MoE; wire gemma4 (+13%) (#268 step 3b)

src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h

@@ -1380,6 +1380,19 @@ std::unique_ptr<MlxArray> fused_moe_expert_kernel(
     int32_t gu_bits, int32_t d_bits, int32_t group_size
 );
 
+// Same as fused_moe_expert_kernel but with GeGLU (gelu tanh approx) instead of
+// SwiGLU for the gate/up activation (gemma4 experts).
+std::unique_ptr<MlxArray> fused_moe_geglu_kernel(
+    const MlxArray& x,
+    const MlxArray& indices,
+    const MlxArray& gate_w, const MlxArray& gate_s, const MlxArray& gate_b,
+    const MlxArray& up_w,   const MlxArray& up_s,   const MlxArray& up_b,
+    const MlxArray& down_w, const MlxArray& down_s, const MlxArray& down_b,
+    const MlxArray& scores,
+    int32_t din, int32_t dff, int32_t k,
+    int32_t gu_bits, int32_t d_bits, int32_t group_size
+);
+
 // Fused MoE forward: gate + switch_mlp + score weighting + optional shared expert
 // Combines ~25 FFI calls into a single C++ function
 // Used by: NemotronH, NemotronNAS

src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp

@@ -895,7 +895,17 @@ namespace {
         g = simd_sum(g);
         u = simd_sum(u);
         if (lane == 0u) {
-            act_g[eslot * Dff + f] = (g / (1.0f + fast::exp(-g))) * u;
+            if (act == 1) {
+                // GeGLU (gelu tanh approx) * up — matches
+                // compiled_geglu_approx_activation (gemma4 experts).
+                float g3 = g * g * g;
+                float inner = 0.7978845608028654f * (g + 0.044715f * g3);
+                float gelu = 0.5f * g * (1.0f + precise::tanh(inner));
+                act_g[eslot * Dff + f] = gelu * u;
+            } else {
+                // SwiGLU (silu) * up.
+                act_g[eslot * Dff + f] = (g / (1.0f + fast::exp(-g))) * u;
+            }
         }
     )";
 
@@ -1074,15 +1084,18 @@ namespace {
     }
 }
 
-std::unique_ptr<MlxArray> fused_moe_expert_kernel(
-    const MlxArray& x,
-    const MlxArray& indices,
+namespace {
+// Shared body for the two fused decode-MoE FFIs. `act` selects the gate/up
+// activation: 0 = SwiGLU (silu), 1 = GeGLU (gelu tanh approx, gemma4). gate/up
+// use `gu_bits` (4/8), down uses `d_bits` (4/6/8); group_size shared.
+std::unique_ptr<MlxArray> run_fused_moe_two_kernel(
+    const MlxArray& x, const MlxArray& indices,
     const MlxArray& gate_w, const MlxArray& gate_s, const MlxArray& gate_b,
     const MlxArray& up_w,   const MlxArray& up_s,   const MlxArray& up_b,
     const MlxArray& down_w, const MlxArray& down_s, const MlxArray& down_b,
     const MlxArray& scores,
     int32_t din, int32_t dff, int32_t k,
-    int32_t gu_bits, int32_t d_bits, int32_t group_size
+    int32_t gu_bits, int32_t d_bits, int32_t group_size, int act
 ) {
     using namespace mlx::core;
     auto T = x.inner.dtype();
@@ -1100,14 +1113,14 @@ std::unique_ptr<MlxArray> fused_moe_expert_kernel(
     // 4-bit, down 6-bit), so each kernel gets its own `bits` template arg.
     std::vector<std::pair<std::string, mlx::core::fast::TemplateArg>> taA = {
         {"T", T}, {"K", k}, {"Din", din}, {"Dff", dff},
-        {"bits", gu_bits}, {"group_size", group_size},
+        {"bits", gu_bits}, {"group_size", group_size}, {"act", act},
     };
     std::vector<std::pair<std::string, mlx::core::fast::TemplateArg>> taB = {
         {"T", T}, {"K", k}, {"Din", din}, {"Dff", dff},
         {"bits", d_bits}, {"group_size", group_size},
     };
 
-    // A) gate/up + swiglu -> act_g[K, Dff] (f32 for the down GEMV).
+    // A) gate/up + activation -> act_g[K, Dff] (f32 for the down GEMV).
     auto& kA = get_moe_gateup_kernel().get();
     std::vector<array> inA = {
         astype(x.inner, T), astype(indices.inner, uint32),
@@ -1137,6 +1150,41 @@ std::unique_ptr<MlxArray> fused_moe_expert_kernel(
     auto summed = sum(rB[0], /*axis=*/0, /*keepdims=*/false);
     return std::make_unique<MlxArray>(std::move(summed));
 }
+}  // namespace
+
+// SwiGLU experts (qwen3_moe, dots.llm1, qwen3_next, ...).
+std::unique_ptr<MlxArray> fused_moe_expert_kernel(
+    const MlxArray& x,
+    const MlxArray& indices,
+    const MlxArray& gate_w, const MlxArray& gate_s, const MlxArray& gate_b,
+    const MlxArray& up_w,   const MlxArray& up_s,   const MlxArray& up_b,
+    const MlxArray& down_w, const MlxArray& down_s, const MlxArray& down_b,
+    const MlxArray& scores,
+    int32_t din, int32_t dff, int32_t k,
+    int32_t gu_bits, int32_t d_bits, int32_t group_size
+) {
+    return run_fused_moe_two_kernel(
+        x, indices, gate_w, gate_s, gate_b, up_w, up_s, up_b,
+        down_w, down_s, down_b, scores, din, dff, k, gu_bits, d_bits,
+        group_size, /*act=*/0);
+}
+
+// GeGLU experts (gemma4): gelu-tanh-approx(gate) * up.
+std::unique_ptr<MlxArray> fused_moe_geglu_kernel(
+    const MlxArray& x,
+    const MlxArray& indices,
+    const MlxArray& gate_w, const MlxArray& gate_s, const MlxArray& gate_b,
+    const MlxArray& up_w,   const MlxArray& up_s,   const MlxArray& up_b,
+    const MlxArray& down_w, const MlxArray& down_s, const MlxArray& down_b,
+    const MlxArray& scores,
+    int32_t din, int32_t dff, int32_t k,
+    int32_t gu_bits, int32_t d_bits, int32_t group_size
+) {
+    return run_fused_moe_two_kernel(
+        x, indices, gate_w, gate_s, gate_b, up_w, up_s, up_b,
+        down_w, down_s, down_b, scores, din, dff, k, gu_bits, d_bits,
+        group_size, /*act=*/1);
+}
 
 // Fused Mamba2 mixer forward for single-token decode.
 // Combines in_proj + conv1d + SSM kernel + MambaRMSNormGated + out_proj into one C++ call.

src/lib/mlxcel-core/src/lib.rs

@@ -1431,6 +1431,30 @@ mod ffi {
             group_size: i32,
         ) -> UniquePtr<MlxArray>;
 
+        /// Like `fused_moe_expert_kernel` but GeGLU (gelu tanh approx) instead
+        /// of SwiGLU for the gate/up activation (gemma4 experts).
+        #[allow(clippy::too_many_arguments)]
+        fn fused_moe_geglu_kernel(
+            x: &MlxArray,
+            indices: &MlxArray,
+            gate_w: &MlxArray,
+            gate_s: &MlxArray,
+            gate_b: &MlxArray,
+            up_w: &MlxArray,
+            up_s: &MlxArray,
+            up_b: &MlxArray,
+            down_w: &MlxArray,
+            down_s: &MlxArray,
+            down_b: &MlxArray,
+            scores: &MlxArray,
+            din: i32,
+            dff: i32,
+            k: i32,
+            gu_bits: i32,
+            d_bits: i32,
+            group_size: i32,
+        ) -> UniquePtr<MlxArray>;
+
         /// Fused gated-delta single-token decode step.
         /// Combines: decay → kv_mem → delta → state_update → output into one C++ call.
         /// Replaces ~26 FFI round-trips with 1.

src/models/gemma4.rs

@@ -462,6 +462,74 @@ impl SwitchGeGLU {
         }
     }
 
+    /// Single-token decode via the fused GeGLU MoE kernel (#268). Returns the
+    /// score-weighted expert sum `[hidden]`, or `None` (caller falls back to
+    /// `forward` + combine) for any unsupported config: non-affine, gate/up not
+    /// 4/8-bit or down not 4/6/8-bit, gate/up bits mismatch, group_size
+    /// mismatch, the Regular variant, or a non-single-token `x`.
+    fn forward_fused_kernel(
+        &self,
+        x: &MlxArray,
+        indices: &MlxArray,
+        scores: &MlxArray,
+    ) -> Option<UniquePtr<MlxArray>> {
+        let gate = self.gate_proj.quantized_parts()?;
+        let up = self.up_proj.quantized_parts()?;
+        let down = self.down_proj.quantized_parts()?;
+        if gate.bits != 4 && gate.bits != 8 {
+            return None;
+        }
+        if down.bits != 4 && down.bits != 8 && down.bits != 6 {
+            return None;
+        }
+        if gate.bits != up.bits
+            || gate.group_size != up.group_size
+            || gate.group_size != down.group_size
+        {
+            return None;
+        }
+        if gate.mode != "affine" || up.mode != "affine" || down.mode != "affine" {
+            return None;
+        }
+        let gw_shape = mlxcel_core::array_shape(gate.weight.as_ref().unwrap());
+        if gw_shape.len() != 3 {
+            return None;
+        }
+        let dff = gw_shape[1];
+        let din = gw_shape[2] * (32 / gate.bits);
+        if down.bits == 6 && dff % 16 != 0 {
+            return None;
+        }
+        let k = *mlxcel_core::array_shape(indices).last()?;
+        let x_elems: i32 = mlxcel_core::array_shape(x).iter().product();
+        if x_elems != din {
+            return None;
+        }
+        let x_flat = mlxcel_core::reshape(x, &[din]);
+        let idx_flat = mlxcel_core::reshape(indices, &[k]);
+        let sc_flat = mlxcel_core::reshape(scores, &[k]);
+        Some(mlxcel_core::fused_moe_geglu_kernel(
+            &x_flat,
+            &idx_flat,
+            gate.weight.as_ref().unwrap(),
+            gate.scales.as_ref().unwrap(),
+            gate.biases.as_ref().unwrap(),
+            up.weight.as_ref().unwrap(),
+            up.scales.as_ref().unwrap(),
+            up.biases.as_ref().unwrap(),
+            down.weight.as_ref().unwrap(),
+            down.scales.as_ref().unwrap(),
+            down.biases.as_ref().unwrap(),
+            &sc_flat,
+            din,
+            dff,
+            k,
+            gate.bits,
+            down.bits,
+            gate.group_size,
+        ))
+    }
+
     fn from_weights(
         weights: &WeightMap,
         prefix: &str,
@@ -653,8 +721,19 @@ impl Experts {
 
         let x_flat = mlxcel_core::reshape(x, &[b * s, h]);
         let indices_flat = mlxcel_core::reshape(top_k_indices, &[b * s, k]);
-        let expert_out = self.switch_geglu.forward(&x_flat, &indices_flat, h);
 
+        // Fused single-token decode GeGLU kernel (#268) behind MLXCEL_FUSED_MOE;
+        // otherwise SwitchGeGLU + weighted combine (also the kernel's fallback).
+        if b * s == 1
+            && std::env::var("MLXCEL_FUSED_MOE").is_ok()
+            && let Some(out) =
+                self.switch_geglu
+                    .forward_fused_kernel(&x_flat, &indices_flat, top_k_weights)
+        {
+            return mlxcel_core::reshape(&out, &[b, s, h]);
+        }
+
+        let expert_out = self.switch_geglu.forward(&x_flat, &indices_flat, h);
         let weights = mlxcel_core::reshape(top_k_weights, &[b * s, k, 1]);
         let weighted = mlxcel_core::multiply(&expert_out, &weights);
         let reduced = mlxcel_core::sum_axis(&weighted, -2, false);