"""Test FMHA with pv_n_tile=16 (N=16 sub-tiles for PV GEMM)."""
import torch
import math
import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..'))

from dsv4.kernels.attention.production import dsv4_attention_per_head


def test_fmha_pv16(hd):
    sk = 128
    scale = 1.0 / math.sqrt(hd)
    
    torch.manual_seed(42)
    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')
    
    o = dsv4_attention_per_head(q, k, v, scale=scale, swa_len=sk)
    
    # Reference
    q_ref = q[0, 0].float()
    k_ref = k.float()
    v_ref = v.float()
    
    s = (q_ref @ k_ref.T) * scale
    p = torch.softmax(s, dim=-1)
    o_ref = (p @ v_ref.T).to(torch.bfloat16)
    
    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")
        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__':
    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)