From 1aa4a91d0138cd9eeeac0bd2fbe9d9e42d517cb1 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Fri, 22 May 2026 18:06:28 +0000
Subject: [PATCH] Diag: test all sizes 128-1024

---
 tests/unit/test_fmha_v3_diag.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/unit/test_fmha_v3_diag.py b/tests/unit/test_fmha_v3_diag.py
index d7d0f4c4..c6be5559 100644
--- 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')