# SPDX-License-Identifier: Apache-2.0
"""Fused MoE kernel."""
import functools
import json
import os
from typing import Any, Callable, Dict, List, Optional, Tuple

import torch
import triton
import triton.language as tl
import aiter.ops.triton.utils.arch_info as arch_info

from aiter.ops.triton.moe_op import fused_moe as invoke_fused_moe_kernel, support_mls
# from vllm import _custom_ops as ops
from aiter.ops.triton.utils.moe_config_utils import get_optimal_moe_config_func
from aiter import silu_and_mul, gelu_and_mul, moe_sum
from aiter import per_token_quant_hip, per_block_quant_wrapper
from aiter import sgl_moe_align_block_size as sgl_moe_align_block_size_aiter
from aiter import moe_align_block_size as moe_align_block_size_aiter
from aiter.jit.utils.torch_guard import torch_compile_guard
from aiter import dtypes,moe_sorting_fwd

device_name = arch_info.get_device()


def get_moe_sum_config(M, top_k, N):
    if M < 32:
        return {"BLOCK_SIZE": 128, "num_warps": 1}
    else:
        return {"BLOCK_SIZE": 512, "num_warps": 4}

# def generate_sum_configs():
#     configs = []
#     for block_n in [32, 64, 128, 256, 512, 1024, 2048, 4096]:
#         for num_warps in [1, 2, 4]:
#             for num_stages in [1, 2]:
#                 config = triton.Config({
#                     'BLOCK_SIZE': block_n,
#                 }, num_warps=num_warps, num_stages=num_stages)
#                 configs.append(config)
#     return configs

# @triton.autotune(
#     key=['M', 'N', 'top_k','compute_type'],
#     configs=generate_sum_configs(),
#     # configs = [
#     #             triton.Config({'BLOCK_SIZE': 64  }, num_warps=1),
#     #             triton.Config({'BLOCK_SIZE': 128 }, num_warps=1),
#     #             triton.Config({'BLOCK_SIZE': 256 }, num_warps=4),
#     #             triton.Config({'BLOCK_SIZE': 512 }, num_warps=4),
#     #           ],
#     perf_debug=True,
# )
@triton.heuristics({
    "n_dividable": lambda args: (args["N"] % args["BLOCK_SIZE"]) == 0,
})
@triton.jit
def moe_sum_kernel(
    output_ptr,             # [M, N]
    input_ptr,              # [M, top_k, N]
    M,
    N: tl.constexpr,
    top_k: tl.constexpr,
    routed_scaling_factor,
    BLOCK_SIZE: tl.constexpr,
    stride_output_m,
    stride_output_n,
    stride_input_m,
    stride_input_k,
    stride_input_n,
    compute_type: tl.constexpr,
    n_dividable: tl.constexpr,
):
    tl.assume(stride_output_m >= 0)
    tl.assume(stride_output_n >= 0)
    tl.assume(stride_input_m >= 0)
    tl.assume(stride_input_k >= 0)
    tl.assume(stride_input_n >= 0)

    num_pid_n = tl.cdiv(N, BLOCK_SIZE)
    pid = tl.program_id(axis=0)
    pid_m = pid // num_pid_n
    pid_n = pid % num_pid_n
    offs_n = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask_n = offs_n < N

    acc = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
    input_row_ptr = input_ptr + pid_m.to(tl.int64) * stride_input_m
    for k in range(top_k):
        input_ptrs = input_row_ptr + (
            k * stride_input_k
            + offs_n * stride_input_n
        )
        if n_dividable:
            x = tl.load(input_ptrs)
        else:
            x = tl.load(input_ptrs, mask=mask_n, other=0.0)
        acc += x.to(tl.float32)

    acc *= routed_scaling_factor
    acc = acc.to(compute_type)
    output_ptrs = output_ptr + (pid_m.to(tl.int64) * stride_output_m + offs_n * stride_output_n)
    if n_dividable:
        tl.store(output_ptrs, acc)
    else:
        tl.store(output_ptrs, acc, mask=mask_n)


def triton_moe_sum(input_tensor,
                   output_tensor,
                   routed_scaling_factor: float = 1.0):
    """
    1D tile version of moe_sum.

    Args:
        input_tensor: [M, top_k, N]
        output_tensor: [M, N]
    """
    M, top_k, N = input_tensor.shape

    assert output_tensor.dtype == torch.float16 or \
           output_tensor.dtype == torch.bfloat16 or \
           output_tensor.dtype == torch.float32

    if output_tensor.dtype == torch.float16:
        compute_type = tl.float16
    elif output_tensor.dtype == torch.bfloat16:
        compute_type = tl.bfloat16
    elif output_tensor.dtype == torch.float32:
        compute_type = tl.float32

    assert input_tensor.is_contiguous()
    assert output_tensor.is_contiguous()
    assert input_tensor.shape[0] == output_tensor.shape[0]
    assert input_tensor.shape[2] == output_tensor.shape[1]

    # 计算grid
    config = get_moe_sum_config(M, top_k, N)
    grid = (M * triton.cdiv(N, config["BLOCK_SIZE"]),)
    # grid = lambda META: (M * triton.cdiv(N, META['BLOCK_SIZE']), )

    moe_sum_kernel[grid](
        output_tensor,
        input_tensor,
        M,
        N,
        top_k,
        routed_scaling_factor,
        stride_output_m=output_tensor.stride(0),
        stride_output_n=output_tensor.stride(1),
        stride_input_m=input_tensor.stride(0),
        stride_input_k=input_tensor.stride(1),
        stride_input_n=input_tensor.stride(2),
        compute_type=compute_type,
        **config,
    )

    return output_tensor


