A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Misc] benchmark_moe supports expert parallel (#22251)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Aug 11, 2025 at 15:13 UTC
384a052971607f1561e734c87c9216f77f47e0fb
1 file changed
+11
-9
benchmarks/kernels/benchmark_moe.py
+11
-9
@@ -22,10 +22,10 @@ from vllm.utils import FlexibleArgumentParser
22
FP8_DTYPE = current_platform.fp8_dtype()
23
24
25
-def ensure_divisibility(numerator, denominator):
25
+def ensure_divisibility(numerator, denominator, text):
26
"""Ensure that numerator is divisible by the denominator."""
27
- assert numerator % denominator == 0, (
28
- "intermediate_size {} is not divisible by tp {}.".format(numerator, denominator)
27
+ assert numerator % denominator == 0, "{} {} is not divisible by tp {}.".format(
28
+ text, numerator, denominator
29
)
30
31
@@ -577,12 +577,10 @@ def main(args: argparse.Namespace):
577
E = config.ffn_config.moe_num_experts
578
topk = config.ffn_config.moe_top_k
579
intermediate_size = config.ffn_config.ffn_hidden_size
580
- shard_intermediate_size = 2 * intermediate_size // args.tp_size
580
elif config.architectures[0] == "JambaForCausalLM":
581
E = config.num_experts
582
topk = config.num_experts_per_tok
583
intermediate_size = config.intermediate_size
585
- shard_intermediate_size = 2 * intermediate_size // args.tp_size
584
elif config.architectures[0] in (
585
"DeepseekV3ForCausalLM",
586
"DeepseekV2ForCausalLM",
@@ -591,17 +589,14 @@ def main(args: argparse.Namespace):
589
E = config.n_routed_experts
590
topk = config.num_experts_per_tok
591
intermediate_size = config.moe_intermediate_size
594
- shard_intermediate_size = 2 * intermediate_size // args.tp_size
592
elif config.architectures[0] in ("Qwen2MoeForCausalLM", "Qwen3MoeForCausalLM"):
593
E = config.num_experts
594
topk = config.num_experts_per_tok
595
intermediate_size = config.moe_intermediate_size
599
- shard_intermediate_size = 2 * intermediate_size // args.tp_size
596
elif config.architectures[0] in ("HunYuanMoEV1ForCausalLM"):
597
E = config.num_experts
598
topk = config.moe_topk[0]
599
intermediate_size = config.moe_intermediate_size[0]
604
- shard_intermediate_size = 2 * intermediate_size // args.tp_size
600
else:
601
# Support for llama4
602
config = config.get_text_config()
@@ -609,8 +604,14 @@ def main(args: argparse.Namespace):
604
E = config.num_local_experts
605
topk = config.num_experts_per_tok
606
intermediate_size = config.intermediate_size
607
+ enable_ep = bool(args.enable_expert_parallel)
608
+ if enable_ep:
609
+ ensure_divisibility(E, args.tp_size, "Number of experts")
610
+ E = E // args.tp_size
611
+ shard_intermediate_size = 2 * intermediate_size
612
+ else:
613
+ ensure_divisibility(intermediate_size, args.tp_size, "intermediate_size")
614
shard_intermediate_size = 2 * intermediate_size // args.tp_size
613
- ensure_divisibility(intermediate_size, args.tp_size)
615
hidden_size = config.hidden_size
616
dtype = torch.float16 if current_platform.is_rocm() else config.torch_dtype
617
use_fp8_w8a8 = args.dtype == "fp8_w8a8"
@@ -742,6 +743,7 @@ if __name__ == "__main__":
743
parser.add_argument(
744
"--tp-size", "-tp", "--tensor-parallel-size", type=int, default=2
745
)
746
+ parser.add_argument("--enable-expert-parallel", "-enable-ep", action="store_true")
747
parser.add_argument(
748
"--dtype", type=str, choices=["auto", "fp8_w8a8", "int8_w8a16"], default="auto"
749
)