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.

This PR was authored with the assistance of Claude.

Author: digantdesai

examples/models/qwen3_5_moe/export.py

@@ -503,6 +503,13 @@ 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 + backend-specific lowering."""
     backend = getattr(args, "backend", "cuda")
@@ -597,10 +604,9 @@ def _export_cuda(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
@@ -625,7 +631,8 @@ def _export_cuda(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)
@@ -637,11 +644,12 @@ def _export_cuda(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 (especially
     # chunk_gated_delta_rule with CHUNK_SIZE=64) for the full range of sequence
     # lengths. Smaller examples cause AOTI to bake in intermediate buffer sizes
     # that reject longer prompts at runtime.
+    _set_batched_moe(model, True)
     print("Exporting prefill method...")
     example_prefill_len = config.max_seq_len - 1
     prefill_tokens = torch.zeros((1, example_prefill_len), dtype=torch.long)

examples/models/qwen3_5_moe/model.py

@@ -492,6 +492,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(
@@ -509,6 +510,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,