# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from unittest.mock import patch

import pytest
import torch

from vllm.model_executor.layers.quantization.utils import fp8_utils, int8_utils
from vllm.model_executor.layers.quantization.utils.quant_utils import get_fp8_min_max
from vllm.platforms import current_platform


@pytest.mark.parametrize(
    "shape", [(31, 128), (32, 128), (63, 256), (64, 256), (16, 512)]
)
@pytest.mark.parametrize("column_major", [False, True])
@pytest.mark.parametrize("tma_aligned", [False, True])
@pytest.mark.parametrize("scale_ue8m0", [False, True])
@pytest.mark.parametrize("group_size", [64, 128])
@pytest.mark.skipif(
    not current_platform.is_cuda_alike(), reason="Only test on CUDA/ROCm."
)
def test_per_token_group_quant_fp8(
    shape, column_major: bool, tma_aligned: bool, scale_ue8m0: bool, group_size: int
):
    device = "cuda"

    torch.manual_seed(42)
    num_tokens, hidden_dim = shape

    x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8

    # cuda path
    out_q, scale = fp8_utils.per_token_group_quant_fp8(
        x,
        group_size,
        column_major_scales=column_major,
        tma_aligned_scales=tma_aligned,
        use_ue8m0=scale_ue8m0,
    )

    # triton ref
    with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
        ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(
            x,
            group_size,
            column_major_scales=column_major,
            use_ue8m0=scale_ue8m0,
        )

    assert torch.allclose(out_q.float(), ref_q.float(), atol=0.15, rtol=0.15)
    assert torch.allclose(scale, ref_s, atol=0.01, rtol=0.01)


@pytest.mark.parametrize(
    "num_tokens,hidden_dim,group_size",
    [
        # No padding: mn=4 (mult of 4), groups_per_row=56 (mult of 4)
        (4, 7168, 128),
        # MN padding only: mn=1, tma_aligned_mn=4
        (1, 7168, 128),
        # MN padding only: mn=3, tma_aligned_mn=4
        (3, 7168, 128),
        # K padding only: groups_per_row=5 (5%4=1)
        (4, 640, 128),
        # K padding only: groups_per_row=6 (6%4=2)
        (4, 768, 128),
        # Single packed column, no padding: k_num_packed=1, mn%4=0
        (4, 384, 128),
        # Both MN and K padding
        (1, 384, 128),
        (3, 640, 128),
        # Larger shapes with no padding
        (64, 7168, 128),
        (128, 14336, 128),
        # Larger shapes with padding
        (127, 7168, 128),
        (253, 640, 128),
    ],
)
@pytest.mark.parametrize("poisoned_scales", [False, True])
@pytest.mark.skipif(
    not current_platform.is_cuda_alike(),
    reason="DeepGEMM not available on this platform",
)
def test_per_token_group_quant_fp8_packed(
    num_tokens, hidden_dim, group_size, poisoned_scales
):
    """Test the packed DeepGEMM quantization kernel against the Triton
    reference (row-major, UE8M0 scales)."""

    device = "cuda"
    torch.manual_seed(42)

    x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8

    mn = num_tokens
    groups_per_row = hidden_dim // group_size
    k_num_packed = (groups_per_row + 3) // 4
    tma_aligned_mn = ((mn + 3) // 4) * 4
    num_scale_elems = mn + (k_num_packed - 1) * tma_aligned_mn

    if poisoned_scales:
        # Call the kernel with poisoned scale buffer to
        # ensure padded indices are correctly zeroed.
        fp8_dtype = current_platform.fp8_dtype()
        fp8_min, fp8_max = get_fp8_min_max()
        out_q = torch.empty_like(x, dtype=fp8_dtype)
        out_s_packed = torch.empty_strided(
            (mn, k_num_packed),
            (1, tma_aligned_mn),
            device=device,
            dtype=torch.int32,
        )
        torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).fill_(0x7F7F7F7F)
        torch.ops._C.per_token_group_fp8_quant_packed(
            x,
            out_q,
            out_s_packed,
            group_size,
            1e-10,
            fp8_min,
            fp8_max,
        )
    else:
        out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
            x,
            group_size=group_size,
            use_ue8m0=True,
        )

    # Triton reference (row-major float32 scales, UE8M0)
    with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
        ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(
            x,
            group_size,
            use_ue8m0=True,
        )

    # Quantized values must match.
    assert torch.equal(out_q, ref_q), "Quantized output mismatch"

    # Verify packed scales (valid exponents + padding zeros).
    ref_s_flat = ref_s.reshape(mn, groups_per_row)
    ref_exponents = (ref_s_flat.view(torch.int32) >> 23) & 0xFF

    expected = torch.zeros(num_scale_elems, dtype=torch.int32, device="cpu")
    for row in range(mn):
        for g in range(groups_per_row):
            pack_col = g // 4
            pos = g % 4
            idx = pack_col * tma_aligned_mn + row
            expected[idx] |= int(ref_exponents[row, g].item()) << (pos * 8)

    actual = torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).cpu()
    assert torch.equal(actual, expected), (
        f"Packed scale storage mismatch.\n"
        f"First diff at index "
        f"{(actual != expected).nonzero(as_tuple=True)[0][0].item()}"
    )


