diff --git a/torch/_inductor/codegen/cuda/cutlass_utils.py b/torch/_inductor/codegen/cuda/cutlass_utils.py index ff60525548b41..62465b08833e9 100644 --- a/torch/_inductor/codegen/cuda/cutlass_utils.py +++ b/torch/_inductor/codegen/cuda/cutlass_utils.py @@ -10,6 +10,7 @@ import sympy import torch +from ... import config from ...config import cuda as inductor_cuda_config from ...ir import Layout @@ -48,6 +49,9 @@ def _gen_cutlass_file( @functools.lru_cache(None) def try_import_cutlass() -> bool: + if config.is_fbcode(): + return True + # Copy CUTLASS python scripts to a temp dir and add the temp dir to Python search path. # This is a temporary hack to avoid CUTLASS module naming conflicts. # TODO(ipiszy): remove this hack when CUTLASS solves Python scripts packaging structure issues. diff --git a/torch/_inductor/config.py b/torch/_inductor/config.py index f8953b874774a..f005bc1c6770d 100644 --- a/torch/_inductor/config.py +++ b/torch/_inductor/config.py @@ -9,6 +9,9 @@ def is_fbcode(): return not hasattr(torch.version, "git_version") +if is_fbcode(): + from triton.fb import build_paths + # add some debug printouts debug = False @@ -753,7 +756,9 @@ class cuda: # The default path only works under PyTorch local development environment. cutlass_dir = os.environ.get( "TORCHINDUCTOR_CUTLASS_DIR", - os.path.abspath( + build_paths.cutlass() + if is_fbcode() + else os.path.abspath( os.path.join(os.path.dirname(torch.__file__), "../third_party/cutlass/") ), )