From 27208be66e6529d016358d15fe87e95810698227 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Tue, 4 Jun 2024 09:58:47 -0700 Subject: [PATCH] [Kernel] Add back batch size 1536 and 3072 to MoE tuning (#5242) --- benchmarks/kernels/benchmark_moe.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index d6fa39a4..2edc6314 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -254,7 +254,9 @@ def main(args: argparse.Namespace): use_fp8 = args.dtype == "fp8" if args.batch_size is None: - batch_sizes = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096] + batch_sizes = [ + 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 1536, 2048, 3072, 4096 + ] else: batch_sizes = [args.batch_size]