import torch
from torch.backends.cuda import sdp_kernel as torch_sdp_kernel
import paddle
from paddle.nn.functional import sdp_kernel as paddle_sdp_kernel
import numpy as np
from numpy.testing import assert_allclose
def create_float_attention_mask(batch_size, seq_lens, dtype='float16'):
    max_seq_len = max(seq_lens)
    
    # Paddle mask with 1e4
    paddle_mask = paddle.zeros([batch_size, 1, max_seq_len, max_seq_len], dtype=dtype)
    for i in range(batch_size):
        seq_len = seq_lens[i]
        mask = paddle.tril(paddle.ones(shape=(seq_len, seq_len), dtype=dtype)) - 1
        paddle_mask[i, 0, :seq_len, :seq_len] = mask * 1e4
    
    # Torch mask with 1e9
    torch_mask = paddle.zeros([batch_size, 1, max_seq_len, max_seq_len], dtype=dtype)
    for i in range(batch_size):
        seq_len = seq_lens[i]
        mask = paddle.tril(paddle.ones(shape=(seq_len, seq_len), dtype=dtype)) - 1
        torch_mask[i, 0, :seq_len, :seq_len] = mask * 1e9
    
    torch_mask = torch_mask.numpy()
    torch_mask = torch.from_numpy(torch_mask).to(dtype=torch.float16).cuda()
    
    return paddle_mask, torch_mask
def test_sdpa_alignment(batch_size=2, seq_len=1024, num_heads=8, head_dim=64):
    np.random.seed(42)
    torch.manual_seed(42)
    paddle.seed(42)
    
    query_np = np.random.randn(batch_size, seq_len, num_heads, head_dim).astype(np.float16)
    key_np = np.random.randn(batch_size, seq_len, num_heads, head_dim).astype(np.float16)
    value_np = np.random.randn(batch_size, seq_len, num_heads, head_dim).astype(np.float16)
    
    paddle_mask, torch_mask = create_float_attention_mask(batch_size, [seq_len] * batch_size, dtype='float16')
    
    query_paddle = paddle.to_tensor(query_np)
    key_paddle = paddle.to_tensor(key_np)
    value_paddle = paddle.to_tensor(value_np)
    
    query_torch = torch.from_numpy(query_np.transpose(0, 2, 1, 3)).cuda()
    key_torch = torch.from_numpy(key_np.transpose(0, 2, 1, 3)).cuda()
    value_torch = torch.from_numpy(value_np.transpose(0, 2, 1, 3)).cuda()
    
    # Flash Attention
    with paddle_sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False), \
         torch_sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
        try:
            paddle_flash = paddle.nn.functional.scaled_dot_product_attention(
                query_paddle, key_paddle, value_paddle, 
                attn_mask=paddle_mask, 
                dropout_p=0.0, 
                is_causal=False
            ).numpy()
            
            torch_flash = torch.nn.functional.scaled_dot_product_attention(
                query_torch, key_torch, value_torch, 
                attn_mask=None, 
                dropout_p=0.0,
                is_causal=True
            ).cpu().numpy().transpose(0, 2, 1, 3)
            
            print("\nComparing Flash Attention:")
            abs_diff = np.abs(paddle_flash - torch_flash)
            print(f"Maximum absolute difference: {np.max(abs_diff)}")
            print(f"Mean absolute difference: {np.mean(abs_diff)}")
            assert_allclose(paddle_flash, torch_flash, rtol=1e-2, atol=1e-2)
            print("Flash Attention results match within tolerance")
        except Exception as e:
            print(f"Flash Attention comparison failed: {e}")
    
    # Math Attention
    with paddle_sdp_kernel(enable_flash=False, enable_math=True, enable_mem_efficient=False), \
         torch_sdp_kernel(enable_flash=False, enable_math=True, enable_mem_efficient=False):
        try:
            paddle_math = paddle.nn.functional.scaled_dot_product_attention(
                query_paddle, key_paddle, value_paddle, 
                attn_mask=paddle_mask, 
                dropout_p=0.0, 
                is_causal=False
            ).numpy()
            
            torch_math = torch.nn.functional.scaled_dot_product_attention(
                query_torch, key_torch, value_torch, 
                attn_mask=torch_mask, 
                dropout_p=0.0,
                is_causal=False
            ).cpu().numpy().transpose(0, 2, 1, 3)
            
            print("\nComparing Math Attention:")
            abs_diff = np.abs(paddle_math - torch_math)
            print(f"Maximum absolute difference: {np.max(abs_diff)}")
            print(f"Mean absolute difference: {np.mean(abs_diff)}")
            assert_allclose(paddle_math, torch_math, rtol=1e-2, atol=1e-2)
            print("Math Attention results match within tolerance")
        except Exception as e:
            print(f"Math Attention comparison failed: {e}")
    
    # Memory Efficient Attention
    with paddle_sdp_kernel(enable_flash=False, enable_math=False, enable_mem_efficient=True), \
         torch_sdp_kernel(enable_flash=False, enable_math=False, enable_mem_efficient=True):
        try:
            paddle_mem = paddle.nn.functional.scaled_dot_product_attention(
                query_paddle, key_paddle, value_paddle, 
                attn_mask=paddle_mask, 
                dropout_p=0.0, 
                is_causal=False
            ).numpy()
            
            torch_mem = torch.nn.functional.scaled_dot_product_attention(
                query_torch, key_torch, value_torch, 
                attn_mask=torch_mask, 
                dropout_p=0.0,
                is_causal=False
            ).cpu().numpy().transpose(0, 2, 1, 3)
            
            print("\nComparing Memory Efficient Attention:")
            abs_diff = np.abs(paddle_mem - torch_mem)
            print(f"Maximum absolute difference: {np.max(abs_diff)}")
            print(f"Mean absolute difference: {np.mean(abs_diff)}")
            assert_allclose(paddle_mem, torch_mem, rtol=1e-2, atol=1e-2)
            print("Memory Efficient Attention results match within tolerance")
        except Exception as e:
            print(f"Memory Efficient Attention comparison failed: {e}")
if __name__ == "__main__":
    for seq_len in [512]:
        print(f"\nTesting with sequence length: {seq_len}")
        test_sdpa_alignment(
            batch_size=2, 
            seq_len=seq_len, 
            num_heads=8, 
            head_dim=64
        )