@pytest.mark.skipif(
    not current_platform.is_cuda_alike(),
    reason="DeepGEMM not available on this platform",
)
def test_per_token_group_quant_fp8_packed_all_zero():
    """All-zero input must produce well-defined UE8M0 scale bytes via the eps
    floor in the kernel's UE8M0 path. Locks down the all-zero behavior before
    optimization.

    The CUDA kernel computes:
        y_s = eps / fp8_max
        y_s = exp2(ceil(log2(fmax(y_s, 1e-10))))
    For all-zero input, eps/fp8_max < 1e-10, so the inner fmax clamps back to
    1e-10, giving exp2(ceil(log2(1e-10))) = exp2(-33) => UE8M0 byte 0x5E (94).
    """

    device = "cuda"
    num_tokens, hidden_dim, group_size = 4, 7168, 128
    x = torch.zeros((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16)

    out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
        x,
        group_size=group_size,
        use_ue8m0=True,
    )

    # Quantized values must be all zero.
    assert torch.equal(
        out_q.view(torch.uint8),
        torch.zeros_like(out_q, dtype=torch.uint8),
    ), "All-zero input should produce all-zero FP8 output"

    # UE8M0 byte produced by the kernel for all-zero input.
    # The kernel's inner fmax(y_s, 1e-10) clamps eps/fp8_max back to 1e-10.
    # 1e-10 as float32 has biased exponent 0x5D and a non-zero mantissa, so
    # the kernel's bit-twiddle (exp_bits + (mant_bits != 0)) rounds up to
    # 0x5E. This matches exp2(ceil(log2(1e-10))) = exp2(-33).
    expected_exp_byte = 0x5E

    mn = num_tokens
    groups_per_row = hidden_dim // group_size
    k_num_packed = (groups_per_row + 3) // 4
    tma_aligned_mn = ((mn + 3) // 4) * 4
    num_scale_elems = mn + (k_num_packed - 1) * tma_aligned_mn

    # All valid scale slots must contain the expected packed value.
    # Padding slots must be zero.
    actual = torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).cpu()

    expected = torch.zeros(num_scale_elems, dtype=torch.int32, device="cpu")
    for row in range(mn):
        for g in range(groups_per_row):
            pack_col = g // 4
            pos = g % 4
            idx = pack_col * tma_aligned_mn + row
            expected[idx] |= expected_exp_byte << (pos * 8)

    assert torch.equal(actual, expected), "All-zero scale bytes mismatch"


