Test

pytorch-labs · Dec 19, 2024 · 9a9a342 · 9a9a342
1 parent 11bd39d
commit 9a9a342
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 6 deletions.
diff --git a/.ci/tritonbench/install.sh b/.ci/tritonbench/install.sh
@@ -10,9 +10,5 @@ fi
 tritonbench_dir=$(dirname "$(readlink -f "$0")")/../..
 cd ${tritonbench_dir}
 
-# probe memory available
-free -h
-
 # Install Tritonbench and all its customized packages
-# Test: only install fa3
-python install.py --fa3
+python install.py --all
diff --git a/install.py b/install.py
@@ -69,7 +69,7 @@ def install_fa3():
     FA3_PATH = REPO_PATH.joinpath("submodules", "flash-attention", "hopper")
     env = os.environ.copy()
     # nvcc will now spawn cicc and will cost ~1G memory
-    env["MAX_JOBS"] = "4"
+    env["MAX_JOBS"] = "8"
     cmd = [sys.executable, "setup.py", "install"]
     subprocess.check_call(cmd, cwd=str(FA3_PATH.resolve()), env=env)