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
pytorch · Nov 1, 2024 · cc094df · cc094df
1 parent a66ce04
commit cc094df
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git 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
@@ -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)]