# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for CPU quantized fused MoE kernels."""

import math
import sys

import pytest
import torch
import torch.nn.functional as F

from vllm.platforms import current_platform
from vllm.utils.torch_utils import set_random_seed

if not current_platform.is_cpu():
    pytest.skip("skipping CPU-only tests", allow_module_level=True)

import vllm._custom_ops as ops  # noqa: E402

if not hasattr(torch.ops._C, "fused_experts_cpu"):
    pytest.skip("fused_experts_cpu op not available", allow_module_level=True)


def _silu_and_mul(x: torch.Tensor) -> torch.Tensor:
    d = x.shape[-1] // 2
    return F.silu(x[..., :d]) * x[..., d:]


def _prepack_experts(w: torch.Tensor) -> torch.Tensor:
    """VNNI-prepack expert weights via ``convert_weight_packed``."""
    return torch.ops._C.convert_weight_packed(w)


# ===========================================================================
# FP8 W8A16 MoE
# ===========================================================================


BLOCK_SIZE = [128, 128]  # [block_n, block_k]

_FP8_INFO = torch.finfo(torch.float8_e4m3fn)
FP8_SCALE = _FP8_INFO.max  # 448.0
FACTOR_FOR_SCALE = 1e-3


def _block_dequant_weight(
    weight: torch.Tensor,
    scales: torch.Tensor,
    block_size: list[int],
) -> torch.Tensor:
    """Block-dequantize FP8 weight [E, N, K] -> float [E, N, K]."""
    E, N, K = weight.shape
    block_n, block_k = block_size
    pad_N = (block_n - N % block_n) % block_n
    pad_K = (block_k - K % block_k) % block_k

    if pad_N > 0 or pad_K > 0:
        weight = F.pad(weight, (0, pad_K, 0, pad_N))

    n_tiles = math.ceil(N / block_n)
    k_tiles = math.ceil(K / block_k)

    weight_block = (
        weight.view(E, n_tiles, block_n, k_tiles, block_k)
        .permute(0, 1, 3, 2, 4)
        .float()
        .contiguous()
    )
    weight_scaled = (
        (weight_block * scales.view(E, n_tiles, k_tiles, 1, 1))
        .permute(0, 1, 3, 2, 4)
        .contiguous()
    )
    if pad_N > 0 or pad_K > 0:
        weight_scaled = weight_scaled.view(E, N + pad_N, K + pad_K)
        weight_scaled = weight_scaled[..., :N, :K].contiguous()
    else:
        weight_scaled = weight_scaled.view(E, N, K)
    return weight_scaled


def ref_w8a16_block_fp8_moe(
    a: torch.Tensor,
    w1: torch.Tensor,
    w2: torch.Tensor,
    w1_s: torch.Tensor,
    w2_s: torch.Tensor,
    topk_weight: torch.Tensor,
    topk_ids: torch.Tensor,
    block_size: list[int],
) -> torch.Tensor:
    """Reference FP8 W8A16 block-scaled fused MoE in pure torch."""
    B, D = a.shape
    topk = topk_ids.size(1)

    w1_dq = _block_dequant_weight(w1, w1_s, block_size)
    w2_dq = _block_dequant_weight(w2, w2_s, block_size)

    a_exp = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D).float()
    out = torch.zeros(B * topk, w2_dq.shape[1], dtype=torch.float32)

    topk_weight_flat = topk_weight.view(-1)
    topk_ids_flat = topk_ids.view(-1)

    for i in range(w1_dq.shape[0]):
        mask = topk_ids_flat == i
        if mask.sum():
            ic0 = torch.matmul(a_exp[mask], w1_dq[i].transpose(0, 1))
            ic1 = _silu_and_mul(ic0)
            out[mask] = torch.matmul(ic1, w2_dq[i].transpose(0, 1))

    return (
        (out.view(B, -1, w2_dq.shape[1]) * topk_weight_flat.view(B, -1, 1))
        .sum(dim=1)
        .to(a.dtype)
    )


