# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import TYPE_CHECKING

import torch

from vllm.logger import init_logger
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.utils.math_utils import round_up

if TYPE_CHECKING:
    from flashinfer.fused_moe.core import ActivationType

logger = init_logger(__name__)


def activation_to_flashinfer_int(activation: MoEActivation) -> int:
    return activation_to_flashinfer_type(activation).value


def activation_to_flashinfer_type(activation: MoEActivation) -> "ActivationType":
    from flashinfer.fused_moe.core import ActivationType

    # silu and gelu are mapped to their gated versions SwiGLU and GeGLU respectively
    ACTIVATION_TO_FI_ACTIVATION = {
        MoEActivation.SILU_NO_MUL: ActivationType.Silu,
        MoEActivation.GELU_NO_MUL: ActivationType.Gelu,
        MoEActivation.SILU: ActivationType.Swiglu,
        # SwiGLU-OAI uses Swiglu; the OAI alpha/beta/clamp come from gemm1_* args.
        MoEActivation.SWIGLUOAI_UNINTERLEAVE: ActivationType.Swiglu,
        MoEActivation.GELU: ActivationType.Geglu,
        MoEActivation.GELU_TANH: ActivationType.Geglu,
        MoEActivation.RELU2_NO_MUL: ActivationType.Relu2,
        MoEActivation.SWIGLUOAI_UNINTERLEAVE: ActivationType.Swiglu,
    }
    return ACTIVATION_TO_FI_ACTIVATION[activation]


