update flash to 2.5.7

ghstack-source-id: 80926e210412d5e00ff714ecb208386b15ab2029 Pull Request resolved: fairinternal/xformers#1084 __original_commit__ = fairinternal/xformers@689b03b
robandpdx-org · May 2, 2024 · 82c6c7b · 82c6c7b
1 parent fd227d3
commit 82c6c7b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/third_party/flash-attention b/third_party/flash-attention
diff --git a/xformers/ops/fmha/flash.py b/xformers/ops/fmha/flash.py
@@ -49,8 +49,8 @@
         from flash_attn.flash_attn_interface import flash_attn_cuda as _C_flashattention
 
         FLASH_VERSION = flash_attn.__version__
-        FLASH_VER_MIN = (2, 5, 2)
-        FLASH_VER_LAST = (2, 5, 6)  # last supported, inclusive
+        FLASH_VER_MIN = (2, 5, 7)
+        FLASH_VER_LAST = (2, 5, 7)  # last supported, inclusive
         flash_ver_parsed = tuple(int(s) for s in FLASH_VERSION.split(".")[:3])
         if (
             flash_ver_parsed < FLASH_VER_MIN or flash_ver_parsed > FLASH_VER_LAST
@@ -145,6 +145,7 @@ def _flash_fwd(
                 cu_seq_lens_q,
                 cu_seq_lens_k,
                 seqused_k,
+                None,  # block_table
                 None,  # alibi_slopes
                 max_seq_len_q,
                 max_seq_len_k,
+2 −2		.github/workflows/publish.yml
+50 −12		csrc/flash_attn/flash_api.cpp
+34 −11		csrc/flash_attn/src/flash_bwd_launch_template.h
+53 −56		csrc/flash_attn/src/flash_fwd_kernel.h
+30 −9		csrc/flash_attn/src/flash_fwd_launch_template.h
+8 −1		csrc/flash_attn/src/softmax.h
+2 −1		csrc/ft_attention/setup.py
+2 −1		csrc/fused_dense_lib/setup.py
+2 −1		csrc/fused_softmax/setup.py
+2 −1		csrc/layer_norm/setup.py
+2 −1		csrc/rotary/setup.py
+2 −1		csrc/xentropy/setup.py
+1 −1		flash_attn/__init__.py
+9 −1		flash_attn/flash_attn_interface.py
+4 −0		flash_attn/modules/mha.py
+5 −5		flash_attn/ops/triton/layer_norm.py
+0 −13		flash_attn/ops/triton/rotary.py
+9 −3		setup.py
+60 −29		tests/test_flash_attn.py
+37 −0		tests/test_rotary.py
+2 −2		training/Dockerfile