D1: Full test with TMEM-P at hd=64,128,256,512

tests/unit/test_fmha_v3_stage_d1.py

@@ -12,7 +12,7 @@ import cuda.bindings.driver as cuda
 from dsv4.kernels.attention.fmha import FmhaKernel
 
 
-def test_head_dim(hd, n_kv):
+def test_head_dim(hd, n_kv=128):
     m = 128
     torch.manual_seed(42)
 
@@ -21,31 +21,28 @@ def test_head_dim(hd, n_kv):
     v = torch.randn(n_kv, hd, dtype=torch.bfloat16, device='cuda')
     c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')
 
-    # FP32 reference (normalized)
+    # FP32 reference
     qf = q[:, :, 0].float()
     kf = k[:, :, 0].float()
     scale = 1.0 / math.sqrt(hd)
-    attn = qf @ kf.T * scale
-    attn_max = attn.max(dim=-1, keepdim=True)[0]
-    attn_exp = torch.exp(attn - attn_max)
+    attn_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0]
+    attn_exp = torch.exp(qf @ kf.T * scale - attn_max)
     attn_sum = attn_exp.sum(dim=-1, keepdim=True)
-    attn_norm = attn_exp / attn_sum
-    ref_norm = attn_norm @ v.float()
-
-    # FP32 reference (un-normalized): O_unnorm = sum(exp(S - max) * V)
+    ref_norm = (attn_exp / attn_sum) @ v.float()
     ref_unnorm = attn_exp @ v.float()
-
-    # Reference LSE: lse = ln(row_sum) + max
-    ref_lse = torch.log(attn_sum.squeeze(-1)) + attn_max.squeeze(-1)  # (m,)
+    ref_lse = (torch.log(attn_sum.squeeze(-1)) + attn_max.squeeze(-1))[0].item()
 
     lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
 
-    kernel = FmhaKernel(head_dim=hd, s_k=n_kv)
+    # Use TMEM-P for all hd (Saves SMEM, and works correctly after qk_mma_tiler fix)
+    # SMEM-P is available for hd>64 but requires more SMEM budget
+    kernel = FmhaKernel(head_dim=hd, s_k=n_kv, use_smem_p=False)
     pv_n_tile = kernel.pv_n_tile
     n_pv_tiles = kernel.n_pv_tiles
 
     stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
 
+    # Compile once
     v_tile = v[:, 0:pv_n_tile].contiguous()
     v_kernel = v_tile.unsqueeze(-1)
     c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
@@ -82,64 +79,34 @@ def test_head_dim(hd, n_kv):
             lse_val = lse_tensor[0, 0, 0].item()
 
     out_unnorm = c[:, :, 0].float()
-
-    # Compare un-normalized O against reference
+    out_norm = out_unnorm / attn_sum
     cos_unnorm = torch.nn.functional.cosine_similarity(
         out_unnorm.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0)
     ).item()
-
-    # Normalize externally: O_norm = O_unnorm / row_sum
-    # row_sum = exp(lse - max) where max is already incorporated in O_unnorm
-    # For the D5 merge, we use exp(lse) directly.
-    # For standalone test: O_norm = O_unnorm * (1/row_sum)
-    # where row_sum per row = O_unnorm row doesn't work. We need lse.
-    # lse = ln(row_sum) + max, so row_sum = exp(lse - max)
-    # But max is the same max used in the softmax, and O_unnorm already has
-    # the exp(-max) scaling baked in. So:
-    # O_norm = O_unnorm / row_sum
-    # We can compute row_sum from O_unnorm by checking what row_sum should be.
-    # Since O_unnorm[i,j] = sum_k(P[i,k] * V[k,j]) where P = exp(S*s - max)
-    # and row_sum = sum_k(exp(S*s - max)),
-    # we can normalize: O_norm[i] = O_unnorm[i] / row_sum[i]
-    # But we can't easily get row_sum from O_unnorm alone.
-    # Use LSE instead: row_sum = exp(lse - max_in_nat) 
-    # where max_in_nat = row_max * ln(2) but we only have lse.
-    # Actually for the merge: we just need exp(lse) * O_unnorm.
-    # For standalone: compute row_sum from attention explicitly.
-    # ref_row_sum = attn_sum.squeeze(-1)  # (m,)
-    # O_norm = O_unnorm / ref_row_sum.unsqueeze(1)
-    # This uses the reference row_sum to normalize — verifies the O_unnorm is correct.
-    out_norm_using_ref = out_unnorm / attn_sum  # (m, hd)
     cos_norm = torch.nn.functional.cosine_similarity(
-        out_norm_using_ref.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0)
+        out_norm.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0)
     ).item()
+    lse_err = abs(lse_val - ref_lse) if lse_val is not None else float('inf')
 
-    # Verify LSE
-    ref_lse_val = ref_lse[0].item()
-    lse_err = abs(lse_val - ref_lse_val) if lse_val is not None else float('inf')
-
-    status = "PASS" if cos_unnorm >= 0.99 else ("WARN" if cos_unnorm >= 0.97 else "FAIL")
-    print(f'hd={hd}, n={n_kv}: cos_unnorm {cos_unnorm:.6f}  cos_norm(ref_sum) {cos_norm:.6f}  lse_err {lse_err:.6f}  {status}')
+    status = "PASS" if cos_unnorm >= 0.99 else "FAIL"
+    print(f'hd={hd}, n={n_kv}: cos_unnorm {cos_unnorm:.6f}  cos_norm {cos_norm:.6f}  lse_err {lse_err:.6f}  {status}')
     return cos_unnorm
 
 
 def test():
     print("=== Stage D1: Parameterized HEAD_DIM ===")
-    print("(Kernel outputs un-normalized O + LSE)\n")
+    print("(Kernel outputs un-normalized O + LSE; TMEM-P path)\n")
 
-    print("--- Regression: HEAD_DIM=64 ---")
     cos64 = test_head_dim(64, 128)
-
-    print("\n--- HEAD_DIM=256 (single PV tile) ---")
+    cos128 = test_head_dim(128, 128)
     cos256 = test_head_dim(256, 128)
-
-    print("\n--- HEAD_DIM=512 (2 PV tiles) ---")
     cos512 = test_head_dim(512, 128)
 
     print("\n=== Summary ===")
-    print(f"hd=64,  n=128: cos_unnorm={cos64:.6f}  {'PASS' if cos64 >= 0.99 else 'FAIL'}")
-    print(f"hd=256, n=128: cos_unnorm={cos256:.6f}  {'PASS' if cos256 >= 0.99 else 'FAIL'}")
-    print(f"hd=512, n=128: cos_unnorm={cos512:.6f}  {'PASS' if cos512 >= 0.99 else 'FAIL'}")
+    print(f"hd=64:  cos={cos64:.6f}  {'PASS' if cos64 >= 0.99 else 'FAIL'}")
+    print(f"hd=128: cos={cos128:.6f}  {'PASS' if cos128 >= 0.99 else 'FAIL'}")
+    print(f"hd=256: cos={cos256:.6f}  {'PASS' if cos256 >= 0.99 else 'FAIL'}")
+    print(f"hd=512: cos={cos512:.6f}  {'PASS' if cos512 >= 0.99 else 'FAIL'}")
 
 
 if __name__ == '__main__':