Diag: test all sizes 128-1024

2026-05-22 18:06:28 +00:00
parent 8586603280
commit 1aa4a91d01
1 changed files with 1 additions and 1 deletions
--- a/tests/unit/test_fmha_v3_diag.py
+++ b/tests/unit/test_fmha_v3_diag.py
@@ -266,7 +266,7 @@ class FmhaV3Diag:


 def test():
-    for n in [128, 256]:
+    for n in [128, 256, 512, 1024]:
        torch.manual_seed(42)
        m, hd = 128, HEAD_DIM
        q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda')