fix: main branch mla data read

Signed-off-by: Kimi Zhao <[email protected]>

Author: YijiaZhao

collector/trtllm/collect_moe.py

@@ -20,8 +20,10 @@
 
 aic_debug = int(os.getenv("aic_moe_debug", "0"))  # noqa: SIM112
 
+moe_tune_path = os.path.dirname(os.path.abspath(__file__))
 
-def cleanup_empty_json_files(directory="/tmp/moe_tune_path"):
+
+def cleanup_empty_json_files(directory):
     if not os.path.exists(directory):
         return
 
@@ -466,9 +468,9 @@ def fp32_to_mxfp4(tensor):
     torch.cuda.synchronize()
 
     if moe_type != "w4a16_mxfp4":
-        cleanup_empty_json_files("/tmp/moe_tune_path")
+        cleanup_empty_json_files(moe_tune_path)
         cache_path = (
-            f"/tmp/moe_tune_path/{moe_type}_{hidden_size}_{inter_size // moe_tp_size}_{num_experts // moe_ep_size}"
+            f"{moe_tune_path}/{moe_type}_{hidden_size}_{inter_size // moe_tp_size}_{num_experts // moe_ep_size}"
         )
         existing_files = glob.glob(f"{cache_path}*")
         cache_loaded = False

src/aiconfigurator/sdk/perf_database.py

@@ -1490,14 +1490,13 @@ def __init__(self, system: str, backend: str, version: str, systems_dir: str = "
                     )  # s
                     target_z_list = [1, 2, 4, 8, 16, 32, 64, 128, 256, 384, 512, 1024, 2048]  # b
 
-                self._extrapolate_data_grid(
-                    data_dict=data_dict,  # tpsize,sb
-                    target_x_list=target_x_list,
-                    target_y_list=target_y_list,
-                    target_z_list=target_z_list,
-                    sqrt_y_value=True,
-                )
-
+                    self._extrapolate_data_grid(
+                        data_dict=data_dict,  # tpsize,sb
+                        target_x_list=target_x_list,
+                        target_y_list=target_y_list,
+                        target_z_list=target_z_list,
+                        sqrt_y_value=True,
+                    )
         # wideep generation mla
         if getattr(self, "_wideep_generation_mla_data", None) is not None:
             for kernel_source in self._wideep_generation_mla_data:
@@ -1581,12 +1580,12 @@ def __init__(self, system: str, backend: str, version: str, systems_dir: str = "
                     2097152 * 8,
                 ]  # s
 
-            self._extrapolate_data_grid(
-                data_dict=data_dict,  # tpsize, bs
-                target_x_list=target_x_list,
-                target_y_list=target_y_list,
-                target_z_list=target_z_list,
-            )
+                self._extrapolate_data_grid(
+                    data_dict=data_dict,  # tpsize, bs
+                    target_x_list=target_x_list,
+                    target_y_list=target_y_list,
+                    target_z_list=target_z_list,
+                )
 
         # post-correction
         self._correct_data()