Use cuda graphs for benchmarking

Summary:
Per https://fb.workplace.com/groups/420659799592399/posts/807860500872325/, it's a lot more accurate than using regular non-cudagraph benchmarking.

I had to change a bunch of use sites of `metrics.latency` because `do_bench_cudagraph` does not support returning quantiles. Could certainly fix it upstream, but that would take more time + it doesn't really seem like quantiles are that useful in TritonBench anyway.

Reviewed By: xuzhao9, sijiac

Differential Revision: D58502780

fbshipit-source-id: 8c97b95097f49ece47ce9b1660af60afae8c25e8

torchbenchmark/operators/addmm/operator.py

@@ -113,18 +113,17 @@ def gbps(
  + (torch.addmm(a, mat1, mat2).numel())
  )
  numel = numel * a.element_size() / 1e9
- gbps = list(map(lambda x: numel / x * 1e3, metrics.latency))
- return statistics.median(gbps)
+ return numel / metrics.latency * 1e3
 
  @register_metric()
  def tflops(
  self, fn_name: str, example_inputs: Any, metrics: BenchmarkOperatorMetrics
- ) -> List[float]:
+ ) -> float:
  _, mat1, mat2 = example_inputs
  m, k = mat1.size()
  k, n = mat2.size()
  flops = m * k * 2 * n
- return [flops / x / 1e12 * 1e3 for x in metrics.latency]
+ return flops / metrics.latency / 1e12 * 1e3
 
  @register_x_val(label="(M, N, K)")
  def get_x_val(self, example_inputs) -> Tuple[int, int, int]:

torchbenchmark/operators/flash_attention/operator.py

@@ -278,7 +278,7 @@ def sdpa_flash_attention(q, k, v):
  @register_metric()
  def tflops(
  self, fn_name: str, example_inputs: Any, metrics: BenchmarkOperatorMetrics
- ) -> List[float]:
+ ) -> float:
  flops_per_matmul = (
  2.0 * self.BATCH * self.H * self.N_CTX * self.N_CTX * self.D_HEAD
  )
