perf(moe): route granite-4.0-h single-token MoE through the fused kernel

docs/benchmark_results/rocm-decode-profile-gfx1151-2026-09-30.md

@@ -95,7 +95,7 @@ Which of those roles each port would actually take over depends on whether mlxce
 
 - **#2063** (`fused_add_rms_norm`, `fused_rope_qk_append`): both paths ship off on every backend (`FUSED_ADD_RMSNORM_DEFAULT` and `FUSED_ROPE_APPEND_DEFAULT` are `false`, `layers.rs:808` and `:820`, after #905 measured no decode win on Metal), and only `llama3.rs` (with `gemma.rs` and `iquestloopcoder.rs`, not profiled) calls them. With the opt-in on, Llama 3.1 would reach only the post-attention join (`llama3.rs:1201`): its `rope_scaling` builds a frequency table, which routes around the RoPE kernel (`llama3.rs:669`).
 - **#2064** (`gumbel_max_sample`, `rejection_sample`): greedy argmax calls neither. A sampled run would reach the whole sampler tail (both kernels are on by default where ported) except the few logit-bias dispatches `--ignore-eos` adds; the table counts the whole tail, so its #2064 figures are slightly high.
-- **#2065** (`moe_gateup`, `moe_down`): reached by `qwen3_moe` (`qwen3_moe.rs:223`, single-token decode). Not by granite, whose `block_sparse_moe` calls `SwitchGLU::forward` and never `forward_fused_kernel`, and not by Nemotron-H, whose `fused_moe_forward` default branch is `gather_qmm`; its kernel branch needs `MLXCEL_FUSED_MOE_RELU2` and `moe_fc1_relu2` (#2069).
+- **#2065** (`moe_gateup`, `moe_down`): reached by `qwen3_moe` (`qwen3_moe.rs:223`, single-token decode). At this profile's source (`929c80ab`) not by granite, whose `block_sparse_moe` called `SwitchGLU::forward` and never `forward_fused_kernel`; since issue #2245 (2026-10-11) granite's `block_sparse_moe` routes single-token decode through `forward_fused_kernel`, so the granite rows in the table above predate the wiring. Not by Nemotron-H, whose `fused_moe_forward` default branch is `gather_qmm`; its kernel branch needs `MLXCEL_FUSED_MOE_RELU2` and `moe_fc1_relu2` (#2069).
 - **#2067** (`ssm_update_kernel`): reached by both hybrids, since `ssm_kernel_available()` gates every single-token SSD step (`granitemoehybrid.rs:440`, `nemotron_h.rs:537` and `:635`).
 - **#2068** (paged attention): not on this path at all. The bench decodes one sequence into a dense `KVCache`; no paged kernel or paged graph fallback appears in any trace.
 

src/lib/mlxcel-surgery/src/config_yaml_compat_tests.rs

