#!/usr/bin/env python3
"""Test FMHA kernel with attention sink bias.

Validates that the kernel's sink bias correction matches PyTorch reference:
  softmax([QK^T * scale, sink_bias])[:N] @ V

Tests HD=64,128,256,512 with and without sinks.
"""
import torch
import math
import sys

def reference_fmha_with_sink(q, k, v, scale, sink_bias=None):
    """PyTorch reference: softmax([QK^T * scale, sink_bias]) @ V.
    
    q: (n_h, T, hd), k: (1, N, hd), v: (1, N, hd)
    sink_bias: (n_h,) FP32 or None
    Returns: (n_h, T, hd) BF16
    """
    n_h, T, hd = q.shape
    N = k.shape[1]
    # QK^T: (n_h, T, N)
    scores = torch.matmul(q, k.transpose(-1, -2)) * scale  # (n_h, T, N)
    
    if sink_bias is not None:
        # Concatenate sink as extra column: (n_h, T, N+1)
        sb = sink_bias.reshape(n_h, 1, 1).expand(-1, T, 1)
        combined = torch.cat([scores, sb], dim=-1)
        attn = torch.softmax(combined.float(), dim=-1)[:, :, :N]  # drop sink column
    else:
        attn = torch.softmax(scores.float(), dim=-1)
    
    out = torch.matmul(attn.bfloat16(), v)  # (n_h, T, hd)
    return out

def test_fmha_sink():
    from dsv4.kernels.attention.production import dsv4_attention
    
    torch.manual_seed(42)
    device = 'cuda'
    passed = 0
    failed = 0
    
    for hd in [64, 128, 256, 512]:
        for N in [9, 32, 128, 256]:
            for use_sink in [False, True]:
                n_h = 4  # small for speed
                T = 1
                scale = 1.0 / math.sqrt(hd)
                
                q = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device=device)
                k = torch.randn(1, N, hd, dtype=torch.bfloat16, device=device)
                v = torch.randn(1, N, hd, dtype=torch.bfloat16, device=device)
                sink = torch.randn(n_h, dtype=torch.float32, device=device) * 2 if use_sink else None
                
                # Production kernel
                try:
                    o_kernel = dsv4_attention(q, k, v, scale=scale, sink_bias=sink)
                except Exception as e:
                    print(f"  FAIL hd={hd} N={N} sink={use_sink}: kernel error: {e}")
                    failed += 1
                    continue
                
                # PyTorch reference
                o_ref = reference_fmha_with_sink(q, k, v, scale, sink)
                
                # Compare
                o_kf = o_kernel.float()
                o_rf = o_ref.float()
                cos = torch.nn.functional.cosine_similarity(o_kf.flatten().unsqueeze(0), 
                                                            o_rf.flatten().unsqueeze(0)).item()
                max_diff = (o_kf - o_rf).abs().max().item()
                
                status = "PASS" if cos > 0.999 else "FAIL"
                if status == "PASS":
                    passed += 1
                else:
                    failed += 1
                print(f"  {status} hd={hd} N={N} sink={use_sink} cos={cos:.6f} max_diff={max_diff:.6f}")
    
    print(f"\n{'='*60}")
    print(f"Results: {passed} PASSED, {failed} FAILED")
    print(f"{'='*60}")
    return failed == 0

if __name__ == "__main__":
    success = test_fmha_sink()
    sys.exit(0 if success else 1)