flash attention 2

2023-07-20 00:00:49 -04:00
parent 06c61d6f13
commit 9b790d359b
2 changed files with 4 additions and 4 deletions
--- a/docker/Dockerfile-base
+++ b/docker/Dockerfile-base
@@ -40,7 +40,7 @@ ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6+PTX"

 RUN git clone https://github.com/Dao-AILab/flash-attention.git && \
    cd flash-attention && \
-    git checkout v1.0.9 && \
+    git checkout v2.0.0 && \
    python3 setup.py bdist_wheel && \
    cd csrc/fused_dense_lib && \
    python3 setup.py bdist_wheel && \