@@ -55,8 +55,7 @@ fn non_finite_factors_reach_operation_validation() {
         let yaml =
             format!("version: 1\noperations: [{{op: scale, pattern: '*', factor: {value}}}]");
         let error = parse_config_str(&yaml, None)
-            .err()
-            .expect("non-finite rejected")
+            .expect_err("non-finite rejected")
             .to_string();
         assert!(
             error.contains("scale") && error.contains("finite"),
@@ -72,8 +71,7 @@ fn non_finite_coefficients_reach_add_validation() {
             "version: 1\noperations: [{{op: add, pattern: '*', source: donor.safetensors, alpha: {value}}}]"
         );
         let error = parse_config_str(&yaml, None)
-            .err()
-            .expect("non-finite rejected")
+            .expect_err("non-finite rejected")
             .to_string();
         assert!(
             error.contains("(add)") && error.contains("alpha must be finite"),
@@ -89,8 +87,7 @@ fn non_finite_ratios_reach_interpolate_validation() {
             "version: 1\noperations: [{{op: interpolate, pattern: '*', source_a: a.safetensors, source_b: b.safetensors, method: lerp, ratio: {value}}}]"
         );
         let error = parse_config_str(&yaml, None)
-            .err()
-            .expect("non-finite rejected")
+            .expect_err("non-finite rejected")
             .to_string();
         assert!(
             error.contains("(interpolate)") && error.contains("ratio must be a finite number"),

src/models/granitemoehybrid.rs

@@ -747,7 +747,7 @@ impl GraniteMoeHybridAttention {
 
 // Dense SwiGLU MLP (`GraniteMoeHybridMLP`).
 
-struct GraniteMoeHybridMLP {
+pub(crate) struct GraniteMoeHybridMLP {
     gate_proj: UnifiedLinear,
     up_proj: UnifiedLinear,
     down_proj: UnifiedLinear,
@@ -764,15 +764,15 @@ impl GraniteMoeHybridMLP {
 
 // Sparse MoE block (`block_sparse_moe`): softmax-over-top-k router + SwitchGLU.
 
-struct GraniteMoeHybridMoE {
-    router: UnifiedLinear,
-    switch_mlp: SwitchGLU,
+pub(crate) struct GraniteMoeHybridMoE {
+    pub(crate) router: UnifiedLinear,
+    pub(crate) switch_mlp: SwitchGLU,
     top_k: i32,
     num_experts: i32,
 }
 
 impl GraniteMoeHybridMoE {
-    fn forward(&self, x: &MlxArray) -> UniquePtr<MlxArray> {
+    pub(crate) fn forward(&self, x: &MlxArray) -> UniquePtr<MlxArray> {
         let orig_shape = mlxcel_core::array_shape(x);
         let hidden_dim = orig_shape[orig_shape.len() - 1];
 
@@ -797,8 +797,34 @@ impl GraniteMoeHybridMoE {
         let topk_logits = mlxcel_core::astype(&topk_logits, mlxcel_core::dtype::FLOAT32);
         let scores = mlxcel_core::softmax(&topk_logits, -1);
 
-        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));
+        // Apply routed experts. Fused single-token decode kernel (#268) on by
+        // default; MLXCEL_FUSED_MOE=0 forces the proven SwitchGLU +
+        // moe_weighted_sum path (also the automatic fallback when the kernel
+        // does not support the config). Granite routes through softmax over
+        // ONLY the top-k logits, which equals softmax-over-all then
+        // renormalize over the top-k (norm_topk_prob), so `scores` already
+        // carries the full combine weights the kernel needs; the launcher
+        // casts scores to f32 itself. The shared expert joins outside this
+        // block (the decoder layer sums it with the MoE output) and is
+        // unchanged (issue #2245).
+        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)
@@ -810,7 +836,7 @@ impl GraniteMoeHybridMoE {
 
 // Shared dense expert (`shared_mlp`): fused gate/up SwiGLU, always summed with MoE.
 
-struct GraniteMoeHybridSharedMLP {
+pub(crate) struct GraniteMoeHybridSharedMLP {
     input_linear: UnifiedLinear,
     output_linear: UnifiedLinear,
     shared_intermediate: i32,
@@ -833,7 +859,7 @@ impl GraniteMoeHybridSharedMLP {
 
 // Per-layer feed-forward: dense MLP, or MoE + shared MLP.
 
-enum FeedForward {
+pub(crate) enum FeedForward {
     Dense(GraniteMoeHybridMLP),
     /// Dense feed-forward kept in the fused `shared_mlp` form (a single
     /// `input_linear` producing gate+up). Used for quantized dense checkpoints,
@@ -1334,7 +1360,7 @@ impl GraniteMoeHybridModel {
 
 /// Build the per-layer feed-forward: dense `mlp`, or `block_sparse_moe` +
 /// `shared_mlp` when `use_moe`.
-fn build_feed_forward(
+pub(crate) fn build_feed_forward(
     weights: &WeightMap,
     args: &ModelArgs,
     prefix: &str,

src/models/granitemoehybrid_tests.rs

@@ -26,7 +26,22 @@
 //! `mlx-community/granite-4.0-h-350m-4bit` checkpoint and is not exercised here
 //! (no Metal device is assumed in unit tests).
 
-use super::granitemoehybrid::ModelArgs;
+use super::granitemoehybrid::{FeedForward, GraniteMoeHybridMoE, ModelArgs, build_feed_forward};
+use crate::models::switch_layers::{
+    HONEST_EXPERT_GROUP_SIZE, insert_honest_affine_swiglu_experts, moe_weighted_sum,
+};
+use mlxcel_core::weights::WeightMap;
+
+fn values(len: usize, start: f32, step: f32) -> Vec<f32> {
+    (0..len).map(|idx| start + idx as f32 * step).collect()
+}
+
+fn insert_tensor(weights: &mut WeightMap, name: &str, values: Vec<f32>, shape: &[i32]) {
+    weights.insert(
+        name.to_string(),
+        mlxcel_core::from_slice_f32(&values, shape),
+    );
+}
 
 /// A trimmed `granite-4.0-h-350m` config (the dense-MLP hybrid). Values mirror
 /// the real checkpoint, including the full 32-entry `layer_types` interleave.
@@ -198,6 +213,213 @@ fn granite_hybrid_quantization_read_from_config() {
     assert_eq!(bf16.bits(), 4);
 }
 
+/// Hidden width of the fused-kernel geometry below: one quantization group per
+/// input row, so one token carries 64 elements. Reuses the honest affine 4-bit
+/// expert triple the Qwen3MoE fused-geometry tests use (issue #1884).
+const FUSED_HIDDEN: i32 = HONEST_EXPERT_GROUP_SIZE;
+const FUSED_EXPERTS: i32 = 2;
+
+/// A granite `GraniteMoeHybridMoE` over two honest affine 4-bit bf16 experts of
+/// intermediate width `dff`, loaded through the family's own `build_feed_forward`,
+/// so the fused-kernel decision under test runs on the exact type and stored
+/// triple a converted granite-4.0-h checkpoint produces.
+fn fused_moe_block() -> GraniteMoeHybridMoE {
+    let mut weights = WeightMap::new();
+    insert_honest_affine_swiglu_experts(
+        &mut weights,
+        "model.layers.0.block_sparse_moe.switch_mlp",
+        FUSED_EXPERTS,
+        FUSED_HIDDEN,
+        FUSED_HIDDEN,
+    );
+    // A non-quantized router and shared expert, so the declared pair can only
+    // reach the routed experts.
+    insert_tensor(
+        &mut weights,
+        "model.layers.0.block_sparse_moe.router.layer.weight",
+        values((FUSED_EXPERTS * FUSED_HIDDEN) as usize, -0.2, 0.004),
+        &[FUSED_EXPERTS, FUSED_HIDDEN],
+    );
+    insert_tensor(
+        &mut weights,
+        "model.layers.0.shared_mlp.input_linear.weight",
+        values((2 * FUSED_HIDDEN * FUSED_HIDDEN) as usize, 0.01, 0.002),
+        &[2 * FUSED_HIDDEN, FUSED_HIDDEN],
+    );
+    insert_tensor(
+        &mut weights,
+        "model.layers.0.shared_mlp.output_linear.weight",
+        values((FUSED_HIDDEN * FUSED_HIDDEN) as usize, 0.02, 0.002),
+        &[FUSED_HIDDEN, FUSED_HIDDEN],
+    );
+    let args: ModelArgs = serde_json::from_value(serde_json::json!({
+        "model_type": "granitemoehybrid",
+        "vocab_size": 8,
+        "hidden_size": FUSED_HIDDEN,
+        "intermediate_size": FUSED_HIDDEN,
+        "num_hidden_layers": 1,
+        "num_attention_heads": 1,
+        "num_key_value_heads": 1,
+        "rms_norm_eps": 1e-5,
+        "embedding_multiplier": 1.0,
+        "attention_multiplier": 1.0,
+        "logits_scaling": 1.0,
+        "residual_multiplier": 1.0,
+        "layer_types": ["attention"],
+        "num_local_experts": FUSED_EXPERTS,
+        "num_experts_per_tok": FUSED_EXPERTS,
+        "shared_intermediate_size": FUSED_HIDDEN,
+        "quantization": { "group_size": 64, "bits": 4 },
+    }))
+    .expect("fused-geometry granite config must parse");
+    let ff = build_feed_forward(&weights, &args, "model.layers.0")
+        .expect("honest affine 4-bit granite experts must load");
+    match ff {
+        FeedForward::Moe { moe, .. } => moe,
+        _ => panic!("num_local_experts 2 must build the MoE feed-forward"),
+    }
+}
+
+/// Before issue #2245 granite's `block_sparse_moe` never called
+/// `forward_fused_kernel`, so the fused single-token MoE kernel (#268) never
+/// ran for the family and every decode step paid the `gather_qmm` fallback
+/// (25.4% of decode GPU time on gfx1151, per the #2061 profile). The kernel
+/// must decline nothing on an honest affine 4-bit expert triple of dff 64 and
+/// compute the `forward` + `moe_weighted_sum` function within the tolerance
+/// that test pins. It launches the kernel, which on a backend without a port
+/// refuses across the bridge, so it runs only when
+/// `fused_moe_kernels_available()` says both ports exist.
+#[test]
+fn granite_hybrid_fused_kernel_matches_the_gather_path_at_one_token() {
+    if !mlxcel_core::fused_moe_kernels_available() {
+        eprintln!("skipping: this backend has no fused MoE kernel port (#1803)");
+        return;
+    }
+    let moe = fused_moe_block();
+    let x = mlxcel_core::astype(
+        &mlxcel_core::from_slice_f32(
+            &values(FUSED_HIDDEN as usize, -0.4, 0.0125),
+            &[1, FUSED_HIDDEN],
+        ),
+        mlxcel_core::dtype::BFLOAT16,
+    );
+    let indices = mlxcel_core::from_slice_i32(&[0, 1], &[1, FUSED_EXPERTS]);
+    let scores = mlxcel_core::from_slice_f32(&[0.25, 0.75], &[1, FUSED_EXPERTS]);
+
+    let fused = moe
+        .switch_mlp
+        .forward_fused_kernel(&x, &indices, &scores)
+        .expect("honest affine 4-bit experts of dff 64 must dispatch the fused kernel");
+    let fused = mlxcel_core::reshape(&fused, &[1, FUSED_HIDDEN]);
+    let gathered = moe_weighted_sum(
+        &moe.switch_mlp.forward(&x, &indices),
+        &scores,
+        mlxcel_core::dtype::BFLOAT16,
+    );
+    // A structural check, not a numeric parity bound (that lives in
+    // `fused_moe_parity_tests`): swapped planes or a dropped expert land near
+    // 1, the bf16 quantized jitter between the two paths two orders lower.
+    let fused = mlxcel_core::utils::array_to_vec_f32(&mlxcel_core::astype(
+        &fused,
+        mlxcel_core::dtype::FLOAT32,
+    ));
+    let gathered = mlxcel_core::utils::array_to_vec_f32(&mlxcel_core::astype(
+        &gathered,
+        mlxcel_core::dtype::FLOAT32,
+    ));
+    let diff_sq: f64 = fused
+        .iter()
+        .zip(&gathered)
+        .map(|(a, b)| f64::from(a - b).powi(2))
+        .sum();
+    let ref_sq: f64 = gathered.iter().map(|b| f64::from(*b).powi(2)).sum();
+    let nrms = (diff_sq / ref_sq.max(1e-20)).sqrt();
+    assert!(
+        fused.iter().all(|v| v.is_finite()) && nrms < 0.1,
+        "the fused kernel must compute the gather path's function: nrms {nrms:e}"
+    );
+}
+
+/// Reproduction for issue #2245: before it, `GraniteMoeHybridMoE::forward`
+/// always took the `SwitchGLU::forward` + `moe_weighted_sum` gather path, so
+/// the fused single-token decode kernel never ran for the family. The
+/// production `forward` must return, at one token, exactly the fused-kernel
+/// computation (bit for bit, since it IS that computation) and not the gather
+/// computation, which the two candidate values below pin. The candidates come
+/// from the block's own router outputs, so the comparison is against the exact
+/// indices and scores `forward` would hand the kernel. Runs only when both
+/// fused MoE ports exist and the `MLXCEL_FUSED_MOE` gate is on, since the
+/// gate is part of the routing under test.
+#[test]
+fn granite_hybrid_forward_routes_one_token_through_the_fused_kernel() {
+    use crate::models::switch_layers::fused_moe_enabled;
+
+    if !mlxcel_core::fused_moe_kernels_available() {
+        eprintln!("skipping: this backend has no fused MoE kernel port (#1803)");
+        return;
+    }
+    if !fused_moe_enabled() {
+        eprintln!("skipping: MLXCEL_FUSED_MOE disables the routing this test pins");
+        return;
+    }
+
+    let moe = fused_moe_block();
+    let x = mlxcel_core::astype(
+        &mlxcel_core::from_slice_f32(
+            &values(FUSED_HIDDEN as usize, -0.4, 0.0125),
+            &[1, FUSED_HIDDEN],
+        ),
+        mlxcel_core::dtype::BFLOAT16,
+    );
+
+    // The router outputs `forward` computes, replicated with the block's own
+    // router so the candidates use the exact indices and scores the
+    // production step would.
+    let logits = moe.router.forward(&x);
+    let kth = 0; // two experts, top_k 2: every expert is routed
+    let indices = mlxcel_core::argpartition(&logits, kth, -1);
+    let indices_shape = mlxcel_core::array_shape(&indices);
+    let topk_indices =
+        mlxcel_core::slice(&indices, &[0, kth], &[indices_shape[0], indices_shape[1]]);
+    let topk_logits = mlxcel_core::take_along_axis(&logits, &topk_indices, -1);
+    let topk_logits = mlxcel_core::astype(&topk_logits, mlxcel_core::dtype::FLOAT32);
+    let scores = mlxcel_core::softmax(&topk_logits, -1);
+
+    let fused = moe
+        .switch_mlp
+        .forward_fused_kernel(&x, &topk_indices, &scores)
+        .expect("honest affine 4-bit experts of dff 64 must dispatch the fused kernel");
+    let fused_candidate = mlxcel_core::reshape(&fused, &[1, FUSED_HIDDEN]);
+    let gather_candidate = moe_weighted_sum(
+        &moe.switch_mlp.forward(&x, &topk_indices),
+        &scores,
+        mlxcel_core::dtype::BFLOAT16,
+    );
+
+    let routed = mlxcel_core::utils::array_to_vec_f32(&mlxcel_core::astype(
+        &moe.forward(&x),
+        mlxcel_core::dtype::FLOAT32,
+    ));
+    let fused = mlxcel_core::utils::array_to_vec_f32(&mlxcel_core::astype(
+        &fused_candidate,
+        mlxcel_core::dtype::FLOAT32,
+    ));
+    let gathered = mlxcel_core::utils::array_to_vec_f32(&mlxcel_core::astype(
+        &gather_candidate,
+        mlxcel_core::dtype::FLOAT32,
+    ));
+
+    assert!(
+        routed.iter().zip(&fused).all(|(a, b)| a == b),
+        "one-token forward must take the fused kernel path, not the gather path"
+    );
+    assert!(
+        routed.iter().zip(&gathered).any(|(a, b)| a != b),
+        "the gather candidate must differ from the fused path for this geometry, \
+         so the bit-exact match above is a routing check"
+    );
+}
+
 /// Regression for the `slice_update`-is-functional footgun: `inject_at_positions`
 /// MUST write the visual features into `h` at the image-token positions. A
 /// dropped `slice_update` return leaves the image slots zero and silently makes

src/models/switch_layers.rs

@@ -1215,6 +1215,10 @@ impl SwitchGLU {
     /// non-SwiGLU activation (the kernel fuses `silu(gate) * up` and nothing
     /// else, so Kimi K3's SiTU experts stay on the gather path). Gated by the
     /// caller (`MLXCEL_FUSED_MOE`).
+    ///
+    /// Used by: AfMoE, BailingMoe, Cohere2Moe, Dbrx, Dots1, GraniteMoeHybrid,
+    ///          Klear, Laguna, Lfm2, Mellum, MiniMax, Mixtral, OLMoE, PhiMoE,
+    ///          Qwen2Moe, Qwen3Moe, and Qwen3VLMoE
     pub fn forward_fused_kernel(
         &self,
         x: &MlxArray,