Route prefill MoE to batched tensor-core kernel in Qwen3.5 export

Add use_batched_moe flag on FusedMoEExperts, toggled by _set_batched_moe
in export.py before each method's torch.export call. Decode (T=1) uses
the vec-mat fused_moe kernel; prefill (T>=2) uses fused_moe_batched_gemm.

Author: digantdesai

examples/models/qwen3_5_moe/export.py

@@ -380,14 +380,20 @@ def _apply_turboquant(model, config):
 # ---------------------------------------------------------------------------
 
 
+def _set_batched_moe(model, enabled):
+    """Toggle batched tensor-core MoE kernel for all MoE layers."""
+    for layer in model.layers:
+        if hasattr(layer, "mlp") and hasattr(layer.mlp, "experts"):
+            layer.mlp.experts.use_batched_moe = enabled
+
+
 def export_and_lower(model, config, args):
     """Export model to .pte via torch.export + CUDA backend.
 
     Exports two methods:
-      - "decode": decode path (T=1), uses native PyTorch recurrent FLA
-        so AOTI can fuse with surrounding ops for maximum decode throughput.
-      - "prefill": prefill path (T>=2), uses chunked FLA triton_op with
-        dynamic sequence length.
+      - "decode": decode path (T=1), vec-mat MoE kernel via fused_moe.
+      - "prefill": prefill path (T>=2), batched tensor-core MoE kernel
+        via fused_moe_batched_gemm, with dynamic sequence length.
 
     Both methods share mutable state buffers (KV cache, conv_state,
     recurrent_state) via share_mutable_buffers=True. The model uses
@@ -412,7 +418,8 @@ def export_and_lower(model, config, args):
     # -O0 compiles ~8x faster than -O1 with no measurable runtime impact.
     inductor_config.aot_inductor.compile_wrapper_opt_level = "O0"
 
-    # --- Decode method (T=1, static shape) ---
+    # --- Decode method (T=1, static shape, vec-mat MoE kernel) ---
+    _set_batched_moe(model, False)
     print("Exporting decode method...")
     decode_tokens = torch.tensor([[0]], dtype=torch.long)
     decode_pos = torch.tensor([0], dtype=torch.long)
@@ -424,10 +431,14 @@ def export_and_lower(model, config, args):
         )
     print("Decode export successful!")
 
-    # --- Prefill method (T>=2, dynamic shape) ---
+    # --- Prefill method (T>=2, dynamic shape, batched tensor-core MoE kernel) ---
+    # Example T must equal max_seq_len-1 so AOTI compiles kernels for the
+    # full range of sequence lengths.
+    _set_batched_moe(model, True)
     print("Exporting prefill method...")
-    prefill_tokens = torch.tensor([[0, 1]], dtype=torch.long)
-    prefill_pos = torch.tensor([0, 1], dtype=torch.long)
+    example_seq_len = config.max_seq_len - 1
+    prefill_tokens = torch.zeros((1, example_seq_len), dtype=torch.long)
+    prefill_pos = torch.arange(example_seq_len, dtype=torch.long)
     seq_dim = Dim("seq_len", min=2, max=config.max_seq_len - 1)
     prefill_dynamic_shapes = (
         {1: seq_dim},  # tokens

examples/models/qwen3_5_moe/model.py

@@ -479,6 +479,7 @@ def __init__(self, config):
         self.intermediate_size = config.moe_intermediate_size
         self.hidden_size = config.hidden_size
         self.group_size = 32
+        self.use_batched_moe = False
 
         self.w1_weight = nn.Parameter(
             torch.empty(
@@ -496,6 +497,19 @@ def __init__(self, config):
         )
 
     def forward(self, x, expert_weights, expert_indices, top_k):
+        if self.use_batched_moe:
+            return torch.ops.triton.fused_moe_batched_gemm(
+                x,
+                self.w1,
+                self.w1_scale,
+                self.w2,
+                self.w2_scale,
+                expert_weights,
+                expert_indices,
+                top_k,
+                self.num_experts,
+                self.group_size,
+            )
         return torch.ops.triton.fused_moe(
             x,
             self.w1,