From 1357397a34ab0845caf3f384578cb084d44d94e5 Mon Sep 17 00:00:00 2001
From: Liana Koleva <43767763+lianakoleva@users.noreply.github.com>
Date: Wed, 29 Oct 2025 01:12:25 -0700
Subject: [PATCH] feat: preview filename from tuning_fused_moe_triton.py
 (#12276)

Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
---
 .../tuning_fused_moe_triton.py                | 42 ++++++++++++-------
 1 file changed, 27 insertions(+), 15 deletions(-)

diff --git a/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py b/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py
index eecc3ca2b..f0c7eec05 100644
--- a/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py
+++ b/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py
@@ -376,6 +376,15 @@ def sort_config(config: BenchmarkConfig) -> BenchmarkConfig:
 
 def save_configs(
     configs: Dict[int, BenchmarkConfig],
+    filename: str,
+) -> None:
+    print(f"Writing best config to {filename}...")
+    with open(filename, "w") as f:
+        json.dump(configs, f, indent=4)
+        f.write("\n")
+
+
+def get_filename(
     num_experts: int,
     shard_intermediate_size: int,
     hidden_size: int,
@@ -404,10 +413,7 @@ def save_configs(
         per_channel_quant,
     )
 
-    print(f"Writing best config to {filename}...")
-    with open(filename, "w") as f:
-        json.dump(configs, f, indent=4)
-        f.write("\n")
+    return filename
 
 
 def main(args: argparse.Namespace):
@@ -541,7 +547,22 @@ def main(args: argparse.Namespace):
                 for config in search_space
                 if block_k % config["BLOCK_SIZE_K"] == 0
             ]
-        print(f"Start tuning over {len(search_space)} configurations...")
+
+        filename = get_filename(
+            E,
+            shard_intermediate_size,
+            hidden_size,
+            topk,
+            dtype,
+            use_fp8_w8a8,
+            use_int8_w8a8,
+            use_int8_w8a16,
+            per_channel_quant,
+            block_shape,
+        )
+        print(
+            f"Start tuning over {len(search_space)} configurations to create {filename}..."
+        )
 
         start = time.perf_counter()
         configs = _distribute(
@@ -569,16 +590,7 @@ def main(args: argparse.Namespace):
         }
         save_configs(
             best_configs,
-            E,
-            shard_intermediate_size,
-            hidden_size,
-            topk,
-            dtype,
-            use_fp8_w8a8,
-            use_int8_w8a8,
-            use_int8_w8a16,
-            per_channel_quant,
-            block_shape,
+            filename,
         )
         end = time.perf_counter()
         print(f"Tuning took {end - start:.2f} seconds")