From 8a7af254a60b69f79fd3c864801852f46cc332aa Mon Sep 17 00:00:00 2001
From: Jeongkyu Shin <inureyes@gmail.com>
Date: Tue, 16 Jun 2026 17:46:53 +0900
Subject: [PATCH] perf(moe): wire qwen3_vl_moe to the fused decode-MoE kernel

SparseMoeBlock::forward in qwen3_vl_moe now gates on single-token
decode (n_tokens == 1) and fused_moe_enabled(), then delegates to
SwitchGLU::forward_fused_kernel (imported from qwen3_moe). The
fallback path (experts.forward + moe_weighted_sum) is unchanged for
prefill and any unsupported bit-width config. Activation is SwiGLU
(act=0); no GeGLU/relu2 template is needed.

Updates the 'Models covered' list in the fused-kernel design doc
to reflect seven wired model paths, moving qwen3_vl_moe out of the
unwired list.
---
 .../fused-moe-decode-kernel-design.md         | 14 ++++----
 src/models/qwen3_vl_moe.rs                    | 33 ++++++++++++++-----
 2 files changed, 32 insertions(+), 15 deletions(-)

diff --git a/docs/benchmark_results/fused-moe-decode-kernel-design.md b/docs/benchmark_results/fused-moe-decode-kernel-design.md
index a88993c3d..1b1a3236e 100644
--- a/docs/benchmark_results/fused-moe-decode-kernel-design.md
+++ b/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 six model paths: qwen3_moe
+The fused single-token decode dispatch is wired into seven model paths: qwen3_moe
 (Qwen3 MoE), qwen3_next (qwen3.5/3.6), dots.llm1 (mixed 4/6-bit), gemma4 (GeGLU),
 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), 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,
-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++
+checkpoint layout), lfm2 (LFM2-MoE; sigmoid-routed, optional expert_bias and
+norm_topk_prob), and qwen3_vl_moe (Qwen3-VL MoE; imports `SwitchGLU` from qwen3_moe,
+SwiGLU activation, text-only decode path). 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, and phimoe. nemotron-h's MoE runs through the separate C++
 `fused_moe_forward` and is wired behind `MLXCEL_FUSED_MOE_RELU2` only.
diff --git a/src/models/qwen3_vl_moe.rs b/src/models/qwen3_vl_moe.rs
index 6191f57e5..0958c53b4 100644
--- a/src/models/qwen3_vl_moe.rs
+++ b/src/models/qwen3_vl_moe.rs
@@ -601,14 +601,31 @@ impl SparseMoeBlock {
             scores = mlxcel_core::divide(&scores, &sum);
         }
 
-        // Apply experts - returns [n_tokens, k, hidden]
-        let expert_out = self.experts.forward(&x_flat, &topk_indices);
-
-        let result = crate::models::switch_layers::moe_weighted_sum(
-            &expert_out,
-            &scores,
-            mlxcel_core::array_dtype(&x_flat),
-        );
+        // Apply experts via fused single-token decode kernel when available,
+        // falling back to gather_qmm + weighted-sum for prefill and any
+        // unsupported bit-width config.
+        let result = {
+            let fused = if mlxcel_core::array_shape(&x_flat)[0] == 1
+                && crate::models::switch_layers::fused_moe_enabled()
+            {
+                self.experts
+                    .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.experts.forward(&x_flat, &topk_indices);
+                    crate::models::switch_layers::moe_weighted_sum(
+                        &expert_out,
+                        &scores,
+                        mlxcel_core::array_dtype(&x_flat),
+                    )
+                }
+            }
+        };
 
         // Reshape back to original shape
         if orig_shape.len() > 2 {