# ruff: noqa
import torch
import tilelang
import triton
from tilelang import language as T
# from utils import assert_tensors_similar
from functools import partial
from aiter.ops.tilelang import tilelang_sparse_fwd, ref_sparse_mla_fwd_interface

def test_sparse_mla_fwd(B=1,
                        S=1,
                        SKV=8192,
                        H=128,
                        HKV=1,
                        DQK=576,
                        DV=512,
                        topk=2048,
                        dtype=torch.bfloat16,
                        check_correctness=True,
                        q_start_s_index=4096,
                        kv_stride=1,
                        sm_scale=None):
    torch.random.manual_seed(0)
    q = torch.randn((S, H, DQK), dtype=dtype, device="cuda")
    kv = torch.randn((SKV, HKV, DQK), dtype=dtype, device="cuda")

    indices = torch.full((S, HKV, topk), -1, dtype=torch.int32, device="cuda")
    for t in range(S):
        for h in range(HKV):
            i_i = torch.randperm(min(max(1, ((t + q_start_s_index) // kv_stride)), SKV))[:topk]
            indices[t, h, :len(i_i)] = i_i

    if dtype == torch.bfloat16:
        dtype_str = "bfloat16"
    else:
        dtype_str = "float16"

    if sm_scale is None:
        sm_scale = (1.0 / (DQK + (DQK - DV))) ** 0.5
    
    tl_out = tilelang_sparse_fwd(q, kv, indices, sm_scale, d_v=DV)

    print(f"{B=} {S=} {SKV=} {H=} {HKV=} {q_start_s_index=} {topk=} {dtype=}")

    if check_correctness:
        # otherwise may cause out of memory
        ref_out = ref_sparse_mla_fwd_interface(q, kv, indices, dtype, q_start_s_index=q_start_s_index, kv_stride=kv_stride, sm_scale=sm_scale)
        # assert_tensors_similar(tl_out, ref_out, eps=1e-2, name="out")
        atol=1e-2
        if dtype_str == "bfloat16":
            atol = 1e-2
        torch.testing.assert_close(tl_out, ref_out, atol=atol, rtol=1e-2)
        print("assert_tensors_similar passed")


if __name__ == "__main__":
    device_id = 0
    torch.cuda.set_device(device_id)
    for dtype in [torch.float16, torch.bfloat16]:
        test_sparse_mla_fwd(
            B=1,
            S=128,
            SKV=8192,
            H=16,
            HKV=1,
            DQK=576,
            DV=512,
            topk=2048,
            dtype=dtype,
            check_correctness=True,
            q_start_s_index=2048)

        test_sparse_mla_fwd(
            B=1,
            S=128,
            SKV=8192,
            H=16,
            HKV=1,
            DQK=576,
            DV=512,
            topk=2048,
            dtype=dtype,
            check_correctness=True,
            q_start_s_index=2048)

        test_sparse_mla_fwd(
            B=1,
            S=84,
            SKV=2048,
            H=16,
            HKV=1,
            DQK=576,
            DV=512,
            topk=2048,
            dtype=dtype,
            check_correctness=True,
            q_start_s_index=177)

        test_sparse_mla_fwd(
            B=1,
            S=284,
            SKV=1024,
            H=16,
            HKV=1,
            DQK=576,
            DV=512,
            topk=2048,
            dtype=dtype,
            check_correctness=True,
            q_start_s_index=0)

        test_sparse_mla_fwd(
            B=1,
            S=144,
            SKV=1024,
            H=16,
            HKV=1,
            DQK=576,
            DV=512,
            topk=2048,
            dtype=dtype,
            check_correctness=True,
            q_start_s_index=751)

        test_sparse_mla_fwd(
            B=1,
            S=256,
            SKV=996,
            H=128,
            HKV=1,
            DQK=576,
            DV=512,
            topk=2048,
            dtype=torch.float16,
            check_correctness=True,
            q_start_s_index=2048)

