"""Minimal multi-tile test via Python."""
import torch, sys, math
sys.path.insert(0, '/root/dsv4-nvfp4-workspace/kernel')

from dsv4.kernels.attention.fmha_multitile_op import fmha_multitile_decode_raw

torch.manual_seed(42)
hd = 64
N = 256
scale = 1.0 / math.sqrt(hd)

q = torch.randn(1, 1, 1, hd, dtype=torch.bfloat16, device='cuda').contiguous()
k = torch.randn(1, 1, N, hd, dtype=torch.bfloat16, device='cuda').contiguous()
v = torch.randn(1, 1, hd, N, dtype=torch.bfloat16, device='cuda').contiguous()

print(f'q align: {q.data_ptr() % 128}, k align: {k.data_ptr() % 128}, v align: {v.data_ptr() % 128}')
print(f'q shape: {q.shape}, k shape: {k.shape}, v shape: {v.shape}')

try:
    o, lse = fmha_multitile_decode_raw(q, k, v, scale)
    print(f'Output[0,0,0,:5]: {o[0,0,0,:5].float()}')
    print(f'LSE: {lse[0,0,0].item():.4f}')

    # Reference
    q_r = q[0,0].float()  # (1, hd)
    k_r = k[0,0].float()  # (N, hd)
    v_r = v[0,0].float().T  # (N, hd) — V is (hd, N), transpose for reference
    s = torch.matmul(q_r, k_r.T) * scale
    s = torch.softmax(s, dim=-1)
    o_ref = torch.matmul(s, v_r)
    cos = torch.nn.functional.cosine_similarity(o[0,0].float().flatten().unsqueeze(0), o_ref.flatten().unsqueeze(0)).item()
    print(f'Cosine vs reference: {cos:.6f}')
    print(f'{"PASS" if cos >= 0.999990 else "FAIL"}')
except Exception as e:
    print(f'FAILED: {e}')