Add more test case for fa3

ROCm · Jan 26, 2025 · e524b68 · e524b68
1 parent 17deba7
commit e524b68
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/test_flash_attn_ck_fa3.py b/tests/test_flash_attn_ck_fa3.py
@@ -68,10 +68,11 @@ def pad_rearrange_dropout_mask_hts_to_bhss(S_dmask, cu_seqlens_q, seqlen_q_round
 @pytest.mark.parametrize("alibi", [False])
 @pytest.mark.parametrize("local", [False])
 @pytest.mark.parametrize("causal", [False, True])
-@pytest.mark.parametrize("d", [64, 128])
+@pytest.mark.parametrize("d", [64, 72, 80, 88, 96, 104, 112, 120, 128])
 @pytest.mark.parametrize(
     "seqlen_q,seqlen_k",
     [
+        (256, 512),
         (1024, 1024),
         (2048, 2048),
     ],