From 18cea7441b11810beca38e9a6216653846fc837e Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Wed, 5 Mar 2025 12:52:36 +0000 Subject: [PATCH] Reset Signed-off-by: Jee Jee Li --- benchmarks/kernels/benchmark_moe.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index bb28c32798e2..dce0bef42039 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -509,6 +509,11 @@ def main(args: argparse.Namespace): intermediate_size = config.moe_intermediate_size shard_intermediate_size = 2 * intermediate_size // args.tp_size block_quant_shape = config.quantization_config['weight_block_size'] + elif config.architectures[0] == "Qwen2MoeForCausalLM": + E = config.num_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // args.tp_size else: # Default: Mixtral. E = config.num_local_experts