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 15, 21:00 UTC
[Misc] Add Qwen2MoeForCausalLM moe tuning support (#14276)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Mar 5, 2025 at 23:11 UTC
7bab4bb048d9af7a465eab0c2a69a84c8b131ec2
1 file changed
+5
benchmarks/kernels/benchmark_moe.py
+5
@@ -509,6 +509,11 @@ def main(args: argparse.Namespace):
509
intermediate_size = config.moe_intermediate_size
510
shard_intermediate_size = 2 * intermediate_size // args.tp_size
511
block_quant_shape = config.quantization_config['weight_block_size']
512
+ elif config.architectures[0] == "Qwen2MoeForCausalLM":
513
+ E = config.num_experts
514
+ topk = config.num_experts_per_tok
515
+ intermediate_size = config.moe_intermediate_size
516
+ shard_intermediate_size = 2 * intermediate_size // args.tp_size
517
else:
518
# Default: Mixtral.
519
E = config.num_local_experts