Add cutlass submodule to match the pytorch version

pytorch-labs · Nov 19, 2024 · c5c194d · c5c194d
1 parent 81d601e
commit c5c194d
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 5 deletions.
diff --git a/.gitmodules b/.gitmodules
@@ -19,3 +19,6 @@
 [submodule "submodules/xformers"]
 	path = submodules/xformers
 	url = https://github.com/facebookresearch/xformers.git
+[submodule "submodules/cutlass"]
+	path = submodules/cutlass
+	url = https://github.com/NVIDIA/cutlass.git
diff --git a/submodules/cutlass b/submodules/cutlass
diff --git a/tritonbench/utils/env_utils.py b/tritonbench/utils/env_utils.py
@@ -24,13 +24,13 @@
 
 
 def set_env():
-    import torch
-    IS_FBCODE = not hasattr(torch.version, "git_version")
     # set cutlass dir
     # by default we use the cutlass version built with fbgemm
-    if not "TORCHINDUCTOR_CUTLASS_DIR" in os.environ and not IS_FBCODE:
-        cutlass_dir = REPO_PATH.joinpath("submodules", "FBGEMM", "external", "cutlass")
-        os.environ["TORCHINDUCTOR_CUTLASS_DIR"] = str(cutlass_dir.absolute())
+    import torch
+    current_cutlass_dir = torch._inductor.config.cuda.cutlass_dir
+    if not os.path.exists(current_cutlass_dir):
+        tb_cutlass_dir = REPO_PATH.joinpath("submodules", "cutlass")
+        torch._inductor.config.cuda.cutlass_dir = str(tb_cutlass_dir)
 
 
 def set_random_seed():