From 5d9590eb9af0c51e0a82719ca06586790f456e50 Mon Sep 17 00:00:00 2001 From: Yitong Huang Date: Mon, 14 Oct 2024 13:47:00 +0800 Subject: [PATCH] Compile FA in local directory to enable docker image generation (#16) --- bazel/flash_attn.BUILD | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/bazel/flash_attn.BUILD b/bazel/flash_attn.BUILD index 6be811b826b..bfeada29778 100644 --- 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)']), )