# @triton.autotune(
#     configs=[
#         triton.Config(
#             {"BLOCK_SIZE_M": bm, "BLOCK_SIZE_N": bn},
#             num_warps=nw,
#         )
#         for bm in (1, 4, 16, 32, 64, 128)
#         for bn in (32, 128, 256, 512,)
#         for nw in (1, 2, 4)
#     ],
#     key=["M", "N", "ACT"],
#     perf_debug=True,
# )
@triton.heuristics(
    {
        "M_DIV": lambda args: (args["M"] % args["BLOCK_SIZE_M"]) == 0,
        "N_DIV": lambda args: (args["N"] % args["BLOCK_SIZE_N"]) == 0,
    }
)
@triton.jit
def activation_and_mul_kernel(
    out_ptr,
    in_ptr,
    M,
    N,
    stride_in0,
    stride_in1,
    stride_out0,
    stride_out1,
    ACT: tl.constexpr,
    M_DIV: tl.constexpr,
    N_DIV: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
):
    tl.assume(M > 0)
    tl.assume(N > 0)
    tl.assume(stride_in0 >= 0)
    tl.assume(stride_in1 >= 0)
    tl.assume(stride_out0 >= 0)
    tl.assume(stride_out1 >= 0)

    pid_m = tl.program_id(1)
    pid_n = tl.program_id(0)

    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)

    mask_m = offs_m < M
    mask_n = offs_n < N
    mask = mask_m[:, None] & mask_n[None, :]

    in_base = in_ptr + offs_m[:, None] * stride_in0 + offs_n[None, :] * stride_in1
    if M_DIV and N_DIV:
        x0 = tl.load(in_base)
        x1 = tl.load(in_base + N * stride_in1)
    else:
        x0 = tl.load(in_base, mask=mask, other=0.0)
        x1 = tl.load(in_base + N * stride_in1, mask=mask, other=0.0)

    x0_f = x0.to(tl.float32)
    x1_f = x1.to(tl.float32)
    if ACT == 0:
        act = x0_f * (1.0 / (1.0 + tl.exp(-x0_f)))
    else:
        act = x0_f * 0.5 * (1.0 + tl.erf(x0_f * 0.7071067811865476))
    y = act * x1_f

    out_ptrs = out_ptr + offs_m[:, None] * stride_out0 + offs_n[None, :] * stride_out1
    if M_DIV and N_DIV:
        tl.store(out_ptrs, y)
    else:
        tl.store(out_ptrs, y, mask=mask)

def get_triton_activation_and_mul_config(M, N):
    if M <= 512:
      return {"BLOCK_SIZE_M": 1, "BLOCK_SIZE_N": 256, "num_warps": 1}
    return {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 128, "num_warps": 4}


@triton.heuristics(
    {
        "M_DIV": lambda args: (args["M"] % args["BLOCK_SIZE_M"]) == 0,
        "N_DIV": lambda args: (args["N"] % args["BLOCK_SIZE_N"]) == 0,
    }
)
@triton.jit
def relu2_kernel(
    out_ptr,
    in_ptr,
    M,
    N,
    stride_in0,
    stride_in1,
    stride_out0,
    stride_out1,
    M_DIV: tl.constexpr,
    N_DIV: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
):
    tl.assume(M > 0)
    tl.assume(N > 0)
    pid_m = tl.program_id(1)
    pid_n = tl.program_id(0)
    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    mask_m = offs_m < M
    mask_n = offs_n < N
    mask = mask_m[:, None] & mask_n[None, :]
    in_base = in_ptr + offs_m[:, None] * stride_in0 + offs_n[None, :] * stride_in1
    if M_DIV and N_DIV:
        x = tl.load(in_base)
    else:
        x = tl.load(in_base, mask=mask, other=0.0)
    xf = x.to(tl.float32)
    y = tl.where(xf > 0.0, xf * xf, 0.0)
    out_ptrs = out_ptr + offs_m[:, None] * stride_out0 + offs_n[None, :] * stride_out1
    if M_DIV and N_DIV:
        tl.store(out_ptrs, y.to(x.dtype))
    else:
        tl.store(out_ptrs, y.to(x.dtype), mask=mask)


def triton_relu2(out: torch.Tensor, inp: torch.Tensor) -> None:
    """Elementwise ReLU² (no gate / no second-path mul). out.shape == inp.shape."""
    assert inp.shape == out.shape
    assert inp.is_contiguous() and out.is_contiguous()
    M, N = inp.shape
    config = get_triton_activation_and_mul_config(M, N)
    grid = (
        triton.cdiv(N, config["BLOCK_SIZE_N"]),
        triton.cdiv(M, config["BLOCK_SIZE_M"]),
    )
    relu2_kernel[grid](
        out,
        inp,
        M,
        N,
        inp.stride(0),
        inp.stride(1),
        out.stride(0),
        out.stride(1),
        **config,
    )


def triton_silu_and_mul(out: torch.Tensor, input: torch.Tensor) -> None:
    assert input.shape[-1] % 2 == 0
    assert input.is_contiguous()
    assert out.is_contiguous()
    M = input.numel() // input.shape[-1]
    N = input.shape[-1] // 2
    input_2d = input.view(M, input.shape[-1])
    out_2d = out.view(M, N)


    # grid = lambda META: (
    #     triton.cdiv(N, META["BLOCK_SIZE_N"]),
    #     triton.cdiv(M, META["BLOCK_SIZE_M"]),
    # )

    config = get_triton_activation_and_mul_config(M, N)
    grid = (
        triton.cdiv(N, config["BLOCK_SIZE_N"]),
        triton.cdiv(M, config["BLOCK_SIZE_M"]),
    )

    activation_and_mul_kernel[grid](
        out_2d,
        input_2d,
        M,
        N,
        input_2d.stride(0),
        input_2d.stride(1),
        out_2d.stride(0),
        out_2d.stride(1),
        ACT=0,
        **config,
    )


