diff --git a/tritonbench/operators/gemm/operator.py b/tritonbench/operators/gemm/operator.py index d32f0a7c..5ac4b87c 100644 --- a/tritonbench/operators/gemm/operator.py +++ b/tritonbench/operators/gemm/operator.py @@ -24,12 +24,14 @@ from .kernels import matmul as kernels from .partition_k import matmul_partition_k + try: from .persistent_matmul import ( matmul_persistent, matmul_tma_persistent, matmul_tma_persistent_cached, ) + HAS_PRESISTENT = True except ModuleNotFoundError: HAS_PRESISTENT = False