save

ROCm · Aug 14, 2024 · 1c1d5f6 · 1c1d5f6
1 parent 26517a9
commit 1c1d5f6
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/python/perf-kernels/flash-attention.py b/python/perf-kernels/flash-attention.py
@@ -1069,7 +1069,11 @@ def varlen_input_helper(Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, dtype, equal_seqlen
 
 
 @pytest.mark.parametrize('Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD', [
+    (1, 1, 1, 512, 256, 128),
+    (1, 1, 1, 256, 128, 160),
     (1, 1, 1, 512, 256, 160),
+    (1, 1, 1, 1024, 512, 160),
+    (1, 1, 1, 512, 256, 256),
     (4, 48, 24, 1024, 1024, 64),
     (1, 24, 6, 8192, 8192, 64),
     (1, 4, 2, 16384, 16384, 128),