@pytest.mark.skipif(
    not current_platform.is_cuda_alike(),
    reason="DeepGEMM not available on this platform",
)
def test_per_token_group_quant_fp8_packed_mantissa_rounds_up():
    """Inputs whose absmax/max_8bit produces a non-power-of-2 force the
    mantissa-rounding-up branch (exp_byte += 1). Locks down this behavior
    before optimization."""

    device = "cuda"
    num_tokens, hidden_dim, group_size = 4, 7168, 128

    # Build a tensor whose per-group absmax = 1.5 * fp8_max * 2^k for various k.
    # fp8_max = torch.finfo(torch.float8_e4m3fn).max = 448.0.
    # Then absmax/fp8_max = 1.5 * 2^k -> non-zero mantissa, triggers ceil
    # rounding to 2^(k+1). Use k=0 for simplicity; the bf16 representation of
    # 1.5*448=672.0 is exact.
    x = torch.full(
        (num_tokens, hidden_dim),
        672.0,
        device=device,
        dtype=torch.bfloat16,
    )

    out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
        x,
        group_size=group_size,
        use_ue8m0=True,
    )

    with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
        ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(
            x,
            group_size,
            use_ue8m0=True,
        )

    assert torch.equal(out_q, ref_q), "Quantized output mismatch"

    mn = num_tokens
    groups_per_row = hidden_dim // group_size
    k_num_packed = (groups_per_row + 3) // 4
    tma_aligned_mn = ((mn + 3) // 4) * 4
    num_scale_elems = mn + (k_num_packed - 1) * tma_aligned_mn

    ref_s_flat = ref_s.reshape(mn, groups_per_row)
    ref_exponents = (ref_s_flat.view(torch.int32) >> 23) & 0xFF
    expected = torch.zeros(num_scale_elems, dtype=torch.int32, device="cpu")
    for row in range(mn):
        for g in range(groups_per_row):
            pack_col = g // 4
            pos = g % 4
            idx = pack_col * tma_aligned_mn + row
            expected[idx] |= int(ref_exponents[row, g].item()) << (pos * 8)

    actual = torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).cpu()
    assert torch.equal(actual, expected), "Scale bytes mismatch"


