auto: pre-test commit

tests/unit/test_fmha_pv16.py

@@ -1,9 +1,4 @@
-"""Test FMHA with pv_n_tile=16 (N=16 sub-tiles for PV GEMM).
-
-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=64, 128 with cosine >= 0.999
-"""
+"""Test FMHA with pv_n_tile=16 (N=16 sub-tiles for PV GEMM)."""
 import torch
 import math
 import sys
@@ -15,37 +10,33 @@ from dsv4.kernels.attention.production import dsv4_attention_per_head
 
 
 def test_fmha_pv16(hd):
-    """Test FMHA with pv_n_tile=16 at given head_dim"""
     sk = 128
     scale = 1.0 / math.sqrt(hd)
     
     torch.manual_seed(42)
-    # 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
+    q = torch.randn(1, 1, 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')
     
-    # FMHA kernel
     o = dsv4_attention_per_head(q, k, v, scale=scale, swa_len=sk)
     
     # Reference
-    q_ref = q.float()  # (1, hd)
-    k_ref = k.float()   # (sk, hd)
-    v_ref = v.float()   # (hd, sk)
+    q_ref = q[0, 0].float()
+    k_ref = k.float()
+    v_ref = v.float()
     
-    s = (q_ref @ k_ref.T) * scale  # (1, sk)
+    s = (q_ref @ k_ref.T) * scale
     p = torch.softmax(s, dim=-1)
-    o_ref = (p @ v_ref.T).to(torch.bfloat16)  # (1, hd)
+    o_ref = (p @ v_ref.T).to(torch.bfloat16)
     
-    # Compare
-    o_f = o.float().flatten()
-    o_ref_f = o_ref.float().flatten()
+    o_f = o[0, 0].float()
+    o_ref_f = o_ref.float()
     
     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}")
     
     if cs < 0.999:
-        print(f"  FAILED: cosine {cs} < 0.999")
+        print(f"  FAILED")
         print(f"  o[0:4]     = {o_f[0:4].tolist()}")
         print(f"  o_ref[0:4] = {o_ref_f[0:4].tolist()}")
         return False