def swap_w13_to_w31(x: torch.Tensor) -> torch.Tensor:
    return (
        x.reshape(-1, 2, x.shape[-2] // 2, x.shape[-1]).flip(dims=[1]).reshape(x.shape)
    )


def rotate_weights_for_fi_trtllm_fp8_per_tensor_moe(
    gemm1_weights: torch.Tensor, gemm2_weights: torch.Tensor, is_gated_activation: bool
):
    """Shuffle weights for FI TRT-LLM Format"""
    from flashinfer import reorder_rows_for_gated_act_gemm, shuffle_matrix_a

    epilogue_tile_m = 128
    num_experts = gemm1_weights.shape[0]
    hidden_size = gemm1_weights.shape[-1]
    intermediate_size = gemm1_weights.shape[1] // 2

    # Reorder rows of W1 for fused gated activation
    gemm1_weights_fp8_interleaved = []
    for i in range(num_experts):
        gemm1_weights_fp8_interleaved.append(
            reorder_rows_for_gated_act_gemm(gemm1_weights[i])
            if is_gated_activation
            else gemm1_weights[i]
        )

    # Stack weights and scales for all experts
    gemm1_weights_fp8_interleaved = torch.stack(gemm1_weights_fp8_interleaved).reshape(
        num_experts, 2 * intermediate_size, hidden_size
    )

    # Shuffle weights and scaling factors for transposed mma output
    gemm1_weights_fp8_shuffled = []
    gemm2_weights_fp8_shuffled = []
    for i in range(num_experts):
        gemm1_weights_fp8_shuffled.append(
            shuffle_matrix_a(
                gemm1_weights_fp8_interleaved[i].view(torch.uint8), epilogue_tile_m
            )
        )

        gemm2_weights_fp8_shuffled.append(
            shuffle_matrix_a(gemm2_weights[i].view(torch.uint8), epilogue_tile_m)
        )

    # Stack weights for all experts
    gemm1_weights.data = torch.stack(gemm1_weights_fp8_shuffled).view(
        torch.float8_e4m3fn
    )
    gemm2_weights.data = torch.stack(gemm2_weights_fp8_shuffled).view(
        torch.float8_e4m3fn
    )


def convert_moe_weights_to_flashinfer_trtllm_block_layout(
    cache_permute_indices: dict[torch.Size, torch.Tensor],
    w13_weight: torch.Tensor,
    w2_weight: torch.Tensor,
    is_gated_act_gemm: bool = True,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Convert expert weights to FlashInfer's block layout.

    This reorders W13 and W2 into the expected epilogue-tiled block layout and
    returns the shuffled weight tensors.
    """
    if w13_weight.dtype != torch.bfloat16 or w2_weight.dtype != torch.bfloat16:
        raise ValueError(
            "Unquantized Moe Backend FlashInfer TRTLLM requires bfloat16 weights"
        )

    from flashinfer.fused_moe.core import (
        _maybe_get_cached_w3_w1_permute_indices,
        get_w2_permute_indices_with_cache,
    )

    epilogue_tile_m = 128
    block_k = 128

    # Reorder rows of W13 and W2 for fused gated activation and convert to the
    # block layout expected by the FlashInfer kernel.
    num_experts = w13_weight.shape[0]

    def _copy_permuted_expert_to_block_layout(
        out: torch.Tensor,
        expert_uint8: torch.Tensor,
        source_indices: torch.Tensor,
    ) -> None:
        expert_blocks = expert_uint8.view(
            expert_uint8.shape[0], out.shape[0], block_k
        ).permute(1, 0, 2)
        torch.index_select(
            expert_blocks,
            1,
            source_indices.to(expert_uint8.device),
            out=out,
        )

    w13_rows, w13_cols = w13_weight[0].view(torch.uint8).shape
    w2_rows, w2_cols = w2_weight[0].view(torch.uint8).shape
    w13_weights_shuffled_tensor = torch.empty(
        (num_experts, w13_cols // block_k, w13_rows, block_k),
        dtype=torch.uint8,
        device=w13_weight.device,
    )
    w2_weights_shuffled_tensor = torch.empty(
        (num_experts, w2_cols // block_k, w2_rows, block_k),
        dtype=torch.uint8,
        device=w2_weight.device,
    )

    for i in range(num_experts):
        w13_expert_uint8 = w13_weight[i].view(torch.uint8)

        permute_indices = _maybe_get_cached_w3_w1_permute_indices(
            cache_permute_indices,
            w13_expert_uint8,
            epilogue_tile_m,
            is_gated_act_gemm=is_gated_act_gemm,
        )
        if is_gated_act_gemm:
            rows = w13_expert_uint8.shape[0]
            permute_indices = (permute_indices + rows // 2) % rows
        _copy_permuted_expert_to_block_layout(
            w13_weights_shuffled_tensor[i],
            w13_expert_uint8,
            permute_indices,
        )

        permute_indices = get_w2_permute_indices_with_cache(
            cache_permute_indices,
            w2_weight[i].view(torch.uint8),
            epilogue_tile_m,
        )
        _copy_permuted_expert_to_block_layout(
            w2_weights_shuffled_tensor[i],
            w2_weight[i].view(torch.uint8),
            permute_indices,
        )

    return (
        w13_weights_shuffled_tensor.view(torch.bfloat16),
        w2_weights_shuffled_tensor.view(torch.bfloat16),
    )


def align_fp4_moe_weights_for_fi(
    w13: torch.Tensor,
    w13_scale: torch.Tensor,
    w2: torch.Tensor,
    w2_scale: torch.Tensor,
    is_act_and_mul: bool,
    min_alignment: int = 16,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]:
    """Pad intermediate size so FlashInfer kernels' alignment constraints hold.

    Some FlashInfer FP4 MoE kernels require the intermediate size
    used for GEMM to be divisible by a small alignment value. When this is
    not satisfied (e.g. with certain tensor-parallel sizes), we pad the
    gate/up and down projection weights along the intermediate dim.
    """

    # Current local intermediate size (per partition) is the K dimension of
    # the down projection.
    num_experts, hidden_size, intermediate = w2.shape
    intermediate *= 2  # because of packed FP4

    padded_intermediate = round_up(intermediate, min_alignment)

    if padded_intermediate == intermediate:
        return w13, w13_scale, w2, w2_scale, intermediate

    logger.info_once(
        "Padding intermediate size from %d to %d for up/down projection weights.",
        intermediate,
        padded_intermediate,
    )

    up_mult = 2 if is_act_and_mul else 1
    padded_gate_up_dim = up_mult * padded_intermediate

    # Pad w13 and w2 along its intermediate dimension.
    padded_w13 = w13.new_zeros((num_experts, padded_gate_up_dim, hidden_size // 2))
    padded_w13[:, : w13.shape[1], :] = w13

    padded_w2 = w2.new_zeros((num_experts, hidden_size, padded_intermediate // 2))
    padded_w2[:, :, : w2.shape[2]] = w2

    padded_w13_scale = w13_scale.new_zeros(
        (num_experts, padded_gate_up_dim, hidden_size // 16)
    )
    padded_w13_scale[:, : w13_scale.shape[1], :] = w13_scale

    padded_w2_scale = w2_scale.new_zeros(
        (num_experts, hidden_size, padded_intermediate // 16)
    )
    padded_w2_scale[:, :, : w2_scale.shape[2]] = w2_scale

    return padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_intermediate


def align_trtllm_fp4_moe_hidden_dim_for_fi(
    w13: torch.Tensor,
    w13_scale: torch.Tensor,
    w2: torch.Tensor,
    w2_scale: torch.Tensor,
    min_alignment: int = 256,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]:
    num_experts, gate_up_dim, packed_hidden_size = w13.shape
    hidden_size = packed_hidden_size * 2
    padded_hidden_size = round_up(hidden_size, min_alignment)

    if padded_hidden_size == hidden_size:
        return w13, w13_scale, w2, w2_scale, hidden_size

    logger.warning_once(
        "Padding hidden size from %d to %d for TRTLLM NVFP4 MoE weights. "
        "This requires activation slicing at runtime and may cause "
        "performance degradation.",
        hidden_size,
        padded_hidden_size,
    )

    padded_w13 = w13.new_zeros((num_experts, gate_up_dim, padded_hidden_size // 2))
    padded_w13[:, :, :packed_hidden_size] = w13

    padded_w13_scale = w13_scale.new_zeros(
        (num_experts, gate_up_dim, padded_hidden_size // 16)
    )
    padded_w13_scale[:, :, : w13_scale.shape[2]] = w13_scale

    padded_w2 = w2.new_zeros((num_experts, padded_hidden_size, w2.shape[2]))
    padded_w2[:, : w2.shape[1], :] = w2

    padded_w2_scale = w2_scale.new_zeros(
        (num_experts, padded_hidden_size, w2_scale.shape[2])
    )
    padded_w2_scale[:, : w2_scale.shape[1], :] = w2_scale

    return padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_hidden_size


def align_moe_weights_for_fi(
    w13: torch.Tensor, w2: torch.Tensor, is_act_and_mul: bool, min_alignment: int = 16
) -> tuple[torch.Tensor, torch.Tensor, int]:
    """Pad intermediate size so FlashInfer kernels' alignment constraints hold.

    Some FlashInfer MoE kernels require the (gated) intermediate size
    used for GEMM to be divisible by a small alignment value. When this is
    not satisfied (e.g. with certain tensor-parallel sizes), we pad the
    gate/up and down projection weights along the intermediate dim.
    """

    # Current local intermediate size (per partition) is the K dimension of
    # the down projection.
    num_experts, hidden_size, intermediate = w2.shape

    padded_intermediate = round_up(intermediate, min_alignment)

    if padded_intermediate == intermediate:
        return w13, w2, intermediate

    logger.info_once(
        "Padding intermediate size from %d to %d for up/down projection weights.",
        intermediate,
        padded_intermediate,
    )

    up_mult = 2 if is_act_and_mul else 1
    padded_gate_up_dim = up_mult * padded_intermediate

    # Pad w13 and w2 along its intermediate dimension.
    padded_w13 = w13.new_zeros((num_experts, padded_gate_up_dim, hidden_size))
    padded_w13[:, : w13.shape[1], :] = w13

    padded_w2 = w2.new_zeros((num_experts, hidden_size, padded_intermediate))
    padded_w2[:, :, :intermediate] = w2

    return padded_w13, padded_w2, padded_intermediate


def _shuffle_deepseek_fp8_moe_weights(
    w13: torch.Tensor,
    w2: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Preprocess DeepSeek FP8 block-scale weights for the FlashInfer TRT-LLM
    kernel using the shuffle + BlockMajorK layout variant.

    Returns 4D weight tensors in BlockMajorK layout
    (E, K/block_k, Mn, block_k)
    """
    from flashinfer import shuffle_matrix_a
    from flashinfer.fused_moe import convert_to_block_layout

    epilogue_tile_m = 64
    block_k = 128
    num_experts = w13.shape[0]

    M13, K13 = w13.shape[1], w13.shape[2]
    M2, K2 = w2.shape[1], w2.shape[2]
    w13_out = torch.empty(
        num_experts, K13 // block_k, M13, block_k, dtype=torch.uint8, device=w13.device
    )
    w2_out = torch.empty(
        num_experts, K2 // block_k, M2, block_k, dtype=torch.uint8, device=w2.device
    )

    for i in range(num_experts):
        t13 = shuffle_matrix_a(w13[i].view(torch.uint8), epilogue_tile_m)
        w13_out[i] = convert_to_block_layout(t13, block_k)

        t2 = shuffle_matrix_a(w2[i].view(torch.uint8), epilogue_tile_m)
        w2_out[i] = convert_to_block_layout(t2, block_k)

    return w13_out.view(torch.float8_e4m3fn), w2_out.view(torch.float8_e4m3fn)


def _shuffle_mxfp8_moe_weights(
    w13: torch.Tensor,
    w2: torch.Tensor,
    w13_scale: torch.Tensor,
    w2_scale: torch.Tensor,
    is_gated: bool,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Preprocess MXFP8 weights and scales for the FlashInfer TRT-LLM kernel.

    Following flashinfer/tests/moe/test_trtllm_gen_fused_moe.py:
      1. reorder_rows_for_gated_act_gemm  (interleave gate/up rows)
      2. shuffle_matrix_a                 (weight data layout shuffle)
      3. shuffle_matrix_sf_a              (scale factor layout shuffle)
    """
    from flashinfer import (
        reorder_rows_for_gated_act_gemm,
        shuffle_matrix_a,
        shuffle_matrix_sf_a,
    )

    epilogue_tile_m = 128
    num_experts = w13.shape[0]
    intermediate_size = w13.shape[1] // 2
    hidden_size = w13.shape[2]

    w13_interleaved: list[torch.Tensor] = []
    w13_scale_interleaved: list[torch.Tensor] = []
    for i in range(num_experts):
        if is_gated:
            w13_interleaved.append(
                reorder_rows_for_gated_act_gemm(
                    w13[i].reshape(2 * intermediate_size, -1)
                )
            )
            w13_scale_interleaved.append(
                reorder_rows_for_gated_act_gemm(
                    w13_scale[i].reshape(2 * intermediate_size, -1)
                )
            )
        else:
            w13_interleaved.append(w13[i])
            w13_scale_interleaved.append(w13_scale[i])

    w13_shuffled: list[torch.Tensor] = []
    w2_shuffled: list[torch.Tensor] = []
    w13_scale_shuffled: list[torch.Tensor] = []
    w2_scale_shuffled: list[torch.Tensor] = []
    for i in range(num_experts):
        w13_shuffled.append(
            shuffle_matrix_a(w13_interleaved[i].view(torch.uint8), epilogue_tile_m)
        )
        w2_shuffled.append(shuffle_matrix_a(w2[i].view(torch.uint8), epilogue_tile_m))
        w13_scale_shuffled.append(
            shuffle_matrix_sf_a(
                w13_scale_interleaved[i]
                .view(torch.uint8)
                .reshape(2 * intermediate_size, -1),
                epilogue_tile_m,
            )
        )
        w2_scale_shuffled.append(
            shuffle_matrix_sf_a(
                w2_scale[i].view(torch.uint8).reshape(hidden_size, -1),
                epilogue_tile_m,
            )
        )

    w13_out = torch.stack(w13_shuffled).view(torch.float8_e4m3fn)
    w2_out = torch.stack(w2_shuffled).view(torch.float8_e4m3fn)
    w13_scale_out = torch.stack(w13_scale_shuffled).reshape(w13_scale.shape)
    w2_scale_out = torch.stack(w2_scale_shuffled).reshape(w2_scale.shape)

    return w13_out, w2_out, w13_scale_out, w2_scale_out


def prepare_fp8_moe_layer_for_fi(
    layer: torch.nn.Module,
    w13: torch.Tensor,
    w2: torch.Tensor,
    w13_scale: torch.Tensor,
    w13_input_scale: torch.Tensor | None,
    w2_scale: torch.Tensor,
    w2_input_scale: torch.Tensor | None,
    is_trtllm: bool = False,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Convert Fp8 MoE weights to flashinfer kernel format

    Note that for trtllm we update the model state dict
    with the scale format needed for these kernels.

    Note that for per-tensor, we update the layer's
    intermediate size if the weights needed padding.
    """

    assert hasattr(layer.moe_config, "is_act_and_mul")
    block_quant = (
        hasattr(layer, "weight_block_size") and layer.weight_block_size is not None
    )
    is_mxfp8 = block_quant and w13_scale.dtype == torch.uint8
    is_deepseek_fp8 = block_quant and not is_mxfp8
    is_gated = layer.activation.is_gated

    # MXFP8 TRT-LLM requires W31 swap + reorder + shuffle.
    if is_mxfp8 and is_trtllm:
        # FlashInfer TRT-LLM SwiGLU expects [up; gate] but vLLM stores
        # [gate; up].  Swap both weights and scales before interleaving.
        if layer.moe_config.is_act_and_mul:
            w13 = swap_w13_to_w31(w13)
            # Scales may be 2D [E, flat] from _quantize_mxfp8_moe_weight;
            # reshape to 3D so swap_w13_to_w31 can flip the two halves,
            # then flatten back.
            if w13_scale.ndim == 2:
                num_rows = w13.shape[1]  # 2 * intermediate_size
                w13_scale = w13_scale.reshape(w13_scale.shape[0], num_rows, -1)
                w13_scale = swap_w13_to_w31(w13_scale)
                w13_scale = w13_scale.reshape(w13_scale.shape[0], -1)
            else:
                w13_scale = swap_w13_to_w31(w13_scale)

        w13, w2, w13_scale, w2_scale = _shuffle_mxfp8_moe_weights(
            w13, w2, w13_scale, w2_scale, is_gated
        )
        return w13, w2, w13_scale, w2_scale

    # Some FI MoE kernels require internal alignment of 16
    # for the gate-up proj. Pad the weights to respect this.
    if not block_quant:
        min_alignment = 16 if is_gated else 128
        w13, w2, new_intermediate = align_moe_weights_for_fi(
            w13,
            w2,
            layer.moe_config.is_act_and_mul,
            min_alignment,
        )
        layer.moe_config.intermediate_size_per_partition = new_intermediate

    # FI kernels require W31 layout rather than W13.
    if layer.moe_config.is_act_and_mul:
        w13 = swap_w13_to_w31(w13)
        if block_quant:
            w13_scale = swap_w13_to_w31(w13_scale)

    # DeepSeekFp8 TRT-LLM: shuffle weights into BlockMajorK layout.
    if is_deepseek_fp8 and is_trtllm:
        w13, w2 = _shuffle_deepseek_fp8_moe_weights(w13, w2)

    # FI TRT-LLM FP8 per-tensor MoE kernel requires weight shuffle
    # and registration of alpha scales.
    if is_trtllm and not block_quant:
        assert w13_input_scale is not None
        assert w2_input_scale is not None

        rotate_weights_for_fi_trtllm_fp8_per_tensor_moe(w13, w2, is_gated)

    # Clamp block scales to avoid NaN from the FlashInfer CUTLASS kernel.
    # Some FP8 models have near-zero block scales (~1e-23) for dead/unused
    # experts. The CUTLASS kernel doesn't handle these correctly on Hopper
    # (SM 9.0), producing NaN instead of near-zero output. Clamping to a
    # small minimum prevents this without affecting model accuracy since
    # these experts' effective weights are already zero.
    if block_quant:
        _FI_CUTLASS_MIN_BLOCK_SCALE = 1e-10
        w13_scale.clamp_(min=_FI_CUTLASS_MIN_BLOCK_SCALE)
        w2_scale.clamp_(min=_FI_CUTLASS_MIN_BLOCK_SCALE)

    return w13, w2, w13_scale, w2_scale