# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for Marlin thread-tile padding of TP-sharded weight shapes.

Run `pytest tests/kernels/quantization/test_marlin_tile_padding.py`.
"""

import pytest
import torch

from vllm import _custom_ops as ops
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
    GPTQ_MARLIN_TILE,
    apply_gptq_marlin_linear,
    marlin_make_empty_g_idx,
    marlin_make_workspace_new,
    marlin_pad_qweight,
    marlin_pad_scales,
    marlin_padded_nk,
    marlin_permute_scales,
    marlin_repacked_nk,
    marlin_zero_points,
)
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
    apply_fp4_marlin_linear,
    is_fp4_marlin_supported,
    prepare_fp4_layer_for_marlin,
)
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
    apply_fp8_marlin_linear,
    apply_mxfp8_marlin_linear,
    is_fp8_marlin_supported,
    prepare_fp8_layer_for_marlin,
    prepare_mxfp8_layer_for_marlin,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
    gptq_pack,
    gptq_quantize_weights,
    quantize_weights,
)
from vllm.platforms import current_platform
from vllm.scalar_type import scalar_types

# (size_n, size_k) rank-local shapes that violate Marlin tile alignment,
# e.g. produced by TP-sharding dims that are valid at TP=1.
ODD_SHAPES = [
    (200, 288),  # N padded
    (256, 208),  # K padded
    (200, 208),  # both padded
    (4640, 512),  # Nemotron-Super-120B q_proj shard at TP=4
]
ALIGNED_SHAPES = [(64, 128), (128, 64), (256, 256), (4608, 4096)]


def _is_tile_aligned(size_n: int, size_k: int) -> bool:
    return (size_n % 64 == 0 and size_k % 128 == 0) or (
        size_n % 128 == 0 and size_k % 64 == 0
    )


@pytest.mark.parametrize("shape", ODD_SHAPES + ALIGNED_SHAPES)
@pytest.mark.parametrize("group_size", [-1, 16, 32, 64, 128])
def test_marlin_padded_nk(shape, group_size):
    size_n, size_k = shape
    padded_n, padded_k = marlin_padded_nk(size_n, size_k, group_size)

    assert padded_n >= size_n and padded_k >= size_k
    assert _is_tile_aligned(padded_n, padded_k)
    if group_size > 0:
        assert padded_k % group_size == 0

    # Aligned shapes must pass through unchanged (zero hot-path cost).
    if _is_tile_aligned(size_n, size_k) and (
        group_size <= 0 or size_k % group_size == 0
    ):
        assert (padded_n, padded_k) == (size_n, size_k)

    # Minimal: no valid shape with a smaller padded area exists.
    area = padded_n * padded_k
    for cand_n in range(size_n, padded_n + 1):
        for cand_k in range(size_k, padded_k + 1):
            if (
                _is_tile_aligned(cand_n, cand_k)
                and (group_size <= 0 or cand_k % group_size == 0)
                and cand_n * cand_k < area
            ):
                pytest.fail(f"({cand_n}, {cand_k}) beats ({padded_n}, {padded_k})")

    # Apply-time derivation from the repacked-tensor shape must round-trip.
    for num_bits in (4, 8):
        pack_factor = 32 // num_bits
        repacked_shape = (
            padded_k // GPTQ_MARLIN_TILE,
            padded_n * GPTQ_MARLIN_TILE // pack_factor,
        )
        repacked = torch.empty(repacked_shape, device="meta")
        assert marlin_repacked_nk(repacked, num_bits) == (padded_n, padded_k)


def test_marlin_pad_helpers_shapes():
    size_n, size_k, group_size = 200, 208, 16
    padded_n, padded_k = marlin_padded_nk(size_n, size_k, group_size)

    qweight = torch.zeros(size_k // 8, size_n, dtype=torch.int32)
    padded = marlin_pad_qweight(qweight, size_n, size_k, padded_n, padded_k)
    assert padded.shape == (padded_k // 8, padded_n)

    scales = torch.ones(size_k // group_size, size_n)
    padded = marlin_pad_scales(scales, size_n, size_k, padded_n, padded_k, group_size)
    assert padded.shape == (padded_k // group_size, padded_n)
    assert padded[:, size_n:].abs().sum() == 0

    channelwise = torch.ones(1, size_n)
    padded = marlin_pad_scales(channelwise, size_n, size_k, padded_n, padded_k, -1)
    assert padded.shape == (1, padded_n)


def _gpu_marlin_unsupported() -> bool:
    return not (
        current_platform.is_cuda() and current_platform.has_device_capability(80)
    )


@pytest.mark.skipif(
    _gpu_marlin_unsupported() or not is_fp8_marlin_supported(),
    reason="FP8 Marlin is not supported on this GPU type.",
)
@pytest.mark.parametrize("shape", ODD_SHAPES)
@pytest.mark.parametrize("use_bias", [False, True])
def test_fp8_marlin_padded_round_trip(shape, use_bias):
    size_n, size_k = shape
    dtype = torch.float16
    layer = torch.nn.Module()
    layer.output_size_per_partition = size_n
    layer.input_size_per_partition = size_k
    layer.orig_dtype = dtype

    weight = torch.randn(size_k, size_n, dtype=dtype, device="cuda") / size_k**0.5
    scale = weight.abs().max() / 448
    weight_fp8 = (weight / scale).to(torch.float8_e4m3fn)
    layer.weight = torch.nn.Parameter(weight_fp8, requires_grad=False)
    layer.weight_scale = torch.nn.Parameter(
        scale.to(torch.float32), requires_grad=False
    )
    bias = None
    if use_bias:
        bias = torch.randn(size_n, dtype=dtype, device="cuda")
        layer.bias = torch.nn.Parameter(bias.clone(), requires_grad=False)

    prepare_fp8_layer_for_marlin(layer, size_k_first=True)

    x = torch.randn(8, size_k, dtype=dtype, device="cuda")
    output = apply_fp8_marlin_linear(
        input=x,
        weight=layer.weight,
        weight_scale=layer.weight_scale,
        workspace=layer.workspace,
        size_n=size_n,
        size_k=size_k,
        bias=layer.bias if use_bias else None,
    )
    ref = x @ (weight_fp8.to(dtype) * scale.to(dtype))
    if use_bias:
        ref = ref + bias

    assert output.shape == (8, size_n)
    torch.testing.assert_close(output, ref, rtol=2e-2, atol=2e-2)


def _dequant_fp4(packed: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
    """Dequantize packed e2m1 nibbles (N, K // 2) -> (N, K) in dtype."""
    lo = (packed & 0b10000000) | ((packed & 0b01110000) >> 2)
    lo = lo.view(torch.float8_e4m3fn).to(dtype) * (2**6)
    hi_bits = packed << 4
    hi = (hi_bits & 0b10000000) | ((hi_bits & 0b01110000) >> 2)
    hi = hi.view(torch.float8_e4m3fn).to(dtype) * (2**6)
    return torch.cat([hi.unsqueeze(2), lo.unsqueeze(2)], 2).view(packed.size(0), -1)


@pytest.mark.skipif(
    _gpu_marlin_unsupported() or not is_fp4_marlin_supported(),
    reason="FP4 Marlin is not supported on this GPU type.",
)
@pytest.mark.parametrize("shape", ODD_SHAPES)
def test_nvfp4_marlin_padded_round_trip(shape):
    size_n, size_k = shape
    group_size = 16
    dtype = torch.float16
    layer = torch.nn.Module()
    layer.output_size_per_partition = size_n
    layer.input_size_per_partition = size_k
    layer.params_dtype = dtype

    packed = torch.randint(
        0, 256, (size_n, size_k // 2), dtype=torch.uint8, device="cuda"
    )
    scales = (torch.rand(size_n, size_k // group_size, device="cuda") + 0.25).to(
        torch.float8_e4m3fn
    )
    global_scale = torch.tensor([0.002], dtype=torch.float32, device="cuda")

    ref_weight = (
        _dequant_fp4(packed, dtype)
        * scales.to(dtype).repeat_interleave(group_size, 1)
        * global_scale.to(dtype)
    )

    layer.weight = torch.nn.Parameter(packed, requires_grad=False)
    layer.weight_scale = torch.nn.Parameter(scales, requires_grad=False)
    layer.weight_global_scale = torch.nn.Parameter(global_scale, requires_grad=False)

    prepare_fp4_layer_for_marlin(layer)

    x = torch.randn(8, size_k, dtype=dtype, device="cuda") / size_k**0.5
    output = apply_fp4_marlin_linear(
        input=x,
        weight=layer.weight,
        weight_scale=layer.weight_scale,
        weight_global_scale=layer.weight_global_scale,
        workspace=layer.workspace,
        size_n=size_n,
        size_k=size_k,
    )
    ref = x @ ref_weight.T

    assert output.shape == (8, size_n)
    torch.testing.assert_close(output, ref, rtol=2e-2, atol=2e-2)


@pytest.mark.skipif(
    _gpu_marlin_unsupported(),
    reason="Marlin is not supported on this GPU type.",
)
@pytest.mark.parametrize("shape", ODD_SHAPES)
@pytest.mark.parametrize("group_size", [-1, 128])
def test_gptq_marlin_padded_round_trip(shape, group_size):
    """Pad-then-repack a GPTQ int4 weight the way MarlinLinearKernel does and
    check the GEMM against the dequantized reference.

    Symmetric int4's quantized zero decodes to -8, so this exercises the
    zero-padded-scales cancellation, not just zero weights.
    """
    size_n, size_k = shape
    if group_size > 0 and size_k % group_size != 0:
        pytest.skip("group must divide the rank-local K (not fixable by padding)")
    dtype = torch.float16
    quant_type = scalar_types.uint4b8
    device = torch.device("cuda")

    weight = torch.randn(size_k, size_n, dtype=dtype, device=device) / size_k**0.5
    w_ref, q_w, s, _, _ = gptq_quantize_weights(
        weight, quant_type, group_size, act_order=False
    )
    qweight = gptq_pack(q_w, quant_type.size_bits, size_k, size_n)

    padded_n, padded_k = marlin_padded_nk(size_n, size_k, group_size)
    qweight = marlin_pad_qweight(qweight, size_n, size_k, padded_n, padded_k)
    marlin_qweight = ops.gptq_marlin_repack(
        b_q_weight=qweight,
        perm=torch.empty(0, dtype=torch.int, device=device),
        size_k=padded_k,
        size_n=padded_n,
        num_bits=quant_type.size_bits,
    )
    s = marlin_pad_scales(s, size_n, size_k, padded_n, padded_k, group_size)
    marlin_s = marlin_permute_scales(
        s, size_k=padded_k, size_n=padded_n, group_size=group_size
    )

    x = torch.randn(8, size_k, dtype=dtype, device=device)
    output = apply_gptq_marlin_linear(
        input=x,
        weight=marlin_qweight,
        weight_scale=marlin_s,
        weight_zp=marlin_make_empty_g_idx(device),
        g_idx=marlin_make_empty_g_idx(device),
        g_idx_sort_indices=marlin_make_empty_g_idx(device),
        workspace=marlin_make_workspace_new(device),
        wtype=quant_type,
        output_size_per_partition=size_n,
        input_size_per_partition=size_k,
        is_k_full=True,
    )
    ref = x @ w_ref

    assert output.shape == (8, size_n)
    torch.testing.assert_close(output, ref, rtol=2e-2, atol=2e-2)


@pytest.mark.skipif(
    _gpu_marlin_unsupported() or not is_fp8_marlin_supported(),
    reason="FP8 Marlin is not supported on this GPU type.",
)
@pytest.mark.parametrize("shape", [(200, 512), (4640, 512)])
def test_fp8_block_marlin_padded_round_trip(shape):
    """Block-quantized FP8 (e.g. Nemotron NVFP4 checkpoints' FP8 layers):
    group_size=128 exercises the lcm K-alignment in marlin_padded_nk and the
    weight_scale_inv group-wise scale padding."""
    size_n, size_k = shape
    block = 128
    dtype = torch.float16
    layer = torch.nn.Module()
    layer.output_size_per_partition = size_n
    layer.input_size_per_partition = size_k
    layer.orig_dtype = dtype
    layer.weight_block_size = [block, block]

    weight = torch.randn(size_n, size_k, dtype=dtype, device="cuda") / size_k**0.5
    n_blocks, k_blocks = (size_n + block - 1) // block, size_k // block
    padded = torch.zeros(n_blocks * block, size_k, dtype=dtype, device="cuda")
    padded[:size_n] = weight
    scales = padded.view(n_blocks, block, k_blocks, block).abs().amax(dim=(1, 3)) / 448
    scales_expanded = scales.repeat_interleave(block, 0)[:size_n].repeat_interleave(
        block, 1
    )
    weight_fp8 = (weight / scales_expanded).to(torch.float8_e4m3fn)

    layer.weight = torch.nn.Parameter(weight_fp8, requires_grad=False)
    layer.weight_scale_inv = torch.nn.Parameter(
        scales.to(torch.float32), requires_grad=False
    )

    prepare_fp8_layer_for_marlin(layer, size_k_first=False)

    x = torch.randn(8, size_k, dtype=dtype, device="cuda")
    output = apply_fp8_marlin_linear(
        input=x,
        weight=layer.weight,
        weight_scale=layer.weight_scale_inv,
        workspace=layer.workspace,
        size_n=size_n,
        size_k=size_k,
        bias=None,
    )
    ref = x @ (weight_fp8.to(dtype) * scales_expanded.to(dtype)).T

    assert output.shape == (8, size_n)
    torch.testing.assert_close(output, ref, rtol=2e-2, atol=2e-2)


@pytest.mark.skipif(
    _gpu_marlin_unsupported() or not is_fp8_marlin_supported(),
    reason="FP8 Marlin is not supported on this GPU type.",
)
@pytest.mark.parametrize("shape", [(200, 288), (4640, 512)])
def test_mxfp8_marlin_padded_round_trip(shape):
    """MXFP8 exercises the e8m0 scale path, where padded 0.0 scales clamp to
    2^-127 instead of zero and must still contribute nothing."""
    size_n, size_k = shape
    group_size = 32
    # The e8m0-scale Marlin kernels are only instantiated for bf16 activations.
    dtype = torch.bfloat16
    layer = torch.nn.Module()
    layer.output_size_per_partition = size_n
    layer.input_size_per_partition = size_k

    weight_fp8 = (torch.randn(size_n, size_k, dtype=dtype, device="cuda") / 4).to(
        torch.float8_e4m3fn
    )
    # e8m0 exponents around 1.0 (127): scales in [2^-6, 2^0]
    scales = torch.randint(
        121, 128, (size_n, size_k // group_size), dtype=torch.uint8, device="cuda"
    )
    ref_weight = weight_fp8.to(dtype) * (
        2.0 ** (scales.to(dtype) - 127)
    ).repeat_interleave(group_size, 1)

    layer.weight = torch.nn.Parameter(weight_fp8, requires_grad=False)
    layer.weight_scale = torch.nn.Parameter(scales, requires_grad=False)

    prepare_mxfp8_layer_for_marlin(layer)

    x = torch.randn(8, size_k, dtype=dtype, device="cuda") / size_k**0.5
    output = apply_mxfp8_marlin_linear(
        input=x,
        weight=layer.weight,
        weight_scale=layer.weight_scale,
        workspace=layer.workspace,
        size_n=size_n,
        size_k=size_k,
    )
    ref = x @ ref_weight.T

    assert output.shape == (8, size_n)
    torch.testing.assert_close(output, ref, rtol=2e-2, atol=2e-2)


@pytest.mark.skipif(
    _gpu_marlin_unsupported(),
    reason="Marlin is not supported on this GPU type.",
)
@pytest.mark.parametrize("shape", [(200, 512), (4640, 512)])
def test_awq_zp_marlin_padded_round_trip(shape):
    """AWQ-style uint4 with runtime zero-points, padded the way
    MarlinLinearKernel does: padded columns rely on (q=0 - zp=0) * scale=0."""
    size_n, size_k = shape
    group_size = 128
    dtype = torch.float16
    quant_type = scalar_types.uint4
    device = torch.device("cuda")

    weight = torch.randn(size_k, size_n, dtype=dtype, device=device) / size_k**0.5
    w_ref, q_w, s, zp = quantize_weights(
        weight, quant_type, group_size, zero_points=True
    )
    qweight = gptq_pack(q_w, quant_type.size_bits, size_k, size_n)

    padded_n, padded_k = marlin_padded_nk(size_n, size_k, group_size)
    qweight = marlin_pad_qweight(qweight, size_n, size_k, padded_n, padded_k)
    marlin_qweight = ops.gptq_marlin_repack(
        b_q_weight=qweight,
        perm=torch.empty(0, dtype=torch.int, device=device),
        size_k=padded_k,
        size_n=padded_n,
        num_bits=quant_type.size_bits,
    )
    s = marlin_pad_scales(s, size_n, size_k, padded_n, padded_k, group_size)
    marlin_s = marlin_permute_scales(
        s, size_k=padded_k, size_n=padded_n, group_size=group_size
    )
    zp = marlin_pad_scales(zp, size_n, size_k, padded_n, padded_k, group_size)
    marlin_zp = marlin_zero_points(
        zp,
        size_k=padded_k // group_size,
        size_n=padded_n,
        num_bits=quant_type.size_bits,
    )

    x = torch.randn(8, size_k, dtype=dtype, device=device)
    output = apply_gptq_marlin_linear(
        input=x,
        weight=marlin_qweight,
        weight_scale=marlin_s,
        weight_zp=marlin_zp,
        g_idx=marlin_make_empty_g_idx(device),
        g_idx_sort_indices=marlin_make_empty_g_idx(device),
        workspace=marlin_make_workspace_new(device),
        wtype=quant_type,
        output_size_per_partition=size_n,
        input_size_per_partition=size_k,
        is_k_full=True,
    )
    ref = x @ w_ref

    assert output.shape == (8, size_n)
    torch.testing.assert_close(output, ref, rtol=2e-2, atol=2e-2)


class _FakeLinear:
    def __init__(self, size_n, size_k, input_size=None):
        self.output_size_per_partition = size_n
        self.input_size_per_partition = size_k
        self.output_size = size_n
        self.input_size = input_size if input_size is not None else size_k


def test_check_marlin_supports_layer_allow_tile_padding():
    from vllm.model_executor.layers.quantization.utils.marlin_utils import (
        check_marlin_supports_layer,
    )

    # Tile-misaligned but group-aligned: rejected strictly, allowed w/ padding
    layer = _FakeLinear(4640, 512, input_size=2048)
    assert not check_marlin_supports_layer(layer, 128)
    assert check_marlin_supports_layer(layer, 128, allow_tile_padding=True)
    assert check_marlin_supports_layer(layer, -1, allow_tile_padding=True)

    # A group straddling the TP shard cannot be fixed by padding
    layer = _FakeLinear(4608, 4672, input_size=18688)
    assert not check_marlin_supports_layer(layer, 128, allow_tile_padding=True)