diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index 7ee0bb6fd..41c7dda1c 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -1380,6 +1380,19 @@ std::unique_ptr 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 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 diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp index f78b651d8..b77059ed9 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp +++ b/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 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 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 fused_moe_expert_kernel( // 4-bit, down 6-bit), so each kernel gets its own `bits` template arg. std::vector> 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> 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 inA = { astype(x.inner, T), astype(indices.inner, uint32), @@ -1137,6 +1150,41 @@ std::unique_ptr fused_moe_expert_kernel( auto summed = sum(rB[0], /*axis=*/0, /*keepdims=*/false); return std::make_unique(std::move(summed)); } +} // namespace + +// SwiGLU experts (qwen3_moe, dots.llm1, qwen3_next, ...). +std::unique_ptr 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 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. diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index ef181f8ef..abae133bd 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -1431,6 +1431,30 @@ mod ffi { group_size: i32, ) -> UniquePtr; + /// 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; + /// 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. diff --git a/src/models/gemma4.rs b/src/models/gemma4.rs index 29451a373..41f8a5e45 100644 --- a/src/models/gemma4.rs +++ b/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> { + 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);