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] Add --save-dir option to benchmark_moe (#23020)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Aug 16, 2025 at 15:26 UTC
6d3da472bc8f202229a8e178671f4fe72037cfb1
1 file changed
+8
-1
benchmarks/kernels/benchmark_moe.py
+8
-1
@@ -3,6 +3,7 @@
3
4
import argparse
5
import json
6
+import os
7
import time
8
from contextlib import nullcontext
9
from datetime import datetime
@@ -542,6 +543,7 @@ def save_configs(
543
use_fp8_w8a8: bool,
544
use_int8_w8a16: bool,
545
block_quant_shape: list[int],
546
+ save_dir: str,
547
) -> None:
548
dtype_str = get_config_dtype_str(
549
dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8
@@ -552,7 +554,8 @@ def save_configs(
554
filename = get_config_file_name(
555
num_experts, shard_intermediate_size // 2, dtype_str, block_quant_shape
556
)
555
-
557
+ os.makedirs(save_dir, exist_ok=True)
558
+ filename = os.path.join(save_dir, filename)
559
print(f"Writing best config to {filename}...")
560
with open(filename, "w") as f:
561
json.dump(configs, f, indent=4)
@@ -707,6 +710,7 @@ def main(args: argparse.Namespace):
710
use_fp8_w8a8,
711
use_int8_w8a16,
712
block_quant_shape,
713
+ args.save_dir,
714
)
715
end = time.time()
716
print(f"Tuning took {end - start:.2f} seconds")
@@ -748,6 +752,9 @@ if __name__ == "__main__":
752
"--dtype", type=str, choices=["auto", "fp8_w8a8", "int8_w8a16"], default="auto"
753
)
754
parser.add_argument("--use-deep-gemm", action="store_true")
755
+ parser.add_argument(
756
+ "--save-dir", type=str, default="./", help="Directory to save tuned results"
757
+ )
758
parser.add_argument("--seed", type=int, default=0)
759
parser.add_argument("--batch-size", type=int, nargs="+", required=False)
760
parser.add_argument("--tune", action="store_true")