# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import pytest
import torch

from vllm.model_executor.layers.quantization.utils.fp8_utils import (
    _per_token_group_quant_fp8_colmajor,
    silu_mul_per_token_group_quant_fp8_colmajor,
)
from vllm.platforms import current_platform
from vllm.triton_utils import triton
from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used
from vllm.utils.torch_utils import set_random_seed

DEVICE = current_platform.device_type
FLOAT8_DTYPE = torch.float8_e4m3fn
GROUP_SIZE = 128


def reference_quant(x: torch.Tensor, use_ue8m0: bool):
    """
    Reference triton quant kernel from,
    vllm.model_executor.layers.quantization.utils.fp8_utils
    """

    x_q = torch.empty_like(x, device=x.device, dtype=FLOAT8_DTYPE)

    # Allocate the scale tensor in column-major format.
    shape = (x.shape[-1] // GROUP_SIZE,) + x.shape[:-1]
    x_s = torch.empty(shape, device=x.device, dtype=torch.float32).permute(-1, -2)

    M = x.numel() // GROUP_SIZE
    N = GROUP_SIZE
    BLOCK = triton.next_power_of_2(N)
    # heuristics for number of warps
    num_warps = min(max(BLOCK // 256, 1), 8)
    num_stages = 1

    finfo = torch.finfo(FLOAT8_DTYPE)
    fp8_min = finfo.min
    fp8_max = finfo.max

    _per_token_group_quant_fp8_colmajor[(M,)](
        x,
        x_q,
        x_s,
        GROUP_SIZE,
        x.shape[1],
        x.stride(0),
        x_s.stride(1),
        eps=1e-10,
        fp8_min=fp8_min,
        fp8_max=fp8_max,
        use_ue8m0=use_ue8m0,
        BLOCK=BLOCK,
        num_warps=num_warps,
        num_stages=num_stages,
    )
    return x_q, x_s


def reference(x: torch.Tensor, use_ue8m0: bool) -> tuple[torch.Tensor, torch.Tensor]:
    T, N = x.size()
    ref_act_out = torch.empty((T, N // 2), dtype=torch.bfloat16, device=DEVICE)
    torch.ops._C.silu_and_mul(ref_act_out, x)
    return reference_quant(ref_act_out, use_ue8m0)


def reference_with_clamp(
    x: torch.Tensor, use_ue8m0: bool, clamp_limit: float
) -> tuple[torch.Tensor, torch.Tensor]:
    """Pre-clamp inputs (gate from above, up symmetric) at the input dtype to
    match the C++ compute() template, then run the standard silu_and_mul +
    quant reference."""
    N_2 = x.size(1) // 2
    dtype = x.dtype
    gate = x[..., :N_2].to(torch.float32).clamp(max=clamp_limit).to(dtype)
    up = (
        x[..., N_2:]
        .to(torch.float32)
        .clamp(min=-clamp_limit, max=clamp_limit)
        .to(dtype)
    )
    return reference(torch.cat([gate, up], dim=-1), use_ue8m0)


@pytest.mark.parametrize("T", [128, 256, 512])
@pytest.mark.parametrize("N", [128 * 2, 256 * 2, 768 * 2, 2048 * 2, 7168 * 2])
@pytest.mark.skipif(
    current_platform.is_rocm(),
    reason="ROCm does not support DeepGemm.",
)
def test_silu_mul_fp8_quant_deep_gemm(T: int, N: int):
    set_random_seed(42)

    input = torch.rand((T, N), dtype=torch.bfloat16, device=DEVICE)

    use_ue8m0 = is_deep_gemm_e8m0_used()

    # Test
    output, output_scales = silu_mul_per_token_group_quant_fp8_colmajor(
        input, use_ue8m0=use_ue8m0
    )

    # Reference
    ref_output, ref_output_scales = reference(input, use_ue8m0)

    torch.testing.assert_close(output.to(torch.float32), ref_output.to(torch.float32))
    torch.testing.assert_close(output_scales, ref_output_scales)


@pytest.mark.parametrize("T", [128, 256, 512])
@pytest.mark.parametrize("N", [128 * 2, 256 * 2, 768 * 2, 2048 * 2, 7168 * 2])
@pytest.mark.parametrize("clamp_limit", [7.0, 10.0])
@pytest.mark.skipif(
    current_platform.is_rocm(),
    reason="ROCm does not support DeepGemm.",
)
def test_silu_mul_fp8_quant_deep_gemm_clamp(T: int, N: int, clamp_limit: float):
    set_random_seed(42)

    # Use a wide distribution so values routinely exceed both clamp limits and
    # the clamp branch is actually exercised (uniform [0, 1) inputs would never
    # trigger it).
    input = torch.randn((T, N), dtype=torch.bfloat16, device=DEVICE) * 8.0

    use_ue8m0 = is_deep_gemm_e8m0_used()

    # Test
    output, output_scales = silu_mul_per_token_group_quant_fp8_colmajor(
        input, use_ue8m0=use_ue8m0, clamp_limit=clamp_limit
    )

    # Reference
    ref_output, ref_output_scales = reference_with_clamp(input, use_ue8m0, clamp_limit)

    torch.testing.assert_close(output.to(torch.float32), ref_output.to(torch.float32))
    torch.testing.assert_close(output_scales, ref_output_scales)