@@ -289,7 +289,7 @@ def tflops(
  tflops *= 2.5 # 2.0(bwd) + 0.5(recompute)
  elif self.mode == BenchmarkMode.FWD_BWD:
  tflops *= 3.5 # 1.0(fwd) + 2.0(bwd) + 0.5(recompute)
- return list(map(lambda x: tflops / x * 1e-9, metrics.latency))
+ return tflops / metrics.latency * 1e-9
 
  def get_bwd_fn(self, fwd_fn: Callable) -> Callable:
  o = fwd_fn()

torchbenchmark/operators/fp8_gemm/fp8_gemm.py

@@ -92,7 +92,7 @@ def nbytes(t):
  m, k = a.shape
  _, n = b.shape
  gb = (nbytes(a) + nbytes(b) + nbytes(c)) / 1e9
- return list(map(lambda x: gb / x * 1e3, metrics.latency))
+ return gb / metrics.latency * 1e3
 
  @register_metric()
  def tflops(
@@ -102,7 +102,7 @@ def tflops(
  m, k = a.size()
  _, n = b.size()
  flops = 2 * m * n * k
- return [flops / x / 1e12 * 1e3 for x in metrics.latency]
+ return flops / metrics.latency / 1e12 * 1e3
 
  def plot(self):
  @triton.testing.perf_report(

torchbenchmark/operators/gather_gemv/operator.py

@@ -24,19 +24,19 @@
 from .triton_gather_gemv import triton_gemv_0 as triton_test_0
 from torch._dynamo.testing import rand_strided
 
+
 class Operator(BenchmarkOperator):
 
  @register_metric()
  def gbps(self, fn_name, example_inputs, metrics: BenchmarkOperatorMetrics):
  arg0_1, arg1_1, arg2_1 = example_inputs
- gbps = (
- lambda ms: 2
+ return (
+ 2
  * arg2_1.size(0) * arg2_1.size(0)
  * arg0_1.element_size()
- / ms
+ / metrics.latency
  * 1e-6
  )
- return list(map(gbps, metrics.latency))
 
  def __init__(self, mode: str, device: str, extra_args: List[str] = []):
  super().__init__(mode=mode, device=device, extra_args=extra_args)

torchbenchmark/operators/gemm/operator.py

@@ -177,8 +177,7 @@ def gbps(
  a, w, bias = example_inputs
  numel = a.numel() + w.numel() + (torch.mm(a, w).numel())
  numel = numel * a.element_size() / 1e9
- gbps = list(map(lambda x: numel / x * 1e3, metrics.latency))
- return statistics.median(gbps)
+ return numel / metrics.latency * 1e3
 
  @register_metric(skip_baseline=True)
  def best_config(
@@ -205,7 +204,7 @@ def tflops(
  flops = m * k * 2 * n + 2 * m * n
  else:
  flops = m * k * 2 * n
- return [flops / x / 1e12 * 1e3 for x in metrics.latency]
+ return flops / metrics.latency / 1e12 * 1e3
 
  @staticmethod
  def _scaled_randn(*args, scale: float, **kwargs) -> torch.Tensor:

torchbenchmark/operators/int4_gemm/int4_gemm.py

@@ -95,7 +95,7 @@ def nbytes(t):
  c = fn()
 
  gb = (sum(nbytes(t) for t in (x, scale_and_zero, c)) + nbytes(w) // 8) / 1e9
- return list(map(lambda ms: gb / ms * 1e3, metrics.latency))
+ return gb / metrics.latency * 1e3
 
  @register_metric()
  def tflops(
@@ -106,7 +106,7 @@ def tflops(
  m = B * m
  _, n = b.size()
  flops = 2 * m * n * k
- return [flops / x / 1e12 * 1e3 for x in metrics.latency]
+ return flops / metrics.latency / 1e12 * 1e3
 
  def plot(self):
  @triton.testing.perf_report(

torchbenchmark/operators/layer_norm/operator.py

@@ -67,16 +67,12 @@ def get_x_val(self, args):
  @register_metric()
  def gbps(self, fn_name, args, metrics: BenchmarkOperatorMetrics) -> float:
  x = args[0]
-
- def gbps(ms):
- base = x.numel() * x.element_size() / ms * 1e-6
- return {
- Mode.FWD: 2 * base,
- Mode.BWD: 3 * base,
- Mode.FWD_BWD: 5 * base,
- }[self.mode]
-
- return list(map(gbps, metrics.latency))
+ base = x.numel() * x.element_size() / metrics.latency * 1e-6
+ return {
+ Mode.FWD: 2 * base,
+ Mode.BWD: 3 * base,
+ Mode.FWD_BWD: 5 * base,
+ }[self.mode]
 
  def plot(self):
  @triton.testing.perf_report(

torchbenchmark/operators/low_mem_dropout/operator.py

@@ -13,21 +13,25 @@
 
 from .kernels import _triton_dropout, _seeded_triton_dropout
 
+
 class Operator(BenchmarkOperator):
  @register_metric()
  def gbps(self, fn_name, example_inputs, metrics: BenchmarkOperatorMetrics):
- gbps = (
- lambda ms: 3 * example_inputs[1].element_size() * example_inputs[1].numel() / ms * 1e-6
+ return (
+ 3
+ * example_inputs[1].element_size()
+ * example_inputs[1].numel()
+ / metrics.latency
+ * 1e-6
  )
- return list(map(gbps, metrics.latency))
 
  @register_metric()
  def tflops(
  self, fn_name: str, example_inputs, metrics: BenchmarkOperatorMetrics
  ):
  p, a = example_inputs
  flops = 2 * len(a)
- return [flops / x / 1e12 * 1e3 for x in metrics.latency]
+ return flops / metrics.latency
 
  @register_benchmark()
  def triton_dropout(self, p, x):

torchbenchmark/operators/softmax/operator.py

@@ -113,15 +113,14 @@ def get_x_val(self, example_inputs) -> int:
  @register_metric()
  def gbps(
  self, fn_name, example_inputs, metrics: BenchmarkOperatorMetrics
- ) -> List[float]:
- gbps = (
- lambda ms: 2
+ ) -> float:
+ return (
+ 2
  * example_inputs[0].nelement()
  * example_inputs[0].element_size()
  * 1e-9
- / (ms * 1e-3)
+ / (metrics.latency * 1e-3)
  )
- return list(map(gbps, metrics.latency))
 
  def plot(self):
  @triton.testing.perf_report(

torchbenchmark/operators/sum/operator.py

@@ -187,13 +187,12 @@ def input_dims(
 
  @register_metric()
  def gbps(self, fn_name, example_inputs, metrics: BenchmarkOperatorMetrics):
- gbps = (
- lambda ms: example_inputs[0].element_size()
+ return (
+ example_inputs[0].element_size()
  * example_inputs[0].numel()
- / ms
+ / metrics.latency
  * 1e-6
  )
- return list(map(gbps, metrics.latency if metrics.latency else [0]))
 
  @register_metric(skip_baseline=True)
  def best_config(

torchbenchmark/operators/vector_add/operator.py

@@ -17,14 +17,13 @@ class Operator(BenchmarkOperator):
 
  @register_metric()
  def gbps(self, fn_name, example_inputs, metrics: BenchmarkOperatorMetrics):
- gbps = (
- lambda ms: 3
+ return (
+ 3
  * example_inputs[0].element_size()
  * example_inputs[0].numel()
- / ms
+ / metrics.latency
  * 1e-6
  )
- return list(map(gbps, metrics.latency))
 
  @register_benchmark()
  def triton_add(self, x: torch.Tensor, y: torch.Tensor):

torchbenchmark/util/triton_op.py

@@ -136,9 +136,9 @@ def dump_autotuner_best_config(kernel: triton.runtime.Autotuner) -> str:
 @dataclass
 class BenchmarkOperatorMetrics:
  # latency in ms
- latency: Optional[List[float]]
+ latency: Optional[float]
  # tflops
- tflops: Optional[List[float]]
+ tflops: Optional[float]
  # speedup over baseline
  speedup: Optional[float]
  # accuracy over baseline
@@ -735,13 +735,13 @@ def _init_extra_metrics() -> Dict[str, Any]:
  if set(["latency", "tflops", "speedup", "compile_time"]) & set(
  self.required_metrics
  ):
- metrics.latency = triton.testing.do_bench(
- fn,
- warmup=warmup,
- rep=rep,
- quantiles=quantiles,
- grad_to_none=self.get_grad_to_none(self.example_inputs),
- )
+ with torch.cuda.stream(torch.cuda.Stream()):
+ metrics.latency = triton.testing.do_bench_cudagraph(
+  fn,
+  rep=rep,
+  return_mode="median",
+  grad_to_none=self.get_grad_to_none(self.example_inputs),
+  )
  if "walltime" in self.required_metrics:
  metrics.walltime = do_bench_walltime(
  fn,
@@ -750,7 +750,7 @@ def _init_extra_metrics() -> Dict[str, Any]:
  )
  if "speedup" in self.required_metrics:
  metrics.speedup = (
- numpy.median(self.baseline_metrics.latency) / numpy.median(metrics.latency)
+ self.baseline_metrics.latency / metrics.latency
  if self.baseline_metrics and self.baseline_metrics.latency
  else None
  )
@@ -950,7 +950,7 @@ def compile_time(
  op_task.run()
  latency_with_compile = op_task.get_attribute("_latency_with_compile_in_task")
  del op_task
- latency_without_compile = numpy.median(metrics.latency)
+ latency_without_compile = metrics.latency
  return latency_with_compile - latency_without_compile
 
  def hw_roofline(self) -> float:
@@ -984,7 +984,7 @@ def _compile_time_in_task(
 
  def tflops(
  self, fn_name: str, example_inputs: Any, metrics: BenchmarkOperatorMetrics
- ) -> List[float]:
+ ) -> float:
  def _get_flops(self, func: Callable) -> float:
  """By default, use the torch.__dispatch__ based flops counter."""
  from torch.utils.flop_counter import FlopCounterMode
@@ -1010,4 +1010,4 @@ def work_func():
  if not fn in self._op_flops:
  self._op_flops[fn] = _get_flops(self, fn)
  op_flops = self._op_flops[fn]
- return list(map(lambda x: op_flops / x / 1e12 * 1e3, metrics.latency))
+ return op_flops / metrics.latency / 1e12 * 1e3