Build fa3

pytorch-labs · Dec 20, 2024 · c1f8bb0 · c1f8bb0
1 parent 57150e9
commit c1f8bb0
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tools/flash_attn/install.py b/tools/flash_attn/install.py
@@ -38,7 +38,7 @@ def install_fa3():
     FA3_PATH = REPO_PATH.joinpath("submodules", "flash-attention", "hopper")
     env = os.environ.copy()
     # nvcc will spawn cicc process and will cost ~1G memory
-    env["MAX_JOBS"] = "8"
-    env["NVCC_THREADS"] = "1"
+    # env["MAX_JOBS"] = "8"
+    # env["NVCC_THREADS"] = "1"
     cmd = [sys.executable, "setup.py", "install"]
     subprocess.check_call(cmd, cwd=str(FA3_PATH.resolve()), env=env)