update fa bazel script

AlibabaPAI · Nov 5, 2024 · a25ce99 · a25ce99
1 parent 118e3ed
commit a25ce99
Showing 1 changed file with 12 additions and 5 deletions.
diff --git a/bazel/flash_attn.BUILD b/bazel/flash_attn.BUILD
@@ -21,11 +21,18 @@ cc_import(
 
 genrule(
     name = "build_flash_attn",
-    srcs = glob(["**"]),
+    srcs = glob(
+        ["**"],
+        exclude=[
+            "dist/**",  # Python whls
+            "build/**",  # Compiled intermediate files
+            "flash_attn.egg-info/**",  # Metadata about the whl package
+            ".*",  # Mainly .git* files
+            ]),
     outs = ["flash_attn_cuda.so"],
-    cmd = '&&'.join(['pushd external/flash_attn/',
-                    'MAX_JOBS=50 FLASH_ATTENTION_FORCE_BUILD=TRUE python setup.py bdist_wheel',
-                    'popd',
-                    'cp external/flash_attn/build/*/*.so $(location flash_attn_cuda.so)']),
+    cmd = "&&".join(["pushd external/flash_attn/",
+                     "MAX_JOBS=50 FLASH_ATTENTION_FORCE_BUILD=TRUE python setup.py bdist_wheel",
+                     "popd",
+                     "cp external/flash_attn/build/*/*.so $(location flash_attn_cuda.so)"]),
     visibility = ["//visibility:public"],
 )