def _make_fp8_moe_weights(
    E: int,
    N: int,
    K: int,
    block_size: list[int],
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Generate random FP8 MoE weights with random block scales."""
    block_n, block_k = block_size

    w1 = (
        (torch.randn(E, 2 * N, K) * FP8_SCALE)
        .clamp(min=-FP8_SCALE, max=FP8_SCALE)
        .to(torch.float8_e4m3fn)
    )
    w2 = (
        (torch.randn(E, K, N) * FP8_SCALE)
        .clamp(min=-FP8_SCALE, max=FP8_SCALE)
        .to(torch.float8_e4m3fn)
    )

    w1_s = (
        torch.randn(E, math.ceil(2 * N / block_n), math.ceil(K / block_k))
        * FACTOR_FOR_SCALE
    )
    w2_s = (
        torch.randn(E, math.ceil(K / block_n), math.ceil(N / block_k))
        * FACTOR_FOR_SCALE
    )
    return w1, w2, w1_s, w2_s


FP8_NUM_TOKENS = [1, 2, 64, 121]
FP8_MOE_CONFIGS = [
    (256, 512, 8, 2),
    (256, 512, 8, 4),
    (512, 256, 8, 2),
    (512, 256, 8, 4),
    (512, 512, 8, 2),
    (512, 512, 8, 4),
    (768, 2048, 8, 2),
    (768, 2048, 8, 4),
    (768, 2048, 128, 8),
]


@pytest.mark.parametrize("M", FP8_NUM_TOKENS)
@pytest.mark.parametrize("N,K,E,topk", FP8_MOE_CONFIGS)
@pytest.mark.parametrize("seed", [0])
def test_w8a16_block_fp8_cpu_fused_moe(M, N, K, E, topk, seed):
    """Test fused_experts_cpu FP8 W8A16 against dequantised torch reference."""
    set_random_seed(seed)

    a = torch.randn(M, K, dtype=torch.bfloat16) / math.sqrt(K)
    w1, w2, w1_s, w2_s = _make_fp8_moe_weights(E, N, K, BLOCK_SIZE)

    score = torch.randn(M, E, dtype=torch.bfloat16)
    score = torch.softmax(score, dim=-1, dtype=torch.float32)
    topk_weight, topk_ids = torch.topk(score, topk)
    topk_ids = topk_ids.to(torch.int32)

    ref_out = ref_w8a16_block_fp8_moe(
        a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, BLOCK_SIZE
    )

    pw1, pw2 = _prepack_experts(w1), _prepack_experts(w2)

    # Test inplace=False against reference
    out = ops.fused_experts_cpu(
        a.clone(),
        pw1,
        pw2,
        topk_weight,
        topk_ids,
        False,
        ops.CPUQuantMethod.FP8_W8A16,
        w1_s,
        w2_s,
        None,
        None,
        BLOCK_SIZE,
        is_vnni=True,
    )
    torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)

    # Test inplace=True produces identical output
    out_inplace = ops.fused_experts_cpu(
        a.clone(),
        pw1,
        pw2,
        topk_weight,
        topk_ids,
        True,
        ops.CPUQuantMethod.FP8_W8A16,
        w1_s,
        w2_s,
        None,
        None,
        BLOCK_SIZE,
        is_vnni=True,
    )
    torch.testing.assert_close(out_inplace, out, atol=0, rtol=0)


# ===========================================================================
# MXFP4 W4A16 MoE
# ===========================================================================


class MXFP4QuantizeUtil:
    """MXFP4 quantization utility."""

    E2M1_max = 6.0
    E2M1_values = [0, 0.5, 1, 1.5, 2, 3, 4, 6]
    E2M1_bounds = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5])
    block_size = 32

    @classmethod
    def quantize(cls, input: torch.Tensor) -> tuple:
        """Quantize BF16 tensor to MXFP4 packed uint8 format."""

        def cast_fp4(x):
            sign = torch.sign(x)
            sign_bit = (2 - sign) // 2
            ord_ = torch.sum(
                (x.abs().unsqueeze(-1) - cls.E2M1_bounds.to(x.device)) > 0, dim=-1
            )
            fp4_val = (sign_bit * 0b1000 + ord_).to(torch.uint8)
            return fp4_val

        def fuse_uint4_to_uint8(x):
            left_side = x[..., 0::2]
            right_side = x[..., 1::2]
            new_data = right_side.clone() << 4
            new_data[..., : left_side.shape[-1]] += left_side
            return new_data

        original_shape = input.shape
        input = input.view(-1, cls.block_size)
        input_amax = input.abs().max(dim=-1, keepdim=True).values
        descale = input_amax / cls.E2M1_max
        min_value = torch.tensor(-127.0, device=descale.device)
        e8m0_scale = torch.ceil(torch.maximum(torch.log2(descale), min_value))

        input = (input / torch.exp2(e8m0_scale)).view(original_shape)
        input_q = cast_fp4(input)
        input_q = fuse_uint4_to_uint8(input_q)
        e8m0_scale = (e8m0_scale + 127).to(torch.uint8)
        return input_q, e8m0_scale

    @classmethod
    def dequantize(cls, quantized_data, dtype: torch.dtype, scale):
        """Dequantize MXFP4 packed tensor back to float."""

        def unfuse_uint8_to_uint4(x):
            left_side = x & 0x0F
            right_side = (x >> 4) & 0x0F
            shape = list(x.shape)
            shape[-1] = shape[-1] * 2
            result = torch.zeros(shape, dtype=torch.uint8, device=x.device)
            result[..., 0::2] = left_side
            result[..., 1::2] = right_side
            return result

        e8m0_scale = scale
        x_unfused = unfuse_uint8_to_uint4(quantized_data)
        sign = 1 - 2 * ((x_unfused & 0b1000) >> 3).to(torch.float32)
        magnitude = (x_unfused & 0b0111).to(torch.long)
        values = torch.tensor(cls.E2M1_values, device=quantized_data.device)
        original_shape = magnitude.shape
        x_float = values[magnitude.reshape(-1)].reshape(original_shape)
        x_float = sign.float() * x_float
        x_float = x_float.reshape(-1, cls.block_size)
        scale_factor = torch.exp2(e8m0_scale.float() - 127)
        scale_factor = scale_factor.reshape(-1, 1)
        x_float = x_float * scale_factor
        return x_float.reshape(original_shape).to(dtype)


def _swiglu(x: torch.Tensor, alpha: float, limit: float) -> torch.Tensor:
    """SwigLU activation used in GPT-OSS.
    Input is interleaved: [gate_0, up_0, gate_1, up_1, ...] in last dim.
    """
    gate = x[..., 0::2]
    up = x[..., 1::2]
    gate_clamped = torch.clamp(gate, max=limit)
    up_clamped = torch.clamp(up, min=-limit, max=limit)
    return gate_clamped * torch.sigmoid(alpha * gate_clamped) * (up_clamped + 1)


def ref_mxfp4_fused_moe(
    a: torch.Tensor,
    w1_dq: torch.Tensor,
    w2_dq: torch.Tensor,
    topk_weight: torch.Tensor,
    topk_ids: torch.Tensor,
    topk: int,
) -> torch.Tensor:
    """Reference MXFP4 fused MoE with SiLU activation."""
    B, D = a.shape
    a_f = a.float()
    out = torch.zeros(B * topk, w2_dq.shape[1], dtype=torch.float32)
    topk_ids_flat = topk_ids.view(-1)

    for i in range(w1_dq.shape[0]):
        mask = topk_ids_flat == i
        if mask.sum() == 0:
            continue
        token_indices = torch.where(mask)[0]
        source_indices = token_indices // topk
        ic0 = torch.matmul(a_f[source_indices], w1_dq[i].float().T)
        ic1 = _silu_and_mul(ic0)
        out[mask] = torch.matmul(ic1, w2_dq[i].float().T)

    return (out.view(B, topk, -1) * topk_weight.unsqueeze(-1)).sum(dim=1).to(a.dtype)


def ref_mxfp4_fused_moe_gptoss(
    a: torch.Tensor,
    w1_dq: torch.Tensor,
    w2_dq: torch.Tensor,
    w1_bias: torch.Tensor,
    w2_bias: torch.Tensor,
    topk_weight: torch.Tensor,
    topk_ids: torch.Tensor,
    alpha: float,
    limit: float,
) -> torch.Tensor:
    """Reference MXFP4 fused MoE with SwigLU+bias (GPT-OSS style)."""
    B, D = a.shape
    topk = topk_ids.shape[1]
    a_f = a.float()
    E = w1_dq.shape[0]
    out = torch.zeros(B * topk, w2_dq.shape[1], dtype=torch.float32)
    topk_ids_flat = topk_ids.view(-1)

    for i in range(E):
        mask = topk_ids_flat == i
        if mask.sum() == 0:
            continue
        token_indices = torch.where(mask)[0]
        source_indices = token_indices // topk
        ic0 = torch.matmul(a_f[source_indices], w1_dq[i].float().T)
        ic0 = ic0 + w1_bias[i].float()
        ic1 = _swiglu(ic0, alpha, limit)
        ic2 = torch.matmul(ic1, w2_dq[i].float().T)
        ic2 = ic2 + w2_bias[i].float()
        out[mask] = ic2

    return (out.view(B, topk, -1) * topk_weight.unsqueeze(-1)).sum(dim=1).to(a.dtype)


def _prepack_mxfp4_experts(
    w: torch.Tensor, w_scale: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """VNNI-prepack MXFP4 weights and repack scales."""
    packed_w = torch.ops._C.convert_weight_packed(w)
    packed_s = torch.ops._C.convert_scale_packed(w_scale)
    return packed_w, packed_s


MXFP4_NUM_TOKENS = [1, 2, 32, 121]
MXFP4_MOE_CONFIGS = [
    (128, 128, 4, 2),
    (256, 256, 8, 4),
    (352, 256, 8, 4),
    (512, 320, 8, 4),
]


@pytest.mark.parametrize("M", MXFP4_NUM_TOKENS)
@pytest.mark.parametrize("N,K,E,topk", MXFP4_MOE_CONFIGS)
@pytest.mark.parametrize("seed", [0])
def test_mxfp4_cpu_fused_moe(M, N, K, E, topk, seed):
    """Test fused_experts_mxfp4_cpu against dequantized torch reference."""
    set_random_seed(seed)
    dtype = torch.bfloat16

    a = torch.randn(M, K, dtype=dtype) / 10

    # Generate and quantize weights
    w1_bf16 = torch.randn(E, 2 * N, K, dtype=dtype) / 10
    w1q, w1s = MXFP4QuantizeUtil.quantize(w1_bf16)
    w1s = w1s.reshape(E, 2 * N, K // 32)
    w1dq = MXFP4QuantizeUtil.dequantize(w1q, dtype, w1s)

    w2_bf16 = torch.randn(E, K, N, dtype=dtype) / 10
    w2q, w2s = MXFP4QuantizeUtil.quantize(w2_bf16)
    w2s = w2s.reshape(E, K, N // 32)
    w2dq = MXFP4QuantizeUtil.dequantize(w2q, dtype, w2s)

    # Routing
    score = torch.randn(M, E, dtype=dtype)
    score = torch.softmax(score, dim=-1, dtype=torch.float32)
    topk_weight, topk_ids = torch.topk(score, topk)
    topk_ids = topk_ids.to(torch.int32)

    # Reference
    ref_out = ref_mxfp4_fused_moe(a, w1dq, w2dq, topk_weight, topk_ids, topk)

    # Pack weights for kernel
    pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s)
    pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s)

    # Kernel
    out = ops.fused_experts_cpu(
        a.clone(),
        pw1,
        pw2,
        topk_weight,
        topk_ids,
        False,  # inplace
        ops.CPUQuantMethod.MXFP4,
        pw1s,  # w1_scale
        pw2s,  # w2_scale
        None,  # w1_zero
        None,  # w2_zero
        None,  # block_size
    )

    torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)


@pytest.mark.parametrize("M", [1, 32])
@pytest.mark.parametrize("N,K,E,topk", [(128, 128, 4, 2), (64, 64, 4, 2)])
@pytest.mark.parametrize("seed", [0])
def test_mxfp4_cpu_fused_moe_bias_swiglu(M, N, K, E, topk, seed):
    """Test fused_experts_mxfp4_cpu with bias and SwigLU activation (GPT-OSS)."""
    set_random_seed(seed)
    dtype = torch.bfloat16
    alpha = 1.702
    limit = 7.0

    a = torch.randn(M, K, dtype=dtype) / 10

    # Generate and quantize weights
    w1_bf16 = torch.randn(E, 2 * N, K, dtype=dtype) / 10
    w1q, w1s = MXFP4QuantizeUtil.quantize(w1_bf16)
    w1s = w1s.reshape(E, 2 * N, K // 32)
    w1dq = MXFP4QuantizeUtil.dequantize(w1q, dtype, w1s)
    w1_b = torch.randn(E, 2 * N, dtype=torch.float32) / 10

    w2_bf16 = torch.randn(E, K, N, dtype=dtype) / 10
    w2q, w2s = MXFP4QuantizeUtil.quantize(w2_bf16)
    w2s = w2s.reshape(E, K, N // 32)
    w2dq = MXFP4QuantizeUtil.dequantize(w2q, dtype, w2s)
    w2_b = torch.randn(E, K, dtype=torch.float32) / 10

    # Routing
    score = torch.randn(M, E, dtype=dtype)
    score = torch.softmax(score, dim=-1, dtype=torch.float32)
    topk_weight, topk_ids = torch.topk(score, topk)
    topk_ids = topk_ids.to(torch.int32)

    # Reference
    ref_out = ref_mxfp4_fused_moe_gptoss(
        a, w1dq, w2dq, w1_b, w2_b, topk_weight, topk_ids, alpha, limit
    )

    # Pack weights for kernel
    pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s)
    pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s)

    # Kernel
    out = ops.fused_experts_cpu(
        a.clone(),
        pw1,
        pw2,
        topk_weight,
        topk_ids,
        False,  # inplace
        ops.CPUQuantMethod.MXFP4,
        pw1s,  # w1_scale
        pw2s,  # w2_scale
        None,  # w1_zero
        None,  # w2_zero
        None,  # block_size
        w1_bias=w1_b,
        w2_bias=w2_b,
        alpha=alpha,
        limit=limit,
    )

    torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)


# ===========================================================================
# INT4 W4A16 MoE
# ===========================================================================


def _pack_int4_gptq(w_int4: torch.Tensor) -> torch.Tensor:
    """Pack INT4 values [N, K] → [N, K//8] int32 along K dim (GPTQ format)."""
    N, K = w_int4.shape
    assert K % 8 == 0
    w = w_int4.to(torch.int32)
    w_packed = torch.zeros(N, K // 8, dtype=torch.int32)
    for j in range(8):
        w_packed |= (w[:, j::8] & 0xF) << (j * 4)
    return w_packed


def _pack_int4_awq(w_int4: torch.Tensor) -> torch.Tensor:
    """Pack INT4 values [..., N] → [..., N//8] int32 along last dim (AWQ format)."""
    # AWQ packing bitshifts: indices {0,4,1,5,2,6,3,7} * 4 bits each
    _AWQ_BITSHIFTS = [0, 16, 4, 20, 8, 24, 12, 28]

    N = w_int4.shape[-1]
    assert N % 8 == 0
    w = w_int4.to(torch.int32)
    w_packed = torch.zeros(*w.shape[:-1], N // 8, dtype=torch.int32)
    for j, shift in enumerate(_AWQ_BITSHIFTS):
        w_packed |= (w[..., j::8] & 0xF) << shift
    return w_packed


def _ref_int4_moe(
    a: torch.Tensor,
    w1_int4: torch.Tensor,
    w2_int4: torch.Tensor,
    w1_zeros: torch.Tensor | None,
    w2_zeros: torch.Tensor | None,
    w1_s: torch.Tensor,
    w2_s: torch.Tensor,
    topk_weight: torch.Tensor,
    topk_ids: torch.Tensor,
    group_size: int,
) -> torch.Tensor:
    """Reference INT4 W4A16 group-quantized fused MoE in pure torch."""
    B = a.shape[0]
    topk = topk_ids.size(1)
    K_out = a.shape[1]

    out = torch.zeros(B, topk, K_out, dtype=torch.float32)
    for b in range(B):
        for t in range(topk):
            eid = topk_ids[b, t].item()
            x = a[b : b + 1].float()

            # Dequantize w1: [K, 2*N], groups along K (input dim)
            K_dim = w1_int4.shape[1]
            w1_dq = torch.zeros(K_dim, w1_int4.shape[2], dtype=torch.float32)
            for g in range(w1_s.shape[1]):
                k_start = g * group_size
                k_end = min((g + 1) * group_size, K_dim)
                zp = w1_zeros[eid, g, :].float() if w1_zeros is not None else 8.0
                w1_dq[k_start:k_end, :] = (
                    w1_int4[eid, k_start:k_end, :].float() - zp
                ) * w1_s[eid, g, :].float()

            ic = torch.matmul(x, w1_dq)  # [1, K] @ [K, 2*N] → [1, 2*N]
            ic = _silu_and_mul(ic)  # [1, N]

            # Dequantize w2: [N, K], groups along N (input dim)
            N_dim = w2_int4.shape[1]
            w2_dq = torch.zeros(N_dim, w2_int4.shape[2], dtype=torch.float32)
            for g in range(w2_s.shape[1]):
                n_start = g * group_size
                n_end = min((g + 1) * group_size, N_dim)
                zp = w2_zeros[eid, g, :].float() if w2_zeros is not None else 8.0
                w2_dq[n_start:n_end, :] = (
                    w2_int4[eid, n_start:n_end, :].float() - zp
                ) * w2_s[eid, g, :].float()

            oc = torch.matmul(ic, w2_dq)  # [1, N] @ [N, K] → [1, K]
            out[b, t] = oc.squeeze(0)

    return (out * topk_weight.unsqueeze(-1)).sum(dim=1).to(a.dtype)


def _make_int4_moe_weights(E, N, K, group_size, quant_algo):
    """Create INT4 MoE weights in GPTQ or AWQ packed format.

    Canonical layout (input × output):
      w1_int4: [E, K, 2*N]  w2_int4: [E, N, K]

    GPTQ packed (pack transposed weight along input/K dim):
      w1_packed: [E, K//8, 2*N]  w2_packed: [E, N//8, K]
      zeros: actual int4 zero points, same packing as weights

    AWQ packed (pack along output/N dim):
      w1_packed: [E, K, 2*N//8]  w2_packed: [E, N, K//8]
      zeros: actual int4 zero points, same packing as weights

    Returns:
        w1_int4, w2_int4,
        w1_packed, w2_packed,
        w1_zeros, w2_zeros,
        w1_zeros_packed, w2_zeros_packed,
        w1_s, w2_s
    """
    w1_int4 = torch.randint(0, 16, (E, K, 2 * N), dtype=torch.int32)
    w2_int4 = torch.randint(0, 16, (E, N, K), dtype=torch.int32)

    num_groups_w1 = K // group_size
    num_groups_w2 = N // group_size
    w1_s = (
        torch.randn(E, num_groups_w1, 2 * N, dtype=torch.bfloat16) * 0.01
    ).abs() + 0.001
    w2_s = (torch.randn(E, num_groups_w2, K, dtype=torch.bfloat16) * 0.01).abs() + 0.001

    if quant_algo == ops.CPUQuantAlgo.GPTQ:
        # Pack: canonical [E, K, 2*N] → transpose [E, 2*N, K] → GPTQ pack
        # [E, 2*N, K//8] → transpose [E, K//8, 2*N]
        w1_t = w1_int4.transpose(1, 2).contiguous()  # [E, 2*N, K]
        w1_packed = (
            torch.stack([_pack_int4_gptq(w1_t[e]) for e in range(E)])
            .transpose(1, 2)
            .contiguous()
        )  # [E, K//8, 2*N]
        w2_t = w2_int4.transpose(1, 2).contiguous()  # [E, K, N]
        w2_packed = (
            torch.stack([_pack_int4_gptq(w2_t[e]) for e in range(E)])
            .transpose(1, 2)
            .contiguous()
        )  # [E, N//8, K]
        w1_zeros = w2_zeros = None
        w1_zeros_packed = torch.full(
            (E, num_groups_w1, 2 * N // 8), 0x77777777, dtype=torch.int32
        )
        w2_zeros_packed = torch.full(
            (E, num_groups_w2, K // 8), 0x77777777, dtype=torch.int32
        )
    else:  # AWQ
        # Asymmetric: actual zero points, packed along output dim.
        w1_zeros = torch.randint(1, 15, (E, num_groups_w1, 2 * N), dtype=torch.int32)
        w2_zeros = torch.randint(1, 15, (E, num_groups_w2, K), dtype=torch.int32)
        w1_packed = torch.stack(
            [_pack_int4_awq(w1_int4[e]) for e in range(E)]
        )  # [E, K, 2*N//8]
        w2_packed = torch.stack(
            [_pack_int4_awq(w2_int4[e]) for e in range(E)]
        )  # [E, N, K//8]
        w1_zeros_packed = torch.stack(
            [_pack_int4_awq(w1_zeros[e]) for e in range(E)]
        )  # [E, K//gs, 2*N//8]
        w2_zeros_packed = torch.stack(
            [_pack_int4_awq(w2_zeros[e]) for e in range(E)]
        )  # [E, N//gs, K//8]

    return (
        w1_int4,
        w2_int4,
        w1_packed,
        w2_packed,
        w1_zeros,
        w2_zeros,
        w1_zeros_packed,
        w2_zeros_packed,
        w1_s,
        w2_s,
    )


INT4_MOE_CONFIGS = [
    # (N, K, E, topk, group_size)
    (256, 512, 8, 2, 128),
    (512, 256, 8, 2, 128),
    (512, 512, 8, 4, 128),
    (768, 2048, 8, 2, 128),
]


@pytest.mark.parametrize("M", [1, 2, 64, 121])
@pytest.mark.parametrize("N,K,E,topk,group_size", INT4_MOE_CONFIGS)
@pytest.mark.parametrize("quant_algo", [ops.CPUQuantAlgo.GPTQ, ops.CPUQuantAlgo.AWQ])
@pytest.mark.parametrize("seed", [0])
def test_int4_w4a16_cpu_fused_moe(M, N, K, E, topk, group_size, quant_algo, seed):
    """Test fused_experts_cpu INT4 W4A16 for both GPTQ and AWQ quant formats."""
    set_random_seed(seed)

    a = torch.randn(M, K, dtype=torch.bfloat16) / (0.5 * K**0.5)
    (
        w1_int4,
        w2_int4,
        w1_packed,
        w2_packed,
        w1_zeros,
        w2_zeros,
        w1_zeros_packed,
        w2_zeros_packed,
        w1_s,
        w2_s,
    ) = _make_int4_moe_weights(E, N, K, group_size, quant_algo)

    score = torch.randn(M, E, dtype=torch.bfloat16)
    score = torch.softmax(score, dim=-1, dtype=torch.float32)
    topk_weight, topk_ids = torch.topk(score, topk)
    topk_ids = topk_ids.to(torch.int32)

    ref_out = _ref_int4_moe(
        a,
        w1_int4,
        w2_int4,
        w1_zeros,
        w2_zeros,
        w1_s,
        w2_s,
        topk_weight,
        topk_ids,
        group_size,
    )

    from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
        prepare_int4_moe_layer_for_cpu,
    )

    (blocked_w1, blocked_w2, blocked_s1, blocked_s2, blocked_z1, blocked_z2) = (
        prepare_int4_moe_layer_for_cpu(
            w1_packed,
            w2_packed,
            w1_s,
            w2_s,
            quant_algo=quant_algo,
            w13_zeros=w1_zeros_packed,
            w2_zeros=w2_zeros_packed,
        )
    )

    out = ops.fused_experts_cpu(
        a.clone(),
        blocked_w1,
        blocked_w2,
        topk_weight,
        topk_ids,
        False,  # inplace
        ops.CPUQuantMethod.INT4_W4A8,
        blocked_s1,
        blocked_s2,
        blocked_z1,
        blocked_z2,
        None,  # block_size
        None,  # w1_bias
        None,  # w2_bias
        None,  # alpha
        None,  # limit
        True,  # is_vnni
    )
    torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)


# ===========================================================================
# INT8 W8A8 MoE
# ===========================================================================


def _quantize_per_channel(w):
    """Symmetric per-channel INT8 quantisation. w: [N, K] -> (int8, scale)."""
    amax = w.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12)
    scale = amax / 127.0
    w_q = (w / scale).round().clamp(-128, 127).to(torch.int8)
    return w_q, scale.float()


def _quantize_per_token(x):
    """Symmetric per-token INT8 quantisation. x: [M, K] -> (int8, scale)."""
    amax = x.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12)
    scale = amax / 127.0
    x_q = (x / scale).round().clamp(-128, 127).to(torch.int8)
    return x_q, scale.float()


def _ref_int8_moe(a, w1, w2, w1_s, w2_s, topk_weight, topk_ids):
    """Reference INT8 W8A8 per-channel fused MoE in pure torch."""
    B, D = a.shape
    topk = topk_ids.size(1)

    out = torch.zeros(B, topk, w2.shape[1], dtype=torch.float32)
    for b in range(B):
        for t in range(topk):
            eid = topk_ids[b, t].item()

            x = a[b : b + 1].float()
            x_q, x_s = _quantize_per_token(x)
            ic = torch.matmul(x_q.float(), w1[eid].float().t())
            ic = ic * x_s * w1_s[eid].view(1, -1)
            ic = _silu_and_mul(ic)

            ic_q, ic_s = _quantize_per_token(ic)
            oc = torch.matmul(ic_q.float(), w2[eid].float().t())
            oc = oc * ic_s * w2_s[eid].view(1, -1)
            out[b, t] = oc.squeeze(0)

    result = (out * topk_weight.unsqueeze(-1)).sum(dim=1)
    return result.to(a.dtype)


def _make_int8_moe_weights(E, N, K):
    factor = 1e-2
    w1_f = (torch.randn(E, 2 * N, K) - 0.5) * 2
    w2_f = (torch.randn(E, K, N) - 0.5) * 2

    w1_q_list, w1_s_list = [], []
    w2_q_list, w2_s_list = [], []
    for e in range(E):
        q, s = _quantize_per_channel(w1_f[e])
        w1_q_list.append(q)
        w1_s_list.append(s)
        q, s = _quantize_per_channel(w2_f[e])
        w2_q_list.append(q)
        w2_s_list.append(s)

    return (
        torch.stack(w1_q_list),
        torch.stack(w2_q_list),
        torch.stack(w1_s_list) * factor,
        torch.stack(w2_s_list) * factor,
    )


INT8_NUM_TOKENS = [1, 2, 64, 121]
INT8_MOE_CONFIGS = [
    # (N, K, E, topk)
    (256, 512, 8, 2),
    (512, 256, 8, 2),
    (512, 512, 8, 4),
    (768, 2048, 8, 2),
]


@pytest.mark.parametrize("M", INT8_NUM_TOKENS)
@pytest.mark.parametrize("N,K,E,topk", INT8_MOE_CONFIGS)
@pytest.mark.parametrize("seed", [0])
@pytest.mark.parametrize("is_vnni", [False, True])
@pytest.mark.parametrize("inplace", [False, True])
def test_int8_w8a8_cpu_fused_moe(M, N, K, E, topk, seed, is_vnni, inplace):
    """Test fused_experts_cpu INT8 W8A8 against torch reference."""
    set_random_seed(seed)

    a = torch.randn(M, K, dtype=torch.bfloat16) / (0.5 * K**0.5)
    w1_q, w2_q, w1_s, w2_s = _make_int8_moe_weights(E, N, K)

    score = torch.randn(M, E, dtype=torch.bfloat16)
    score = torch.softmax(score, dim=-1, dtype=torch.float32)
    topk_weight, topk_ids = torch.topk(score, topk)
    topk_ids = topk_ids.to(torch.int32)

    ref_out = _ref_int8_moe(a, w1_q, w2_q, w1_s, w2_s, topk_weight, topk_ids)

    w1 = _prepack_experts(w1_q) if is_vnni else w1_q
    w2 = _prepack_experts(w2_q) if is_vnni else w2_q

    out = ops.fused_experts_cpu(
        a.clone(),
        w1,
        w2,
        topk_weight,
        topk_ids,
        inplace,
        ops.CPUQuantMethod.INT8_W8A8,
        w1_s,
        w2_s,
        None,  # w1_zero
        None,  # w2_zero
        None,  # block_size
        None,  # w1_bias
        None,  # w2_bias
        None,  # alpha
        None,  # limit
        is_vnni,
    )
    torch.testing.assert_close(
        ref_out.bfloat16(),
        out,
        atol=2e-1,
        rtol=2e-1,
    )


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