@pytest.mark.parametrize(
    "num_tokens,hidden_dim",
    [
        (1, 7168),  # mn padded 1 -> 4
        (2, 7168),  # mn padded 2 -> 4
        (3, 7168),  # mn padded 3 -> 4
        (5, 7168),  # mn padded 5 -> 8
        (127, 7168),  # mn padded 127 -> 128
        (253, 640),  # both mn and groups padded
        (1, 384),  # extreme: 1 group, 1 mn row -> both axes padded
    ],
)
@pytest.mark.skipif(
    not current_platform.is_cuda_alike(),
    reason="DeepGEMM not available on this platform",
)
def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q(
    num_tokens, hidden_dim
):
    """When output_q is allocated with shape (tma_aligned_mn, k) instead of
    (mn, k), the kernel must overwrite the padded mn rows with zeros so
    callers can use ``torch.empty`` instead of ``torch.zeros``."""

    device = "cuda"
    group_size = 128
    torch.manual_seed(42)
    x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8

    mn = num_tokens
    groups_per_row = hidden_dim // group_size
    k_num_packed = (groups_per_row + 3) // 4
    tma_aligned_mn = ((mn + 3) // 4) * 4

    fp8_dtype = current_platform.fp8_dtype()
    fp8_min, fp8_max = get_fp8_min_max()
    # Allocate output_q with the padded mn extent and pre-fill with 0xFF
    # so the kernel cannot rely on a clean buffer.
    out_q = torch.empty((tma_aligned_mn, hidden_dim), device=device, dtype=fp8_dtype)
    out_q.view(torch.uint8).fill_(0xFF)

    out_s_packed = torch.empty_strided(
        (mn, k_num_packed),
        (1, tma_aligned_mn),
        device=device,
        dtype=torch.int32,
    )

    torch.ops._C.per_token_group_fp8_quant_packed(
        x, out_q, out_s_packed, group_size, 1e-10, fp8_min, fp8_max
    )

    # Live rows must match the Triton reference.
    with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
        ref_q, _ = fp8_utils.per_token_group_quant_fp8(x, group_size, use_ue8m0=True)
    assert torch.equal(out_q[:mn], ref_q), "Live region mismatch"

    # Padded rows must be all-zero; without this, downstream TMA loads would
    # see uninitialised data.
    if tma_aligned_mn > mn:
        padded_bytes = out_q[mn:tma_aligned_mn].view(torch.uint8)
        assert padded_bytes.eq(0).all(), (
            f"Padded rows [{mn}, {tma_aligned_mn}) not zeroed; "
            f"{padded_bytes.ne(0).sum().item()} non-zero bytes"
        )


@pytest.mark.skipif(
    not current_platform.is_cuda_alike(),
    reason="packed FP8 per-token-group quant kernel requires a CUDA-alike GPU",
)
def test_per_token_group_quant_fp8_packed_large_mn():
    """Regression test for https://github.com/vllm-project/vllm/issues/45099.

    Some background: gridDim.x and gridDim.y have different limits of 2^31 - 1 and
    2^16 - 1, respectively.
    Prior code introduced a bug where it incorrectly assumed grid.x and y both have
    2^31 - 1 limits and mixed them up, which doesn't surface until the kernel is
    launched with a large mn that exceeds grid.y limit (2^16 - 1).

    This issue doesn't surface often because each forward pass only processes a
    bounded token batch, not the full context.
    Quantizing tensors with more rows than that will fail at launch with
    "CUDA error: invalid argument".
    This is a differential test that compares fp8 output against Triton output
    reference when token size sits just above the gridDim.y 2^16 - 1 limit.
    """

    device = "cuda"
    group_size = 128
    # hidden 2048 -> 2048/128 = 16 groups per row -> kx=16, ry=1: one grid row per mn
    # row, so any mn > 65535 overflowed grid.y before the fix.
    num_tokens, hidden_dim = 65537, 2048
    torch.manual_seed(42)
    x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8

    out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
        x,
        group_size=group_size,
        use_ue8m0=True,
    )

    with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
        ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(
            x, group_size, use_ue8m0=True
        )

    assert torch.equal(out_q, ref_q), "Quantized output mismatch"

    # Vectorized packed-scale check; the per-element loop used by the smaller
    # tests is too slow at this size. groups_per_row is a multiple of 4 here,
    # so there is no K padding and the packed view lines up.
    mn = num_tokens
    groups_per_row = hidden_dim // group_size
    k_num_packed = (groups_per_row + 3) // 4
    assert groups_per_row % 4 == 0
    ref_exponents = (ref_s.reshape(mn, groups_per_row).view(torch.int32) >> 23) & 0xFF
    exp = ref_exponents.view(mn, k_num_packed, 4)
    expected = (
        exp[..., 0] | (exp[..., 1] << 8) | (exp[..., 2] << 16) | (exp[..., 3] << 24)
    )
    assert torch.equal(out_s_packed.cpu(), expected.cpu()), "Packed scale mismatch"


@pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)])
@pytest.mark.parametrize("group_size", [64, 128])
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_per_token_group_quant_int8(shape, group_size: int):
    device = "cuda"

    torch.manual_seed(42)
    num_tokens, hidden_dim = shape

    x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8

    # cuda path
    out_q, scale = int8_utils.per_token_group_quant_int8(
        x,
        group_size,
    )

    # triton ref
    with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
        ref_q, ref_s = int8_utils.per_token_group_quant_int8(
            x,
            group_size,
        )

    assert torch.allclose(out_q.float(), ref_q.float(), atol=0.15, rtol=0.15)
    assert torch.allclose(scale, ref_s, atol=0.01, rtol=0.01)