auto: pre-test commit

tests/unit/test_fmha_pv16.py

@@ -2,7 +2,7 @@
 
 This tests the CuTeDSL FMHA kernel with the Layout D bug fix:
 - pv_n_tile=16 avoids the tcgen05.mma N=64 bug (missing TMEM columns)
-- Should work for HD=16, 64, 128, 256 with cosine ≥ 0.999
+- Should work for HD=64, 128 with cosine >= 0.999
 """
 import torch
 import math
@@ -11,77 +11,51 @@ import os
 
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..'))
 
-from dsv4.kernels.attention.fmha import FmhaKernel
 from dsv4.kernels.attention.production import dsv4_attention_per_head
 
 
-def test_fmha_pv16_hd64():
-    """HD=64 with pv_n_tile=16 (4 PV sub-tiles)"""
-    hd = 64
+def test_fmha_pv16(hd):
+    """Test FMHA with pv_n_tile=16 at given head_dim"""
     sk = 128
-    n_h = 1
     scale = 1.0 / math.sqrt(hd)
     
     torch.manual_seed(42)
-    q = torch.randn(1, sk, hd, dtype=torch.bfloat16, device='cuda')  # head-packed
+    # dsv4_attention_per_head expects q: (T, hd), k: (s_k, hd), v: (hd, s_k)
+    q = torch.randn(1, hd, dtype=torch.bfloat16, device='cuda')  # T=1 decode
     k = torch.randn(sk, hd, dtype=torch.bfloat16, device='cuda')
     v = torch.randn(hd, sk, dtype=torch.bfloat16, device='cuda')
-    o = torch.zeros(1, sk, hd, dtype=torch.bfloat16, device='cuda')
     
     # FMHA kernel
-    dsv4_attention_per_head(q, k, v, o, sk, scale, swa_len=sk, is_causal=False)
+    o = dsv4_attention_per_head(q, k, v, scale=scale, swa_len=sk)
     
     # Reference
-    q_ref = q[0].float()  # (sk, hd)
-    k_ref = k.float()      # (sk, hd)
-    v_ref = v.float().T     # (hd, sk) → need (sk, hd) for matmul
+    q_ref = q.float()  # (1, hd)
+    k_ref = k.float()   # (sk, hd)
+    v_ref = v.float()   # (hd, sk)
     
-    s = q_ref @ k_ref.T * scale  # (sk, sk)
-    s_max = s.max(dim=-1, keepdim=True).values
-    p = torch.softmax(s - s_max, dim=-1)
-    o_ref = (p @ v_ref.T).to(torch.bfloat16)  # (sk, hd) → bf16
+    s = (q_ref @ k_ref.T) * scale  # (1, sk)
+    p = torch.softmax(s, dim=-1)
+    o_ref = (p @ v_ref.T).to(torch.bfloat16)  # (1, hd)
     
-    # Compare row 0
-    o_row0 = o[0, 0].float()
-    o_ref0 = o_ref[0].float()
+    # Compare
+    o_f = o.float().flatten()
+    o_ref_f = o_ref.float().flatten()
     
-    cs = torch.nn.functional.cosine_similarity(o_row0.unsqueeze(0), o_ref0.unsqueeze(0)).item()
+    cs = torch.nn.functional.cosine_similarity(o_f.unsqueeze(0), o_ref_f.unsqueeze(0)).item()
     print(f"HD={hd} pv_n_tile=16: cosine={cs:.8f}")
-    assert cs > 0.999, f"Cosine {cs} < 0.999"
-    print("PASSED")
-
-
-def test_fmha_pv16_hd128():
-    """HD=128 with pv_n_tile=16 (8 PV sub-tiles)"""
-    hd = 128
-    sk = 128
-    n_h = 1
-    scale = 1.0 / math.sqrt(hd)
     
-    torch.manual_seed(42)
-    q = torch.randn(1, sk, hd, dtype=torch.bfloat16, device='cuda')
-    k = torch.randn(sk, hd, dtype=torch.bfloat16, device='cuda')
-    v = torch.randn(hd, sk, dtype=torch.bfloat16, device='cuda')
-    o = torch.zeros(1, sk, hd, dtype=torch.bfloat16, device='cuda')
-    
-    dsv4_attention_per_head(q, k, v, o, sk, scale, swa_len=sk, is_causal=False)
-    
-    # Reference
-    q_ref = q[0].float()
-    k_ref = k.float()
-    v_ref = v.float()
-    s = q_ref @ k_ref.T * scale
-    p = torch.softmax(s - s.max(dim=-1, keepdim=True).values, dim=-1)
-    o_ref = (p @ v_ref.T).to(torch.bfloat16)
-    
-    o_row0 = o[0, 0].float()
-    o_ref0 = o_ref[0].float()
-    cs = torch.nn.functional.cosine_similarity(o_row0.unsqueeze(0), o_ref0.unsqueeze(0)).item()
-    print(f"HD={hd} pv_n_tile=16: cosine={cs:.8f}")
-    assert cs > 0.999, f"Cosine {cs} < 0.999"
-    print("PASSED")
+    if cs < 0.999:
+        print(f"  FAILED: cosine {cs} < 0.999")
+        print(f"  o[0:4]     = {o_f[0:4].tolist()}")
+        print(f"  o_ref[0:4] = {o_ref_f[0:4].tolist()}")
+        return False
+    print(f"  PASSED")
+    return True
 
 
 if __name__ == '__main__':
-    test_fmha_pv16_hd64()
-    test_fmha_pv16_hd128()
+    all_pass = True
+    for hd in [16, 64, 128]:
+        if not test_fmha_pv16(hd):
+            all_pass = False
+    sys.exit(0 if all_pass else 1)