extras for the various flash attn subdirs and build those in the base module as it is a slow step

This commit is contained in:
Wing Lian
2024-01-26 00:40:39 -05:00
parent a1d168d314
commit 34de5b3bd5
3 changed files with 13 additions and 3 deletions

View File

@@ -53,7 +53,13 @@ setup(
"flash-attn==2.3.3",
],
"fused-dense-lib": [
"fused-dense-lib @ git+https://github.com/Dao-AILab/flash-attention@v2.3.3#subdirectory=csrc/fused_dense_lib",
"fused-dense-lib @ git+https://github.com/Dao-AILab/flash-attention@v2.3.3#subdirectory=csrc/fused_dense_lib",
],
"dropout-layer-norm": [
"dropout-layer-norm @ git+https://github.com/Dao-AILab/flash-attention@v2.3.3#subdirectory=csrc/layer_norm",
],
"xentropy-cuda-lib": [
"xentropy_cuda_lib @ git+https://github.com/Dao-AILab/flash-attention.git@2.3.3#&subdirectory=csrc/xentropy",
],
"deepspeed": [
"deepspeed>=0.13.1",