[benchmark] Add triton version in the moe tuned config (vllm-project#24769)

Signed-off-by: Jee Jee Li <[email protected]>

Author: jeejeelee

benchmarks/kernels/benchmark_moe.py

@@ -560,7 +560,7 @@ def save_configs(
     filename = os.path.join(save_dir, filename)
     print(f"Writing best config to {filename}...")
     with open(filename, "w") as f:
-        json.dump(configs, f, indent=4)
+        json.dump({"triton_version": triton.__version__, **configs}, f, indent=4)
         f.write("\n")
 
 

vllm/model_executor/layers/fused_moe/fused_moe.py

@@ -720,7 +720,10 @@ def get_moe_configs(
                 logger.info("Using configuration from %s for MoE layer.",
                             config_file_path)
                 # If a configuration has been found, return it
-                return {int(key): val for key, val in json.load(f).items()}
+                tuned_config = json.load(f)
+                # Delete triton_version from tuned_config
+                tuned_config.pop("triton_version", None)
+                return {int(key): val for key, val in tuned_config.items()}
 
     # If no optimized configuration is available, we will use the default
     # configuration

vllm/triton_utils/importing.py

@@ -68,7 +68,7 @@ class TritonPlaceholder(types.ModuleType):
 
     def __init__(self):
         super().__init__("triton")
-        self.__version__ = "3.3.0"
+        self.__version__ = "3.4.0"
         self.jit = self._dummy_decorator("jit")
         self.autotune = self._dummy_decorator("autotune")
         self.heuristics = self._dummy_decorator("heuristics")