[Refactor] tuning_fused_moe for MLLM and small refactor (#11224)

Co-authored-by: Cursor Agent <cursoragent@cursor.com>

benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py

@@ -419,97 +419,25 @@ def get_filename(
 def main(args: argparse.Namespace):
     print(args)
 
-    config = AutoConfig.from_pretrained(args.model, trust_remote_code=True)
-    if config.architectures[0] == "DbrxForCausalLM":
-        E = config.ffn_config.moe_num_experts // args.ep_size
-        topk = config.ffn_config.moe_top_k
-        intermediate_size = config.ffn_config.ffn_hidden_size
-        shard_intermediate_size = (
-            2 * intermediate_size // (args.tp_size // args.ep_size)
-        )
-    elif config.architectures[0] == "JambaForCausalLM":
-        E = config.num_experts // args.ep_size
-        topk = config.num_experts_per_tok
-        intermediate_size = config.intermediate_size
-        shard_intermediate_size = (
-            2 * intermediate_size // (args.tp_size // args.ep_size)
-        )
-    elif config.architectures[0] in [
-        "Qwen2MoeForCausalLM",
-        "Qwen3MoeForCausalLM",
-        "Qwen3NextForCausalLM",
-    ]:
-        E = config.num_experts // args.ep_size
-        topk = config.num_experts_per_tok
-        intermediate_size = config.moe_intermediate_size
-        shard_intermediate_size = (
-            2 * intermediate_size // (args.tp_size // args.ep_size)
-        )
-    elif config.architectures[0] in ["DeepseekV2ForCausalLM", "DeepseekV3ForCausalLM"]:
-        E = (config.n_routed_experts // args.ep_size) + (
-            0
-            if args.disable_shared_experts_fusion
-            or config.architectures[0] not in ["DeepseekV3ForCausalLM"]
-            else 1
-        )
-        topk = config.num_experts_per_tok
-        intermediate_size = config.moe_intermediate_size
-        shard_intermediate_size = (
-            2 * intermediate_size // (args.tp_size // args.ep_size)
-        )
-    elif config.architectures[0] == "Llama4ForConditionalGeneration":
-        E = config.text_config.num_local_experts // args.ep_size + (
-            0 if args.disable_shared_experts_fusion else 1
-        )
-        topk = config.text_config.num_experts_per_tok
-        intermediate_size = config.text_config.intermediate_size
-        shard_intermediate_size = (
-            2 * intermediate_size // (args.tp_size // args.ep_size)
-        )
-    elif config.architectures[0] in [
-        "Grok1ForCausalLM",
-        "Grok1ImgGen",
-        "Grok1AForCausalLM",
-    ]:
-        E = config.num_local_experts // args.ep_size
-        topk = config.num_experts_per_tok
-        intermediate_size = config.moe_intermediate_size
-        shard_intermediate_size = (
-            2 * intermediate_size // (args.tp_size // args.ep_size)
-        )
-    elif config.architectures[0] in [
-        "BailingMoEForCausalLM",
-        "BailingMoeForCausalLM",
-        "BailingMoeV2ForCausalLM",
-    ]:
-        E = config.num_experts // args.ep_size
-        topk = config.num_experts_per_tok
-        intermediate_size = config.moe_intermediate_size
-        shard_intermediate_size = (
-            2 * intermediate_size // (args.tp_size // args.ep_size)
-        )
-    elif config.architectures[0] in ["Glm4MoeForCausalLM"]:
-        E = config.n_routed_experts // args.ep_size
-        topk = config.num_experts_per_tok
-        intermediate_size = config.moe_intermediate_size
-        shard_intermediate_size = (
-            2 * intermediate_size // (args.tp_size // args.ep_size)
-        )
-    else:
-        # Default: Mixtral
-        E = config.num_local_experts // args.ep_size
-        topk = config.num_experts_per_tok
-        intermediate_size = config.intermediate_size
-        shard_intermediate_size = (
-            2 * intermediate_size // (args.tp_size // args.ep_size)
+    def _calculate_shard_intermediate_size(intermediate_size: int) -> int:
+        # In EP mode, use original intermediate_size; otherwise apply TP sharding
+        return (
+            intermediate_size
+            if args.ep_size > 1
+            else 2 * intermediate_size // args.tp_size
         )
 
-    hidden_size = getattr(config, "hidden_size", None) or config.text_config.hidden_size
-    dtype = config.torch_dtype
-    use_fp8_w8a8 = args.dtype == "fp8_w8a8"
-    use_int8_w8a8 = args.dtype == "int8_w8a8"
-    use_int8_w8a16 = args.dtype == "int8_w8a16"
-    per_channel_quant = args.per_channel_quant
+    # Check EP mode constraint: tp_size must be 1 when ep_size > 1
+    if args.ep_size > 1 and args.tp_size != 1:
+        raise ValueError(
+            f"When using Expert Parallelism (ep_size={args.ep_size}), "
+            f"tp_size must be set to 1, but got tp_size={args.tp_size}. "
+            f"Please set --tp-size 1 when using --ep-size > 1."
+        )
+
+    config = AutoConfig.from_pretrained(args.model, trust_remote_code=True)
+
+    # Determine block shape for quantization
     block_shape = None
     if (
         hasattr(config, "quantization_config")
@@ -518,6 +446,82 @@ def main(args: argparse.Namespace):
         block_shape = config.quantization_config["weight_block_size"]
         assert len(block_shape) == 2
 
+    architecture = config.architectures[0]
+    # replace config with text_config for encoder-decoder models after getting block_shape and architecture
+    if hasattr(config, "text_config"):
+        config = config.get_text_config()
+
+    if architecture == "DbrxForCausalLM":
+        E = config.ffn_config.moe_num_experts
+        topk = config.ffn_config.moe_top_k
+        intermediate_size = config.ffn_config.ffn_hidden_size
+        shard_intermediate_size = _calculate_shard_intermediate_size(intermediate_size)
+    elif architecture == "JambaForCausalLM":
+        E = config.num_experts
+        topk = config.num_experts_per_tok
+        intermediate_size = config.intermediate_size
+        shard_intermediate_size = _calculate_shard_intermediate_size(intermediate_size)
+    elif architecture in [
+        "Qwen2MoeForCausalLM",
+        "Qwen3MoeForCausalLM",
+        "Qwen3NextForCausalLM",
+        "Qwen3VLMoeForConditionalGeneration",
+    ]:
+        E = config.num_experts // args.ep_size
+        topk = config.num_experts_per_tok
+        intermediate_size = config.moe_intermediate_size
+        shard_intermediate_size = _calculate_shard_intermediate_size(intermediate_size)
+    elif architecture in ["DeepseekV2ForCausalLM", "DeepseekV3ForCausalLM"]:
+        E = (
+            config.n_routed_experts + (0 if args.disable_shared_experts_fusion else 1)
+            if architecture == "DeepseekV3ForCausalLM"
+            else config.n_routed_experts
+        )
+        topk = config.num_experts_per_tok
+        intermediate_size = config.moe_intermediate_size
+        shard_intermediate_size = _calculate_shard_intermediate_size(intermediate_size)
+    elif architecture == "Llama4ForConditionalGeneration":
+        E = config.num_local_experts + (0 if args.disable_shared_experts_fusion else 1)
+        topk = config.num_experts_per_tok
+        intermediate_size = config.intermediate_size
+        shard_intermediate_size = _calculate_shard_intermediate_size(intermediate_size)
+    elif architecture in [
+        "Grok1ForCausalLM",
+        "Grok1ImgGen",
+        "Grok1AForCausalLM",
+    ]:
+        E = config.num_local_experts // args.ep_size
+        topk = config.num_experts_per_tok
+        intermediate_size = config.moe_intermediate_size
+        shard_intermediate_size = _calculate_shard_intermediate_size(intermediate_size)
+    elif architecture in [
+        "BailingMoEForCausalLM",
+        "BailingMoeForCausalLM",
+        "BailingMoeV2ForCausalLM",
+    ]:
+        E = config.num_experts // args.ep_size
+        topk = config.num_experts_per_tok
+        intermediate_size = config.moe_intermediate_size
+        shard_intermediate_size = _calculate_shard_intermediate_size(intermediate_size)
+    elif architecture in ["Glm4MoeForCausalLM"]:
+        E = config.n_routed_experts
+        topk = config.num_experts_per_tok
+        intermediate_size = config.moe_intermediate_size
+        shard_intermediate_size = _calculate_shard_intermediate_size(intermediate_size)
+    else:
+        # Default: Mixtral
+        E = config.num_local_experts // args.ep_size
+        topk = config.num_experts_per_tok
+        intermediate_size = config.intermediate_size
+        shard_intermediate_size = _calculate_shard_intermediate_size(intermediate_size)
+
+    hidden_size = config.hidden_size
+    dtype = config.torch_dtype
+    use_fp8_w8a8 = args.dtype == "fp8_w8a8"
+    use_int8_w8a8 = args.dtype == "int8_w8a8"
+    use_int8_w8a16 = args.dtype == "int8_w8a16"
+    per_channel_quant = args.per_channel_quant
+
     if args.batch_size is None:
         batch_sizes = [
             1,