From 10b4b4a6b63ce58440133216ee1b75ad73ff07a2 Mon Sep 17 00:00:00 2001
From: Xu Zhao <i@xuzhao.net>
Date: Wed, 18 Dec 2024 23:54:12 -0500
Subject: [PATCH] Save memory to build

---
 .ci/tritonbench/install.sh | 3 +++
 install.py                 | 5 ++++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/.ci/tritonbench/install.sh b/.ci/tritonbench/install.sh
index c30ab81c..2e4a156f 100644
--- a/.ci/tritonbench/install.sh
+++ b/.ci/tritonbench/install.sh
@@ -10,5 +10,8 @@ fi
 tritonbench_dir=$(dirname "$(readlink -f "$0")")/../..
 cd ${tritonbench_dir}
 
+# probe memory available
+free -h
+
 # Install Tritonbench and all its customized packages
 python install.py --all
diff --git a/install.py b/install.py
index e0e227a4..ab6ded70 100644
--- a/install.py
+++ b/install.py
@@ -67,8 +67,11 @@ def install_fa2(compile=False):
 
 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"
     cmd = [sys.executable, "setup.py", "install"]
-    subprocess.check_call(cmd, cwd=str(FA3_PATH.resolve()))
+    subprocess.check_call(cmd, cwd=str(FA3_PATH.resolve()), env=env)
 
 
 def install_liger():