test: removed dropout test for varlen

ROCm · Nov 7, 2024 · 547b828 · 547b828
1 parent ea3e72d
commit 547b828
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/flash_attn/flash_attn_interface.py b/flash_attn/flash_attn_interface.py
@@ -1325,6 +1325,8 @@ def flash_attn_varlen_kvpacked_func(
     max_seqlen_q,
     max_seqlen_k,
     dropout_p=0.0,
+    dropout_philox_seed=0,
+    dropout_philox_offset=0,
     softmax_scale=None,
     causal=False,
     window_size=(-1, -1),  # -1 means infinite context window
@@ -1398,6 +1400,8 @@ def flash_attn_varlen_kvpacked_func(
         max_seqlen_q,
         max_seqlen_k,
         dropout_p,
+        dropout_philox_seed,
+        dropout_philox_offset,
         softmax_scale,
         causal,
         window_size,

diff --git a/tests/test_flash_attn_triton_amd.py b/tests/test_flash_attn_triton_amd.py
@@ -1268,6 +1268,9 @@ def test_flash_attn_varlen_output(
     seqlen_q, seqlen_k, d, dropout_p, causal, local, alibi, deterministic, mha_type, dtype, kvpacked, softcap
 ):
     if USE_TRITON_ROCM:
+
+        if dropout_p != 0.0:
+            pytest.skip("varlen does not support dropout on AMD Triton Backend currently.")
 
         if USE_REF and dropout_p != 0.0:
             pytest.skip("USE_REF not supported when Dropout is enabled.")
@@ -1347,6 +1350,8 @@ def test_flash_attn_varlen_output(
             max_seqlen_q,
             max_seqlen_k,
             dropout_p,
+            dropout_philox_seed=dropout_philox_seed,
+            dropout_philox_offset=dropout_philox_offset,
             causal=causal,
             window_size=window_size,
             softcap=softcap,