Compile FA in local directory to enable docker image generation (#16)

AlibabaPAI · Oct 14, 2024 · 5d9590e · 5d9590e
1 parent fab18e0
commit 5d9590e
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/bazel/flash_attn.BUILD b/bazel/flash_attn.BUILD
@@ -23,8 +23,8 @@ genrule(
     name = "build_flash_attn",
     srcs = ["setup.py"],
     outs = ["flash_attn_cuda.so"],
-    cmd = ';'.join(['pushd external/flash_attn/',
-                    'FLASH_ATTENTION_FORCE_BUILD=TRUE python setup.py bdist_wheel',
+    cmd = ';'.join(['pushd third_party/flash-attention/',
+                    'MAX_JOBS=50 FLASH_ATTENTION_FORCE_BUILD=TRUE python setup.py bdist_wheel 2>&1 | tee build.log',
                     'popd',
-                    'cp external/flash_attn/build/*/*.so $(location flash_attn_cuda.so)']),
+                    'cp third_party/flash-attention/build/*/*.so $(location flash_attn_cuda.so)']),
 )