[Kernel] Re-tune Mixtral MoE configurations for FP8 on H100 (#5238)

2024-06-05 10:59:14 -07:00
parent eb8fcd2666
commit 51a08e7d8f
3 changed files with 118 additions and 117 deletions
--- a/benchmarks/kernels/benchmark_moe.py
+++ b/benchmarks/kernels/benchmark_moe.py
@@ -255,7 +255,8 @@ def main(args: argparse.Namespace):

    if args.batch_size is None:
        batch_sizes = [
-            1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 1536, 2048, 3072, 4096
+            1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128, 256, 512, 1024, 1536,
+            2048, 3072, 4096
        ]
    else:
        batch_sizes = [args.batch_size]