def triton_gelu_and_mul(out: torch.Tensor, input: torch.Tensor) -> None:
    assert input.shape[-1] % 2 == 0
    assert input.is_contiguous()
    assert out.is_contiguous()
    M = input.numel() // input.shape[-1]
    N = input.shape[-1] // 2
    input_2d = input.view(M, input.shape[-1])
    out_2d = out.view(M, N)

    # grid = lambda META: (
    #     triton.cdiv(M, META["BLOCK_SIZE_M"]),
    #     triton.cdiv(N, META["BLOCK_SIZE_N"]),
    # )
    config = get_triton_activation_and_mul_config(M, N)
    grid = (triton.cdiv(M, config['BLOCK_SIZE_M']) * triton.cdiv(N, config['BLOCK_SIZE_N']),)
    activation_and_mul_kernel[grid](
        out_2d,
        input_2d,
        M,
        N,
        input_2d.stride(0),
        input_2d.stride(1),
        out_2d.stride(0),
        out_2d.stride(1),
        ACT=1,
        **config,
    )


def ceil_div(a, b):
    return (a + b - 1) // b

@triton.jit
def moe_align_block_size_stage1(
    topk_ids_ptr,
    tokens_cnts_ptr,
    num_experts: tl.constexpr,
    numel: tl.constexpr,
    tokens_per_thread: tl.constexpr,
):
    pid = tl.program_id(0)

    start_idx = pid * tokens_per_thread

    off_c = (pid + 1) * num_experts

    for i in range(tokens_per_thread):
        if start_idx + i < numel:
            idx = tl.load(topk_ids_ptr + start_idx + i)
            token_cnt = tl.load(tokens_cnts_ptr + off_c + idx)
            tl.store(tokens_cnts_ptr + off_c + idx, token_cnt + 1)


@triton.jit
def moe_align_block_size_stage2(
    tokens_cnts_ptr,
    num_experts: tl.constexpr,
):
    pid = tl.program_id(0)

    last_cnt = 0
    for i in range(1, num_experts + 1):
        token_cnt = tl.load(tokens_cnts_ptr + i * num_experts + pid)
        last_cnt = last_cnt + token_cnt
        tl.store(tokens_cnts_ptr + i * num_experts + pid, last_cnt)


@triton.jit
def moe_align_block_size_stage3(
    total_tokens_post_pad_ptr,
    tokens_cnts_ptr,
    cumsum_ptr,
    num_experts: tl.constexpr,
    block_size: tl.constexpr,
):
    last_cumsum = 0
    off_cnt = num_experts * num_experts
    for i in range(1, num_experts + 1):
        token_cnt = tl.load(tokens_cnts_ptr + off_cnt + i - 1)
        last_cumsum = last_cumsum + tl.cdiv(token_cnt, block_size) * block_size
        tl.store(cumsum_ptr + i, last_cumsum)
    tl.store(total_tokens_post_pad_ptr, last_cumsum)


