From cc094dfec91c8f028d32a0fcaadc5a0ec09759e7 Mon Sep 17 00:00:00 2001 From: Xu Zhao Date: Thu, 31 Oct 2024 17:43:58 -0700 Subject: [PATCH] Add isolate mode Summary: When benchmarking across multiple operators, we can optionally isolate each operator run in a child process. Reviewed By: FindHao Differential Revision: D65154665 fbshipit-source-id: 9c9a21a76897084b061374cb3f7d8524a4aaac9b --- torchbenchmark/operators/fp8_gemm/fp8_gemm.py | 1 + torchbenchmark/operators/fp8_gemm_blockwise/operator.py | 1 + 2 files changed, 2 insertions(+) diff --git a/torchbenchmark/operators/fp8_gemm/fp8_gemm.py b/torchbenchmark/operators/fp8_gemm/fp8_gemm.py index 4ba5e283b7..d05f6d193c 100644 --- a/torchbenchmark/operators/fp8_gemm/fp8_gemm.py +++ b/torchbenchmark/operators/fp8_gemm/fp8_gemm.py @@ -47,6 +47,7 @@ def __init__( self, tb_args: argparse.Namespace, extra_args: Optional[List[str]] = None ): super().__init__(tb_args, extra_args) + self.use_cuda_graphs = True self.extra_args = parse_args(extra_args) def get_input_iter(self): diff --git a/torchbenchmark/operators/fp8_gemm_blockwise/operator.py b/torchbenchmark/operators/fp8_gemm_blockwise/operator.py index 265dfee975..78d67ba2a4 100644 --- a/torchbenchmark/operators/fp8_gemm_blockwise/operator.py +++ b/torchbenchmark/operators/fp8_gemm_blockwise/operator.py @@ -121,6 +121,7 @@ def __init__( self, tb_args: argparse.Namespace, extra_args: Optional[List[str]] = None ): super().__init__(tb_args, extra_args) + self.use_cuda_graphs = True addmm_args = parse_args(self.extra_args) if addmm_args.m and addmm_args.n and addmm_args.k: self.shapes = [(addmm_args.m, addmm_args.n, addmm_args.k)]