feat: add support for flash_attn>=2.6.0 (#70)

feifeibear · Aug 26, 2024 · 4bfd585 · 4bfd585
1 parent 478c3a2
commit 4bfd585
Show file tree

Hide file tree

Showing 24 changed files with 121 additions and 7 deletions.
diff --git a/benchmark/benchmark_longctx.py b/benchmark/benchmark_longctx.py
@@ -157,6 +157,7 @@ def benchmark(num_iter=100, forward_only=True, log=True, profile=False):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=False,
@@ -171,6 +172,7 @@ def benchmark(num_iter=100, forward_only=True, log=True, profile=False):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=False,
@@ -194,6 +196,7 @@ def benchmark(num_iter=100, forward_only=True, log=True, profile=False):
                         dropout_p=dropout_p,
                         causal=causal,
                         window_size=(-1, -1),
+                        softcap=0.0,
                         alibi_slopes=None,
                         deterministic=deterministic,
                         return_attn_probs=False,
@@ -215,6 +218,7 @@ def benchmark(num_iter=100, forward_only=True, log=True, profile=False):
                     dropout_p=dropout_p,
                     causal=causal,
                     window_size=(-1, -1),
+                    softcap=0.0,
                     alibi_slopes=None,
                     deterministic=deterministic,
                     return_attn_probs=False,

diff --git a/benchmark/benchmark_longctx_qkvpacked.py b/benchmark/benchmark_longctx_qkvpacked.py
@@ -117,6 +117,7 @@ def benchmark(num_iter=100, forward_only=True, log=True):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=False,
@@ -134,6 +135,7 @@ def benchmark(num_iter=100, forward_only=True, log=True):
                     dropout_p=dropout_p,
                     causal=causal,
                     window_size=(-1, -1),
+                    softcap=0.0,
                     alibi_slopes=None,
                     deterministic=deterministic,
                     return_attn_probs=False,
@@ -147,6 +149,7 @@ def benchmark(num_iter=100, forward_only=True, log=True):
                 dropout_p=dropout_p,
                 causal=causal,
                 window_size=(-1, -1),
+                softcap=0.0,
                 alibi_slopes=None,
                 deterministic=deterministic,
                 return_attn_probs=False,

diff --git a/benchmark/benchmark_qkvpacked_func.py b/benchmark/benchmark_qkvpacked_func.py
@@ -56,6 +56,7 @@ def benchmark(f, num_iter=100, forward_only=True, log=True):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=False,
@@ -65,6 +66,7 @@ def benchmark(f, num_iter=100, forward_only=True, log=True):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=False,
@@ -82,6 +84,7 @@ def benchmark(f, num_iter=100, forward_only=True, log=True):
                     dropout_p=dropout_p,
                     causal=causal,
                     window_size=(-1, -1),
+                    softcap=0.0,
                     alibi_slopes=None,
                     deterministic=deterministic,
                     return_attn_probs=False,
@@ -95,6 +98,7 @@ def benchmark(f, num_iter=100, forward_only=True, log=True):
                 dropout_p=dropout_p,
                 causal=causal,
                 window_size=(-1, -1),
+                softcap=0.0,
                 alibi_slopes=None,
                 deterministic=deterministic,
                 return_attn_probs=False,

diff --git a/benchmark/benchmark_ring_func.py b/benchmark/benchmark_ring_func.py
@@ -79,6 +79,7 @@ def benchmark(f, num_iter=100, forward_only=True, log=True):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=False,
@@ -90,6 +91,7 @@ def benchmark(f, num_iter=100, forward_only=True, log=True):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=False,
@@ -109,6 +111,7 @@ def benchmark(f, num_iter=100, forward_only=True, log=True):
                     dropout_p=dropout_p,
                     causal=causal,
                     window_size=(-1, -1),
+                    softcap=0.0,
                     alibi_slopes=None,
                     deterministic=deterministic,
                     return_attn_probs=False,
@@ -126,6 +129,7 @@ def benchmark(f, num_iter=100, forward_only=True, log=True):
                 dropout_p=dropout_p,
                 causal=causal,
                 window_size=(-1, -1),
+                softcap=0.0,
                 alibi_slopes=None,
                 deterministic=deterministic,
                 return_attn_probs=False,

diff --git a/benchmark/benchmark_varlen_qkvpacked_func.py b/benchmark/benchmark_varlen_qkvpacked_func.py
@@ -64,6 +64,7 @@ def benchmark(f, num_iter=100, forward_only=True, log=True):
                     dropout_p=dropout_p,
                     causal=causal,
                     window_size=(-1, -1),
+                    softcap=0.0,
                     alibi_slopes=None,
                     deterministic=deterministic,
                     return_attn_probs=False,
@@ -78,6 +79,7 @@ def benchmark(f, num_iter=100, forward_only=True, log=True):
                 dropout_p=dropout_p,
                 causal=causal,
                 window_size=(-1, -1),
+                softcap=0.0,
                 alibi_slopes=None,
                 deterministic=deterministic,
                 return_attn_probs=False,

diff --git a/scripts/run_gqa.sh b/scripts/run_gqa.sh
diff --git a/scripts/run_qkvpack_compare.sh b/scripts/run_qkvpack_compare.sh
diff --git a/setup.py b/setup.py
@@ -2,11 +2,11 @@
 
 setup(
     name="yunchang",
-    version="0.2",
+    version="0.3",
     author="Jiarui Fang, Zilin Zhu, Yang Yu",
     url="https://github.com/feifeibear/long-context-attention",
     packages=find_packages(exclude=['test', 'benchmark']),
     install_requires=[
-        'flash-attn',
+        'flash-attn>=2.6.0',
     ],
 )
diff --git a/test/test_hybrid_attn.py b/test/test_hybrid_attn.py
@@ -110,6 +110,7 @@ def log(msg, a, rank0_only=False):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=True,
@@ -137,6 +138,7 @@ def log(msg, a, rank0_only=False):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=True,

diff --git a/test/test_hybrid_qkvpacked_attn.py b/test/test_hybrid_qkvpacked_attn.py
@@ -111,6 +111,7 @@ def test(ring_impl_type="basic"):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=True,
@@ -124,6 +125,7 @@ def test(ring_impl_type="basic"):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=True,

diff --git a/test/test_ring_flash_attn_func.py b/test/test_ring_flash_attn_func.py
@@ -73,6 +73,7 @@ def log(msg, a, rank0_only=False):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=True,
@@ -88,6 +89,7 @@ def log(msg, a, rank0_only=False):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=True,

diff --git a/test/test_ring_flash_attn_varlen_func.py b/test/test_ring_flash_attn_varlen_func.py
@@ -99,6 +99,7 @@ def extract_lse(lse, cu_seqlens):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=True,
@@ -114,6 +115,7 @@ def extract_lse(lse, cu_seqlens):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=True,

diff --git a/test/test_stripe_flash_attn_func.py b/test/test_stripe_flash_attn_func.py
@@ -80,6 +80,7 @@ def extract_local(value, rank, world_size, dim=1):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=True,
@@ -93,6 +94,7 @@ def extract_local(value, rank, world_size, dim=1):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=True,

diff --git a/test/test_ulysses_attn.py b/test/test_ulysses_attn.py
@@ -93,6 +93,7 @@ def log(msg, a, rank0_only=False):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=True,
@@ -119,6 +120,7 @@ def log(msg, a, rank0_only=False):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=True,

diff --git a/test/test_zigzag_ring_flash_attn_func.py b/test/test_zigzag_ring_flash_attn_func.py
@@ -80,6 +80,7 @@ def extract_local(value, rank, world_size, dim=1):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=True,
@@ -93,6 +94,7 @@ def extract_local(value, rank, world_size, dim=1):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=True,

diff --git a/test/test_zigzag_ring_flash_attn_varlen_func.py b/test/test_zigzag_ring_flash_attn_varlen_func.py
@@ -104,6 +104,7 @@ def extract_lse(lse, cu_seqlens):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=True,
@@ -119,6 +120,7 @@ def extract_lse(lse, cu_seqlens):
         dropout_p=dropout_p,
         causal=causal,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=deterministic,
         return_attn_probs=True,

diff --git a/yunchang/hybrid/async_attn_layer.py b/yunchang/hybrid/async_attn_layer.py
@@ -50,6 +50,7 @@ def forward(
         softmax_scale=None,
         causal=False,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=False,
         return_attn_probs=False,
@@ -148,6 +149,7 @@ def forward(
                 softmax_scale=softmax_scale,
                 causal=causal,
                 window_size=window_size,
+                softcap=softcap,
                 alibi_slopes=alibi_slopes,
                 deterministic=deterministic,
                 return_attn_probs=return_attn_probs,

diff --git a/yunchang/hybrid/attn_layer.py b/yunchang/hybrid/attn_layer.py
@@ -49,6 +49,7 @@ def forward(
         softmax_scale=None,
         causal=False,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=False,
         return_attn_probs=False,
@@ -84,6 +85,7 @@ def forward(
                 softmax_scale=softmax_scale,
                 causal=causal,
                 window_size=window_size,
+                softcap=softcap,
                 alibi_slopes=alibi_slopes,
                 deterministic=deterministic,
                 return_attn_probs=return_attn_probs,
@@ -108,6 +110,7 @@ def forward(
                 softmax_scale=softmax_scale,
                 causal=causal,
                 window_size=window_size,
+                softcap=softcap,
                 alibi_slopes=alibi_slopes,
                 deterministic=deterministic,
                 return_attn_probs=return_attn_probs,
@@ -166,6 +169,7 @@ def forward(
         softmax_scale=None,
         causal=False,
         window_size=(-1, -1),
+        softcap=0.0,
         alibi_slopes=None,
         deterministic=False,
         return_attn_probs=False,
@@ -198,6 +202,7 @@ def forward(
             softmax_scale=softmax_scale,
             causal=causal,
             window_size=window_size,
+            softcap=softcap,
             alibi_slopes=alibi_slopes,
             deterministic=deterministic,
             return_attn_probs=return_attn_probs,