# sgl-kernel/tests/test_fp8_blockwise_gemm_deterministic.py
import os
import torch

from sgl_kernel.gemm import fp8_blockwise_scaled_mm


def _m_major_view_2d(t: torch.Tensor) -> torch.Tensor:
    # M major
    base = t.t().contiguous()
    return base.t()


def _col_major_view_k_n(mat_k_n: torch.Tensor) -> torch.Tensor:
    # Col major
    # by materializing a transposed contiguous base and returning the transpose view.
    base = mat_k_n.t().contiguous()     # (n, k), row-major contiguous
    return base.t()                     # (k, n), column-major view (stride(0) == 1)


def _make_fp8_with_scales(m, k, n, device, dtype_out):
    torch.manual_seed(1234)
    assert k % 128 == 0 and n % 128 == 0

    fp8_dtype = torch.float8_e4m3fn
    fp8_max = torch.finfo(fp8_dtype).max

    # A: (m, k), per-128 block scales along K
    a_bf16 = torch.randn(m, k, device=device, dtype=torch.bfloat16)
    a_scales_blocks = (
        a_bf16.abs().view(m, k // 128, 128).amax(dim=2) / fp8_max
    ).clamp_min(1e-8).to(torch.float32).contiguous()
    a_scale_full = a_scales_blocks.repeat_interleave(128, dim=1)
    a_fp8 = torch.clamp(
        (a_bf16 / a_scale_full).to(torch.float32), min=-fp8_max, max=fp8_max
    ).to(fp8_dtype).contiguous()
    # M major
    a_scales = _m_major_view_2d(a_scales_blocks)
    assert a_scales.size() == (m, k // 128) and a_scales.stride(0) == 1

    # B: (k, n), per-128x128 block scales
    b_bf16 = torch.randn(k, n, device=device, dtype=torch.bfloat16)
    b_scales_blocks = (
        b_bf16.abs().view(k // 128, 128, n // 128, 128).amax(dim=(1, 3))
    ).clamp_min(1e-8).to(torch.float32).contiguous()
    b_scale_full = b_scales_blocks.repeat_interleave(128, dim=0).repeat_interleave(128, dim=1)
    b_fp8_k_n = torch.clamp(
        (b_bf16 / b_scale_full).to(torch.float32), min=-fp8_max, max=fp8_max
    ).to(fp8_dtype)
    # Col major
    b_fp8 = _col_major_view_k_n(b_fp8_k_n)
    assert b_fp8.size() == (k, n) and b_fp8.stride(0) == 1

    b_scales = _m_major_view_2d(b_scales_blocks)
    assert b_scales.size() == (k // 128, n // 128) and b_scales.stride(0) == 1

    return a_fp8, a_scales, b_fp8, b_scales, dtype_out


def _run_once(m, k, n, device, dtype_out):
    a_fp8, a_scales, b_fp8, b_scales, out_dtype = _make_fp8_with_scales(
        m, k, n, device, dtype_out
    )
    # print(a_fp8.shape, b_fp8.shape, a_scales.shape, b_scales.shape)
    return fp8_blockwise_scaled_mm(a_fp8, b_fp8, a_scales, b_scales, out_dtype)


def test_fp8_blockwise_gemm_deterministic_sm90():
    if not torch.cuda.is_available():
        return
    dev = torch.device("cuda")
    os.environ["SGLANG_ENABLE_DETERMINISTIC_INFERENCE"] = "1"

    shapes = [(1024, 4096, 1024), (2048, 2048, 2048)]
    dtypes = [torch.bfloat16, torch.half]

    for (m, k, n) in shapes:
        for out_dtype in dtypes:
            y0 = _run_once(m, k, n, dev, out_dtype)
            for _ in range(5):
                y = _run_once(m, k, n, dev, out_dtype)
                assert torch.equal(y0, y), f"Non-deterministic result for shape {(m,k,n)} dtype {out_dtype}"