Enable sleeping between trials (#363) (#363)

Summary:
Pull Request resolved: #363

Adds command line support for a `--sleep` argument that can be used to inject a sleep in between trials. This is potentially useful on Blackwell where due to the floating clock measurements sometimes favor a earlier running kernel.

While most of this should be covered by just specifying running for at least 2 seconds, this can help ensure a more consistent run.

Reviewed By: xuzhao9, FindHao

Differential Revision: D81152631

Author: njriasan

run.py

@@ -44,7 +44,7 @@ def _run(args: argparse.Namespace, extra_args: List[str]) -> BenchmarkOperatorRe
         extra_args=extra_args,
     )
     try:
-        opbench.run(args.warmup, args.iter)
+        opbench.run(args.warmup, args.rep, sleep=args.sleep)
     finally:
         metrics = opbench.output
         if not args.skip_print:

tritonbench/utils/parser.py

@@ -60,6 +60,12 @@ def get_parser(args=None):
         default=DEFAULT_REP,
         help="The rep time for each benchmark run.",
     )
+    parser.add_argument(
+        "--sleep",
+        type=float,
+        default=0.0,
+        help="The amount of time (in seconds) to sleep between benchmark runs.",
+    )
     parser.add_argument(
         "--csv",
         action="store_true",

tritonbench/utils/triton_op.py

@@ -81,6 +81,7 @@ class BenchmarkOperatorBackend:
 DEFAULT_WARMUP = 25
 DEFAULT_REP = 100
 DEFAULT_QUANTILES = [0.5, 0.1, 0.9]
+DEFAULT_SLEEP = 0.0
 REGISTERED_BENCHMARKS: Dict[str, OrderedDict[str, BenchmarkOperatorBackend]] = {}
 REGISTERED_METRICS: defaultdict[str, List[str]] = defaultdict(list)
 OVERRIDDEN_METRICS: defaultdict[str, List[str]] = defaultdict(list)
@@ -808,7 +809,11 @@ def fwd_no_grad_fn():
             return fwd_no_grad_fn
 
     def run(
-        self, warmup=DEFAULT_WARMUP, rep=DEFAULT_REP, quantiles=DEFAULT_QUANTILES
+        self,
+        warmup=DEFAULT_WARMUP,
+        rep=DEFAULT_REP,
+        quantiles=DEFAULT_QUANTILES,
+        sleep=DEFAULT_SLEEP,
     ) -> None:
         """Benchmarking the operator and returning its metrics."""
         metrics = []
@@ -910,6 +915,9 @@ def _reduce_benchmarks(acc, bm_name: str):
                     )
                     if baseline:
                         self.baseline_metrics = acc[bm_name]
+                    if sleep:
+                        logging.debug(f"Sleeping for {sleep} seconds before next run")
+                        time.sleep(sleep)
                     return acc
 
                 y_vals: Dict[str, BenchmarkOperatorMetrics] = functools.reduce(