import itertools
import sys

import pytest
import torch

from sglang.jit_kernel.utils import get_ci_test_range
from sglang.test.ci.ci_register import register_cuda_ci

register_cuda_ci(est_time=5, suite="stage-b-kernel-unit-1-gpu-large")
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)


def sglang_jit_fused_add_rmsnorm(
    input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float
) -> None:
    from sglang.jit_kernel.norm import fused_add_rmsnorm

    fused_add_rmsnorm(input, residual, weight, eps)


def flashinfer_fused_add_rmsnorm(
    input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float
) -> None:
    from flashinfer.norm import fused_add_rmsnorm

    fused_add_rmsnorm(input, residual, weight, eps=eps)


BS_LIST = [2**n for n in range(0, 14)]
BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)]
BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 4109])
HIDDEN_SIZE_LIST = get_ci_test_range(
    [512, 1024, 1536, 2048, 3072, 4096, 5120, 6144, 7168, 8192],
    [512, 2048, 8192],
)
DEVICE = "cuda"
DTYPE = torch.bfloat16


@pytest.mark.parametrize(
    "batch_size,hidden_size", list(itertools.product(BS_LIST, HIDDEN_SIZE_LIST))
)
def test_fused_add_rmsnorm(batch_size: int, hidden_size: int) -> None:
    input = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=DTYPE)
    residual = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=DTYPE)
    weight = torch.randn(hidden_size, device=DEVICE, dtype=DTYPE)

    input_sglang = input.clone()
    residual_sglang = residual.clone()
    input_flashinfer = input.clone()
    residual_flashinfer = residual.clone()
    sglang_jit_fused_add_rmsnorm(
        input_sglang, residual_sglang, weight, torch.finfo(torch.bfloat16).eps
    )
    flashinfer_fused_add_rmsnorm(
        input_flashinfer, residual_flashinfer, weight, torch.finfo(torch.bfloat16).eps
    )
    torch.testing.assert_close(input_sglang, input_flashinfer, atol=1e-2, rtol=1e-2)
    torch.testing.assert_close(
        residual_sglang, residual_flashinfer, atol=1e-2, rtol=1e-2
    )


if __name__ == "__main__":
    sys.exit(pytest.main([__file__, "-v", "-s"]))
