diff --git a/test/test_gpu/main.py b/test/test_gpu/main.py index 190085e2..2738ba77 100644 --- a/test/test_gpu/main.py +++ b/test/test_gpu/main.py @@ -37,8 +37,6 @@ TEST_OPERATORS = set(list_operators_by_collection(op_collection="default")) - SKIP_OPS -print(f"Testing operators: {TEST_OPERATORS}") - def check_ci_output(op): from tritonbench.utils.triton_op import REGISTERED_BENCHMARKS diff --git a/tritonbench/operators/addmm/hstu.py b/tritonbench/operators/addmm/hstu.py index 7af2a424..ddd08073 100644 --- a/tritonbench/operators/addmm/hstu.py +++ b/tritonbench/operators/addmm/hstu.py @@ -9,6 +9,7 @@ with add_path(str(SUBMODULE_PATH.joinpath("generative-recommenders"))): from generative_recommenders.ops.triton.triton_addmm import _addmm_fwd + class _AddMmFunction(torch.autograd.Function): @staticmethod # pyre-ignore[14] diff --git a/tritonbench/utils/triton_op.py b/tritonbench/utils/triton_op.py index a9d5b06e..785583c7 100644 --- a/tritonbench/utils/triton_op.py +++ b/tritonbench/utils/triton_op.py @@ -521,11 +521,13 @@ def __call__(cls, *args, **kwargs): obj.__post__init__() return obj + def _translate_mode(tb_args): def _has_and_true(attr): if hasattr(tb_args, attr) and getattr(tb_args, attr): return True return False + if _has_and_true("fwd"): tb_args.mode = "fwd" if _has_and_true("bwd"): @@ -535,6 +537,7 @@ def _has_and_true(attr): if _has_and_true("fwd_no_grad"): tb_args.mode = "fwd_no_grad" + class BenchmarkOperator(metaclass=PostInitProcessor): mode: Mode = Mode.FWD test: str = "eval"