# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Adapted from https://github.com/vllm-project/vllm/blob/2c58742dff8613a3bd7496f2008ce927e18d38d1/tests/kernels/mamba/test_mamba_mixer2.py


import pytest
import torch

from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import (
    update_environment_variables,
)
from sglang.srt.distributed.parallel_state import (
    get_default_distributed_backend,
    init_distributed_environment,
    initialize_model_parallel,
)
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils import get_device, get_device_count
from sglang.test.ci.ci_register import register_cuda_ci, register_xpu_ci
from sglang.test.test_utils import publish_build_topology

register_cuda_ci(est_time=30, stage="base-b", runner_config="2-gpu-large")
register_xpu_ci(est_time=60, suite="nightly-xpu-2-gpu", nightly=True)

NUM_GPUS = 2


@pytest.mark.parametrize("batch_size", [8])
@pytest.mark.parametrize("seq_len", [128])
@pytest.mark.parametrize(
    "hidden_size_n_groups",
    [
        (64, 1),  # hidden_size be divisible by num_gpus
        (100, 4),  # and n_groups must divide hidden_size
    ],
)
@pytest.mark.parametrize("dtype", [torch.float16])
def test_mixer2_gated_norm_multi_gpu(
    batch_size: int,
    seq_len: int,
    hidden_size_n_groups: tuple[int, int],
    dtype: torch.dtype,
    device: str = get_device(),
):
    if device not in ["cuda", "xpu"]:
        pytest.skip("Test only supports CUDA and XPU devices")

    assert get_device_count() >= NUM_GPUS, (
        f"This test requires at least {NUM_GPUS} GPUs, but only {get_device_count()} available"
    )

    hidden_size, n_groups = hidden_size_n_groups
    num_processes = NUM_GPUS

    def run_torch_spawn(fn, nprocs):
        # need to use torch.mp.spawn otherwise will have problems with
        # torch.distributed and cuda
        torch.multiprocessing.spawn(
            fn,
            args=(
                num_processes,
                batch_size,
                seq_len,
                hidden_size,
                n_groups,
                dtype,
                device,
            ),
            nprocs=nprocs,
        )

    run_torch_spawn(mixer2_gated_norm_tensor_parallel, NUM_GPUS)


def mixer2_gated_norm_tensor_parallel(
    local_rank: int,
    world_size: int,
    batch_size: int,
    seq_len: int,
    hidden_size: int,
    n_groups: int,
    dtype: torch.dtype,
    device: str,
):
    torch.manual_seed(0)

    device = torch.device(get_device(local_rank))
    torch.get_device_module(device).set_device(device)
    torch.set_default_device(device)
    torch.set_default_dtype(dtype)

    update_environment_variables(
        {
            "RANK": str(local_rank),
            "LOCAL_RANK": str(local_rank),
            "WORLD_SIZE": str(world_size),
            "MASTER_ADDR": "localhost",
            "MASTER_PORT": "12345",
        }
    )

    # nccl on CUDA, xccl on XPU, ...; the parameter default is always "nccl".
    init_distributed_environment(
        world_size=world_size,
        rank=local_rank,
        local_rank=local_rank,
        backend=get_default_distributed_backend(device.type),
    )
    publish_build_topology(tp_size=world_size, world_rank=local_rank)
    initialize_model_parallel()

    # create random weights an inputs
    weight = torch.rand((hidden_size,), dtype=dtype, device=device)
    hidden_states = torch.randn(batch_size, seq_len, hidden_size)
    gate_states = torch.randn(batch_size, seq_len, hidden_size)

    import sglang.srt.layers.attention.mamba.mixer2_rms_norm_gated as m2

    # Force the TP topology through the context (the weight loader reads
    # get_parallel().attn_tp_rank, Mixer2RMSNormGated reads tp_size / tp_rank);
    # avoids calling initialize_dp_attention.
    with get_parallel().override(
        attn_tp_rank=local_rank, tp_size=world_size, tp_rank=local_rank
    ):
        # create gated-norm with TP
        mixer = m2.Mixer2RMSNormGated(
            full_hidden_size=hidden_size,
            full_n_groups=n_groups,
        )
        mixer.weight.weight_loader(mixer.weight, weight)

    # m2 reads tp via get_parallel().tp_size/rank — force it through the context.
    with get_parallel().override(tp_size=1, tp_rank=0):
        # create gated-norm without TP to compute reference
        mixer_single_gpu = m2.Mixer2RMSNormGated(
            full_hidden_size=hidden_size,
            full_n_groups=n_groups,
        )
        # assign weight to single-gpu mixer
        mixer_single_gpu.weight.data = weight

    # generate and compare
    N = hidden_size // world_size
    output = mixer(
        hidden_states[..., local_rank * N : (local_rank + 1) * N],
        gate_states[..., local_rank * N : (local_rank + 1) * N],
    )
    ref_output = mixer_single_gpu(hidden_states, gate_states)
    torch.testing.assert_close(
        output,
        ref_output[..., local_rank * N : (local_rank + 1) * N],
        atol=5e-3,
        rtol=1e-3,
    )


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main([__file__]))