@triton.jit
def moe_align_block_size_stage4(
    topk_ids_ptr,
    sorted_token_ids_ptr,
    expert_ids_ptr,
    tokens_cnts_ptr,
    cumsum_ptr,
    num_experts: tl.constexpr,
    block_size: tl.constexpr,
    numel: tl.constexpr,
    tokens_per_thread: tl.constexpr,
):
    pid = tl.program_id(0)
    start_idx = tl.load(cumsum_ptr + pid)
    end_idx = tl.load(cumsum_ptr + pid + 1)

    for i in range(start_idx, end_idx, block_size):
        tl.store(expert_ids_ptr + i // block_size, pid)

    start_idx = pid * tokens_per_thread
    off_t = pid * num_experts

    for i in range(start_idx, tl.minimum(start_idx + tokens_per_thread,
                                         numel)):
        expert_id = tl.load(topk_ids_ptr + i)
        token_cnt = tl.load(tokens_cnts_ptr + off_t + expert_id)
        rank_post_pad = token_cnt + tl.load(cumsum_ptr + expert_id)
        tl.store(sorted_token_ids_ptr + rank_post_pad, i)
        tl.store(tokens_cnts_ptr + off_t + expert_id, token_cnt + 1)


# Triton implementation based on:
# https://github.com/sgl-project/sglang/commit/ba5112ff691d791a9e38c6c71f59324a5fcb49d0
def moe_align_block_size_triton(
    topk_ids: torch.Tensor,
    num_experts: int,
    block_size: int,
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
) -> None:
    numel = topk_ids.numel()
    grid = (num_experts, )
    tokens_cnts = torch.zeros((num_experts + 1, num_experts),
                              dtype=torch.int32,
                              device=topk_ids.device)
    cumsum = torch.zeros((num_experts + 1, ),
                         dtype=torch.int32,
                         device=topk_ids.device)
    tokens_per_thread = ceil_div(numel, num_experts)

    moe_align_block_size_stage1[grid](
        topk_ids,
        tokens_cnts,
        num_experts,
        numel,
        tokens_per_thread,
    )
    moe_align_block_size_stage2[grid](
        tokens_cnts,
        num_experts,
    )
    moe_align_block_size_stage3[(1, )](
        num_tokens_post_pad,
        tokens_cnts,
        cumsum,
        num_experts,
        block_size,
    )
    moe_align_block_size_stage4[grid](
        topk_ids,
        sorted_token_ids,
        expert_ids,
        tokens_cnts,
        cumsum,
        num_experts,
        block_size,
        numel,
        tokens_per_thread,
    )


def moe_align_block_size(
    topk_ids: torch.Tensor,
    block_size: int,
    num_experts: int,
    expert_map: torch.Tensor = None
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Aligns the token distribution across experts to be compatible with block
    size for matrix multiplication.

    Parameters:
    - topk_ids: A tensor of shape [total_tokens, top_k] representing the
        top-k expert indices for each token.
    - block_size: The block size used in block matrix multiplication.
    - num_experts: The total number of experts.
    - expert_map: A tensor of shape [num_experts] that maps the expert index
        from the global space to the local index space of the current
        expert parallel shard. If the expert is not in the current expert
        parallel shard, the mapping is set to -1.

    Returns:
    - sorted_token_ids: A tensor containing the sorted token indices according
        to their allocated expert.
    - expert_ids: A tensor indicating the assigned expert index for each block.
    - num_tokens_post_padded: The total number of tokens after padding,
        ensuring divisibility by block_size.

    This function pads the number of tokens that each expert needs to process
    so that it is divisible by block_size.
    Padding ensures that during block matrix multiplication, the dimensions
    align correctly.

    Example:
    Given topk_ids = [[2, 3, 4], [1, 2, 4], [1, 3, 4], [1, 2, 3]],
    block_size = 4, and num_experts = 4:
    - We initially have 12 tokens (after repeating 'top_k' times) and 4 experts,
        with each expert needing to process 3 tokens.
    - As block_size is 4, we pad 1 token for each expert.
    - First, flatten topk_ids to [2, 3, 4, 1, 2, 4, 1, 3, 4, 1, 2, 3].
    - Then append padding tokens [12, 12, 12, 12] for each block.
    - After sorting by expert index, we obtain token_ids
        [3, 6, 9, 12, 0, 4, 10, 12, 1, 7, 11, 12, 2, 5, 8, 12].
        Tokens 12 are non-existent (padding) and are ignored in
        the subsequent matrix multiplication.
    - The padding ensures that the total number of tokens is now divisible
        by block_size for proper block matrix operations.
    """
    max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
    sorted_ids = torch.empty((max_num_tokens_padded, ),
                             dtype=torch.int32,
                             device=topk_ids.device)
    sorted_ids.fill_(topk_ids.numel())
    max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
    # Expert ids must be zeroed out to prevent index out of bounds error while
    # mapping global expert ids to local expert ids in expert parallelism.
    expert_ids = torch.zeros((max_num_m_blocks, ),
                             dtype=torch.int32,
                             device=topk_ids.device)
    num_tokens_post_pad = torch.empty((1),
                                      dtype=torch.int32,
                                      device=topk_ids.device)
    if num_experts >= 224:
        if num_experts != 256:
            moe_align_block_size_triton(
                topk_ids,
                num_experts,
                block_size,
                sorted_ids,
                expert_ids,
                num_tokens_post_pad,
            )
        else:
            # Currently requires num_experts=256
            # Note: sgl_moe_align_block_size has 10x performance compared to moe_align_block_size,
            # but this functionis removed since vllm.0.9.2. vLLM will check this issue.
            # if hasattr(ops, 'sgl_moe_align_block_size'):
            #     # ops.sgl_moe_align_block_size(
            #         topk_ids,
            #         num_experts,
            #         block_size,
            #         sorted_ids,
            #         expert_ids,
            #         num_tokens_post_pad,
            #     )
            # else:
            #     ops.moe_align_block_size(
            #         topk_ids,
            #         num_experts,
            #         block_size,
            #         sorted_ids,
            #         expert_ids,
            #         num_tokens_post_pad,
            #     )
            sgl_moe_align_block_size_aiter(
                topk_ids,
                num_experts,
                block_size,
                sorted_ids,
                expert_ids,
                num_tokens_post_pad,
            )
    else:
        # ops.moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids,
        #                          expert_ids, num_tokens_post_pad)
        moe_align_block_size_aiter(topk_ids, num_experts, block_size, sorted_ids,
                                   expert_ids, num_tokens_post_pad)
    if expert_map is not None:
        expert_ids = expert_map[expert_ids]

    return sorted_ids, expert_ids, num_tokens_post_pad

def moe_sorting_ck(
    topk_ids,
    topk_weights,
    num_experts,
    model_dim,
    moebuf_dtype,
    block_size=32,
    expert_mask=None,
):
    device = topk_ids.device
    M, topk = topk_ids.shape
    topk = topk_ids.shape[1]
    max_num_tokens_padded = topk_ids.numel() + num_experts * block_size - topk
    max_num_m_blocks = int((max_num_tokens_padded + block_size - 1) // block_size)
    sorted_ids = torch.empty((max_num_tokens_padded,), dtype=dtypes.i32, device=device)
    sorted_weights = torch.empty(
        (max_num_tokens_padded,), dtype=dtypes.fp32, device=device
    )
    sorted_expert_ids = torch.empty(
        (max_num_m_blocks,), dtype=dtypes.i32, device=device
    )
    tokens_positions_per_expert = torch.empty(
        (num_experts*2,), dtype=dtypes.i32, device=device
    )
    num_valid_ids = torch.empty((1), dtype=dtypes.i32, device=device)
    moe_buf = torch.empty((M, model_dim), dtype=moebuf_dtype, device=device)

    # for now, moe_sorting_fwd only support int32 topk_ids
    if topk_ids.dtype != dtypes.i32:
        topk_ids = topk_ids.to(dtypes.i32)

    moe_sorting_fwd(
        topk_ids,
        topk_weights,
        sorted_ids,
        sorted_weights,
        sorted_expert_ids,
        tokens_positions_per_expert,
        num_valid_ids,
        moe_buf,
        num_experts,
        block_size,
        expert_mask,
    )
    return sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, tokens_positions_per_expert, moe_buf


def inplace_fused_experts(hidden_states: torch.Tensor,
                          w1: torch.Tensor,
                          w2: torch.Tensor,
                          topk_weights: torch.Tensor,
                          topk_ids: torch.Tensor,
                          activation: Optional[str] = None,
                          use_fp8_w8a8: bool = False,
                          use_int8_w8a8: bool = False,
                          use_int8_w8a16: bool = False,
                          use_int4_w4a16: bool = False,
                          global_num_experts: int = -1,
                          expert_map: Optional[torch.Tensor] = None,
                          w1_scale: Optional[torch.Tensor] = None,
                          w2_scale: Optional[torch.Tensor] = None,
                          w1_zp: Optional[torch.Tensor] = None,
                          w2_zp: Optional[torch.Tensor] = None,
                          a1_scale: Optional[torch.Tensor] = None,
                          a2_scale: Optional[torch.Tensor] = None,
                          block_shape: Optional[List[int]] = None) -> None:
    if activation is None:
        activation = "silu"
    fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids, True,
                       activation, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16,
                       use_int4_w4a16, global_num_experts, expert_map,
                       w1_scale, w2_scale, w1_zp, w2_zp, a1_scale, a2_scale,
                       block_shape)

def outplace_fused_experts(
        hidden_states: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        activation: Optional[str] = None,
        use_fp8_w8a8: bool = False,
        use_int8_w8a8: bool = False,
        use_int8_w8a16: bool = False,
        use_int4_w4a16: bool = False,
        global_num_experts: int = -1,
        expert_map: Optional[torch.Tensor] = None,
        w1_scale: Optional[torch.Tensor] = None,
        w2_scale: Optional[torch.Tensor] = None,
        w1_zp: Optional[torch.Tensor] = None,
        w2_zp: Optional[torch.Tensor] = None,
        a1_scale: Optional[torch.Tensor] = None,
        a2_scale: Optional[torch.Tensor] = None,
        block_shape: Optional[List[int]] = None) -> torch.Tensor:
    if activation is None:
        activation = "silu"
    return fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids,
                              False, activation, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16,
                              use_int4_w4a16, global_num_experts, expert_map,
                              w1_scale, w2_scale, w1_zp, w2_zp, a1_scale,
                              a2_scale, block_shape)

def fused_experts(hidden_states: torch.Tensor,
                  w1: torch.Tensor,
                  w2: torch.Tensor,
                  topk_weights: torch.Tensor,
                  topk_ids: torch.Tensor,
                  inplace: bool = False,
                  activation: Optional[str] = None,
                  use_fp8_w8a8: bool = False,
                  use_int8_w8a8: bool = False,
                  use_int8_w8a16: bool = False,
                  use_int4_w4a16: bool = False,
                  global_num_experts: int = -1,
                  expert_map: Optional[torch.Tensor] = None,
                  w1_scale: Optional[torch.Tensor] = None,
                  w2_scale: Optional[torch.Tensor] = None,
                  w1_zp: Optional[torch.Tensor] = None,
                  w2_zp: Optional[torch.Tensor] = None,
                  a1_scale: Optional[torch.Tensor] = None,
                  a2_scale: Optional[torch.Tensor] = None,
                  block_shape: Optional[List[int]] = None) -> torch.Tensor:
    if activation is None:
        activation = 'silu'
    if inplace:
        torch.ops.vllm.inplace_fused_experts(
            hidden_states, w1, w2, topk_weights, topk_ids, activation,
            use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, global_num_experts,
            expert_map, w1_scale, w2_scale, w1_zp, w2_zp, a1_scale, a2_scale,
            block_shape)
        return hidden_states
    else:
        return torch.ops.vllm.outplace_fused_experts(
            hidden_states, w1, w2, topk_weights, topk_ids, activation,
            use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, global_num_experts,
            expert_map, w1_scale, w2_scale, w1_zp, w2_zp, a1_scale, a2_scale,
            block_shape)

def fused_moe_fake(
        hidden_states: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        odtype:torch.dtype, #compute or output type for i8& f8
        inplace: bool = False,
        activation: str = "silu",
        use_fp8_w8a8: bool = False,
        use_int8_w8a8: bool = False,
        use_int8_w8a16: bool = False,
        use_int4_w4a16: bool = False,
        use_int4_w4a8: bool = False,
        per_channel_quant: bool = False,
        global_num_experts: int = -1,
        expert_map: Optional[torch.Tensor] = None,
        w1_scale: Optional[torch.Tensor] = None,
        w2_scale: Optional[torch.Tensor] = None,
        w1_zp: Optional[torch.Tensor] = None,
        w2_zp: Optional[torch.Tensor] = None,
        a1_scale: Optional[torch.Tensor] = None,
        a2_scale: Optional[torch.Tensor] = None,
        block_shape: Optional[List[int]] = None,
        routed_scaling_factor: Optional[float] = 1.0,
        fn_key: Optional[str] = None,
) -> torch.Tensor:
    device = topk_ids.device
    M, topk = topk_ids.shape
    dtype = odtype
    # E, model_dim, inter_dim = get_inter_dim(w1.shape, w2.shape)
    # FIXME: W2.size must be same as hidden_dim
    moe_buf = torch.empty(hidden_states.shape, dtype=dtype, device=device)
    return moe_buf
@functools.lru_cache()
def _bottom_moe_use_mls():
    return support_mls()


@torch_compile_guard(gen_fake=fused_moe_fake)
def fused_experts_impl(hidden_states: torch.Tensor,
                       w1: torch.Tensor,
                       w2: torch.Tensor,
                       topk_weights: torch.Tensor,
                       topk_ids: torch.Tensor,
                       odtype:torch.dtype, #compute or output type for i8& f8
                       inplace: bool = False,
                       activation: str = "silu",
                       use_fp8_w8a8: bool = False,
                       use_int8_w8a8: bool = False,
                       use_int8_w8a16: bool = False,
                       use_int4_w4a16: bool = False,
                       use_int4_w4a8: bool = False,
                       per_channel_quant: bool = False,
                       global_num_experts: int = -1,
                       expert_map: Optional[torch.Tensor] = None,
                       w1_scale: Optional[torch.Tensor] = None,
                       w2_scale: Optional[torch.Tensor] = None,
                       w1_zp: Optional[torch.Tensor] = None,
                       w2_zp: Optional[torch.Tensor] = None,
                       a1_scale: Optional[torch.Tensor] = None,
                       a2_scale: Optional[torch.Tensor] = None,
                       block_shape: Optional[List[int]] = None,
                       routed_scaling_factor: Optional[float] = 1.0)-> torch.Tensor:
    if routed_scaling_factor is None:
        routed_scaling_factor = 1.0

    # Check constraints.
    if use_int4_w4a16 or use_int4_w4a8:
        assert hidden_states.shape[1] // 2 == w1.shape[
            2], "Hidden size mismatch"
    else:
        assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch"

    assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
    assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
    assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
    assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
    assert hidden_states.dtype in [
        torch.float32, torch.float16, torch.bfloat16, torch.int8,  torch.float8_e4m3fn
    ]

    num_tokens, _ = hidden_states.shape
    E, N, _ = w1.shape
    if global_num_experts == -1:
        global_num_experts = E
    top_k_num = topk_ids.shape[1]
    # We execute the fused_moe kernel in chunks to circumvent this issue:
    # https://github.com/vllm-project/vllm/issues/5938
    CHUNK_SIZE = int(os.environ.get("TRITON_FUSED_MOE_CHUNK_SIZE", "16384"))
    M = min(num_tokens, CHUNK_SIZE)

    moe_config_func = get_optimal_moe_config_func(
                        hidden_states, w1, topk_ids,
                        use_int8_w8a16=use_int8_w8a16,
                        use_int8_w8a8=use_int8_w8a8,
                        use_fp8_w8a8=use_fp8_w8a8,
                        use_int4_w4a16=use_int4_w4a16,
                        use_int4_w4a8=use_int4_w4a8,
                        use_mxfp4_w4a4=False, #always false in wna16
                        block_shape=block_shape,
                        is_bottom=False)
    moe_config_func2 = get_optimal_moe_config_func(
                        hidden_states, w2, topk_ids,
                        use_int8_w8a16=use_int8_w8a16,
                        use_int8_w8a8=use_int8_w8a8,
                        use_fp8_w8a8=use_fp8_w8a8,
                        use_int4_w4a16=use_int4_w4a16,
                        use_int4_w4a8=use_int4_w4a8,
                        use_mxfp4_w4a4=False, #always false in wna16
                        block_shape=block_shape,
                        is_bottom=True)
    config = moe_config_func(M)
    config2, max_block_m = moe_config_func2(M)

    # config["COMBINE_SCALE_LOAD"] = False
    # config2["COMBINE_SCALE_LOAD"] = False
    # config["USE_MLS_LOAD"] = True
    # config2["USE_MLS_LOAD"] = True

    bottom_moe_a_use_mls = (
        _bottom_moe_use_mls()
        and not use_int4_w4a8
        and config2 is not None
        and config2.get("USE_MLS_LOAD", False))
    # bottom_moe_a_use_mls = False

    topk = top_k_num
    max_padded_tokens = (
        min(M * topk, E + 1) * (max_block_m - 1) if bottom_moe_a_use_mls else 0
    )
    total_tokens = M * topk + max_padded_tokens

    # We can reuse the memory between these because by the time we need
    # cache3, we're done with cache1
    if expert_map is not None:
        cache13 = torch.zeros(total_tokens * max(N, w2.shape[1]),
                            device=hidden_states.device,
                            dtype=odtype)
    else:
        cache13 = torch.empty(total_tokens * max(N, w2.shape[1]),
                            device=hidden_states.device,
                            dtype=odtype)
    intermediate_cache3 = cache13[:M * topk * w2.shape[1]].view(
        (M, topk, w2.shape[1]))

    if hidden_states.dtype == torch.bfloat16:
        compute_type = tl.bfloat16
    elif hidden_states.dtype == torch.float16:
        compute_type = tl.float16
    elif hidden_states.dtype == torch.float32:
        compute_type = tl.float32
    elif hidden_states.dtype == torch.int8 or hidden_states.dtype == torch.float8_e4m3fn:
        if odtype == torch.bfloat16:
            compute_type = tl.bfloat16
        else:
            compute_type = tl.float16
    else:
        raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")

    if inplace:
        out_hidden_states = hidden_states
    else:
        out_hidden_states = torch.empty(hidden_states.shape, device=hidden_states.device, dtype=odtype)

    for chunk in range((num_tokens // CHUNK_SIZE) + 1):
        begin_chunk_idx, end_chunk_idx = (chunk * CHUNK_SIZE,
                                          min((chunk + 1) * CHUNK_SIZE,
                                              num_tokens))
        curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
        tokens_in_chunk, _ = curr_hidden_states.shape

        if tokens_in_chunk == 0:
            break

        if tokens_in_chunk < CHUNK_SIZE and chunk > 0:
            # Adjust the intermediate cache size and config for the last
            # chunk. Note that in most cases we only have one chunk
            # so the cache size and config are already set correctly and
            # do not need to be adjusted.
            config = moe_config_func(tokens_in_chunk)
            config2, max_block_m = moe_config_func2(tokens_in_chunk)
            bottom_moe_a_use_mls = (
                _bottom_moe_use_mls()
                and config2 is not None
                and config2.pop("USE_MLS_LOAD", False)
                and (block_shape is not None and (use_int8_w8a8 or use_fp8_w8a8)))
            intermediate_cache3 = intermediate_cache3[:tokens_in_chunk]

        padded_tokens = (
            min(tokens_in_chunk * topk, E + 1) * (config["BLOCK_SIZE_M"] - 1)
            if bottom_moe_a_use_mls
            else 0
        )
        total_tokens = tokens_in_chunk * topk + padded_tokens
        intermediate_cache1 = cache13[: total_tokens * N].view(
            (total_tokens, N),
        )
        if expert_map is not None:
            intermediate_cache2 = torch.zeros(
                (total_tokens, N // 2 if activation != "relu2" else N),
                device=hidden_states.device,
                dtype=odtype)
        else:
            intermediate_cache2 = torch.empty(
                (total_tokens, N // 2 if activation != "relu2" else N),
                device=hidden_states.device,
                dtype=odtype)

        curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
        curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]

        ck_sorting = expert_map is None # TODO: check why expert_map failed issue.
        sorted_weights = None
        if not ck_sorting:
            sorted_token_ids, expert_ids, num_tokens_post_padded = (
                moe_align_block_size(curr_topk_ids, config['BLOCK_SIZE_M'],
                                    global_num_experts, expert_map))
        else:
            sorted_token_ids, sorted_weights, expert_ids, num_tokens_post_padded, \
                _tokens_positions_per_expert, _moe_buf = (
                    moe_sorting_ck(curr_topk_ids, curr_topk_weights, global_num_experts,
                                w2.shape[1], odtype, config["BLOCK_SIZE_M"], expert_map)
            )

        if (use_int8_w8a8 or use_fp8_w8a8 or use_int4_w4a8) and per_channel_quant:
            quant_dtype = torch.float8_e4m3fn if use_fp8_w8a8 else torch.int8
            if curr_hidden_states.dtype == torch.float16 or curr_hidden_states.dtype==torch.bfloat16:
                input_q,input_scale = per_token_quant_hip(curr_hidden_states,quant_dtype=quant_dtype)
            else:
                input_q,input_scale = curr_hidden_states,a1_scale
            invoke_fused_moe_kernel(input_q,
                                    w1,
                                    intermediate_cache1,
                                    input_scale,
                                    w1_scale,
                                    w1_zp,
                                    curr_topk_weights,
                                    curr_topk_ids,
                                    sorted_token_ids,
                                    sorted_weights,
                                    expert_ids,
                                    num_tokens_post_padded,
                                    False,
                                    top_k_num,
                                    compute_type=compute_type,
                                    use_fp8_w8a8=use_fp8_w8a8,
                                    use_int8_w8a8=use_int8_w8a8,
                                    use_int8_w8a16=use_int8_w8a16,
                                    use_int4_w4a16=use_int4_w4a16,
                                    use_int4_w4a8=use_int4_w4a8,
                                    per_channel_quant=per_channel_quant,
                                    block_shape=block_shape,
                                    c_sorted=bottom_moe_a_use_mls,
                                    ck_sorting=ck_sorting,
                                    ck_topk=top_k_num,
                                    config=config)
        elif block_shape is not None and (use_int8_w8a8 or use_int4_w4a8 or use_fp8_w8a8):
            quant_dtype = torch.float8_e4m3fn if use_fp8_w8a8 else torch.int8
            if curr_hidden_states.dtype == torch.float16 or curr_hidden_states.dtype==torch.bfloat16:
                input_q, input_scale = per_block_quant_wrapper((1,block_shape[1]))(per_token_quant_hip)(curr_hidden_states,quant_dtype=quant_dtype)
            else:
                input_q, input_scale = curr_hidden_states,a1_scale
            invoke_fused_moe_kernel(input_q,
                                    w1,
                                    intermediate_cache1,
                                    input_scale,
                                    w1_scale,
                                    w1_zp,
                                    curr_topk_weights,
                                    curr_topk_ids,
                                    sorted_token_ids,
                                    sorted_weights,
                                    expert_ids,
                                    num_tokens_post_padded,
                                    False,
                                    top_k_num,
                                    compute_type=compute_type,
                                    use_fp8_w8a8=use_fp8_w8a8,
                                    use_int8_w8a8=use_int8_w8a8,
                                    use_int8_w8a16=use_int8_w8a16,
                                    use_int4_w4a16=use_int4_w4a16,
                                    use_int4_w4a8=use_int4_w4a8,
                                    per_channel_quant=per_channel_quant,
                                    block_shape=block_shape,
                                    c_sorted=bottom_moe_a_use_mls,
                                    ck_sorting=ck_sorting,
                                    ck_topk=top_k_num,
                                    config=config)
        else:
            invoke_fused_moe_kernel(curr_hidden_states,
                                    w1,
                                    intermediate_cache1,
                                    a1_scale,
                                    w1_scale,
                                    w1_zp,
                                    curr_topk_weights,
                                    curr_topk_ids,
                                    sorted_token_ids,
                                    sorted_weights,
                                    expert_ids,
                                    num_tokens_post_padded,
                                    False,
                                    top_k_num,
                                    compute_type=compute_type,
                                    use_fp8_w8a8=use_fp8_w8a8,
                                    use_int8_w8a8=use_int8_w8a8,
                                    use_int8_w8a16=use_int8_w8a16,
                                    use_int4_w4a16=use_int4_w4a16,
                                    use_int4_w4a8=use_int4_w4a8,
                                    per_channel_quant=per_channel_quant,
                                    block_shape=block_shape,
                                    c_sorted=bottom_moe_a_use_mls,
                                    ck_sorting=ck_sorting,
                                    ck_topk=top_k_num,
                                    config=config)
        if activation == "silu":
            triton_silu_and_mul(intermediate_cache2,
                                intermediate_cache1.view(-1, N))
            # torch.ops._C.silu_and_mul(intermediate_cache2,
            #                     intermediate_cache1.view(-1, N))
        elif activation == "gelu":
            triton_gelu_and_mul(intermediate_cache2,
                                intermediate_cache1.view(-1, N))
        elif activation == "relu2":
            triton_relu2(
                intermediate_cache2,
                intermediate_cache1.view(-1, N),
            )
        else:
            raise ValueError(f"Unsupported FusedMoe activation: {activation}")
        if expert_map != None:
            # for EP mode, intermediate_cache1 and intermediate_cache3 need be zeros inited
            # since intermediate_cache1 and intermediate_cache3 shared same buffer,
            # to make sure intermediate_cache3 is zeros inited,
            # intermediate_cache1 need inited to zeros after silu_and_mul
            intermediate_cache1.fill_(0)

        if (use_int8_w8a8 or use_fp8_w8a8 or use_int4_w4a8) and per_channel_quant:
            quant_dtype = torch.float8_e4m3fn if use_fp8_w8a8 else torch.int8
            bridge_q, bridge_scale = per_token_quant_hip(intermediate_cache2, quant_dtype=quant_dtype)
            invoke_fused_moe_kernel(bridge_q,
                                    w2,
                                    intermediate_cache3,
                                    bridge_scale,
                                    w2_scale,
                                    w2_zp,
                                    curr_topk_weights,
                                    curr_topk_ids,
                                    sorted_token_ids,
                                    sorted_weights,
                                    expert_ids,
                                    num_tokens_post_padded,
                                    True,
                                    1,
                                    compute_type=compute_type,
                                    use_fp8_w8a8=use_fp8_w8a8,
                                    use_int8_w8a8=use_int8_w8a8,
                                    use_int8_w8a16=use_int8_w8a16,
                                    use_int4_w4a16=use_int4_w4a16,
                                    use_int4_w4a8=use_int4_w4a8,
                                    per_channel_quant=per_channel_quant,
                                    block_shape=block_shape,
                                    bottom_a_use_mls_load=bottom_moe_a_use_mls,
                                    ck_sorting=ck_sorting,
                                    ck_topk=top_k_num,
                                    config=config2)
        elif block_shape is not None and (use_int8_w8a8 or use_int4_w4a8 or use_fp8_w8a8):
            quant_dtype = torch.float8_e4m3fn if use_fp8_w8a8 else torch.int8
            bridge_q, bridge_scale = per_block_quant_wrapper((1,block_shape[1]))(per_token_quant_hip)(intermediate_cache2, quant_dtype=quant_dtype)
            invoke_fused_moe_kernel(bridge_q,
                                    w2,
                                    intermediate_cache3,
                                    bridge_scale,
                                    w2_scale,
                                    w2_zp,
                                    curr_topk_weights,
                                    curr_topk_ids,
                                    sorted_token_ids,
                                    sorted_weights,
                                    expert_ids,
                                    num_tokens_post_padded,
                                    True,
                                    1,
                                    compute_type=compute_type,
                                    use_fp8_w8a8=use_fp8_w8a8,
                                    use_int8_w8a8=use_int8_w8a8,
                                    use_int8_w8a16=use_int8_w8a16,
                                    use_int4_w4a16=use_int4_w4a16,
                                    use_int4_w4a8=use_int4_w4a8,
                                    per_channel_quant=per_channel_quant,
                                    block_shape=block_shape,
                                    bottom_a_use_mls_load=bottom_moe_a_use_mls,
                                    ck_sorting=ck_sorting,
                                    ck_topk=top_k_num,
                                    config=config2)
        else:
            invoke_fused_moe_kernel(intermediate_cache2,
                        w2,
                        intermediate_cache3,
                        a2_scale,
                        w2_scale,
                        w2_zp,
                        curr_topk_weights,
                        curr_topk_ids,
                        sorted_token_ids,
                        sorted_weights,
                        expert_ids,
                        num_tokens_post_padded,
                        True,
                        1,
                        compute_type=compute_type,
                        use_fp8_w8a8=use_fp8_w8a8,
                        use_int8_w8a8=use_int8_w8a8,
                        use_int8_w8a16=use_int8_w8a16,
                        use_int4_w4a16=use_int4_w4a16,
                        use_int4_w4a8=use_int4_w4a8,
                        per_channel_quant=per_channel_quant,
                        block_shape=block_shape,
                        bottom_a_use_mls_load=bottom_moe_a_use_mls,
                        ck_sorting=ck_sorting,
                        ck_topk=top_k_num,
                        config=config2)

        mode_use_triton_moe_sum = out_hidden_states.dtype == torch.float16 or  \
                                  out_hidden_states.dtype == torch.bfloat16 or \
                                  out_hidden_states.dtype == torch.float32
        if mode_use_triton_moe_sum:
            triton_moe_sum(intermediate_cache3.view(*intermediate_cache3.shape),
                           out_hidden_states[begin_chunk_idx:end_chunk_idx],
                           routed_scaling_factor=routed_scaling_factor)
        else:
            moe_sum(intermediate_cache3.view(*intermediate_cache3.shape),
                        out_hidden_states[begin_chunk_idx:end_chunk_idx])
            if routed_scaling_factor != 1.0:
                out_hidden_states[begin_chunk_idx:end_chunk_idx].mul_(
                    routed_scaling_factor)
        if end_chunk_idx < num_tokens and expert_map != None:
            # if has next chunk, intermediate_cache3 need init to zeros
            intermediate_cache3.fill_(0)
    return out_hidden_states
