fix benchmark hot path

Author: yf225

benchmarks/run.py

@@ -24,6 +24,7 @@
 import sys
 from typing import Any
 from typing import Callable
+import time
 
 # Maps tritonbench op names to Helion kernel examples
 # Can map to a single kernel or a list of kernel variants
@@ -349,21 +350,17 @@ def helion_method(
                 # so that each input size can go through its own autotuning.
                 from helion.runtime.kernel import Kernel
 
-                for attr_name in dir(mod):
-                    attr = getattr(mod, attr_name)
-                    if isinstance(attr, Kernel):
-                        attr.reset()
+                # Force autotuning unless HELION_USE_DEFAULT_CONFIG=1 is set
+                # This ensures we run autotuning even if the kernel has pre-specified configs
+                if os.environ.get("HELION_USE_DEFAULT_CONFIG", "0") != "1":
+                    # Find all Kernel objects in the module and force autotuning
+                    for attr_name in dir(mod):
+                        attr = getattr(mod, attr_name)
+                        if isinstance(attr, Kernel):
+                            attr.reset()
+                            attr.settings.force_autotune = True
 
                 def _inner() -> Callable[..., Any] | object:
-                    # Force autotuning unless HELION_USE_DEFAULT_CONFIG=1 is set
-                    # This ensures we run autotuning even if the kernel has pre-specified configs
-                    if os.environ.get("HELION_USE_DEFAULT_CONFIG", "0") != "1":
-                        # Find all Kernel objects in the module and force autotuning
-                        for attr_name in dir(mod):
-                            attr = getattr(mod, attr_name)
-                            if isinstance(attr, Kernel):
-                                attr.settings.force_autotune = True
-
                     result = kfunc(*args)
                     if callable(result):
                         return result()

benchmarks/run_input_shard.sh

@@ -11,9 +11,9 @@ attempt=0
 while true; do
     attempt=$((attempt + 1))
     echo "Attempt $attempt: Running benchmark for shard $((SHARD+1))/${WORLD_SIZE}..."
-    
-    HELION_FORCE_DISK_CACHE=1 CUDA_VISIBLE_DEVICES=$((RANK_OFFSET+SHARD)) python benchmarks/run.py --input-shard $((SHARD+1))/${WORLD_SIZE} --metrics tflops,gbps,speedup >"$OUTPUT_FILE" 2>&1
-    
+
+    CUDA_VISIBLE_DEVICES=$((RANK_OFFSET+SHARD)) python benchmarks/run.py --input-shard $((SHARD+1))/${WORLD_SIZE} --metrics tflops,gbps,speedup >"$OUTPUT_FILE" 2>&1
+
     exit_code=$?
     if [ $exit_code -eq 0 ]; then
         echo "Success! Benchmark completed for shard $((SHARD+1))/${WORLD_SIZE}"