import pytest
import torch
import itertools
import math

from aiter import dtypes
from op_tests.utility.scalar_type import ScalarType, scalar_types
from op_tests.utility.utils import quantize_weights
from op_tests.utility.utils import torch_moe as torch_score_moe
#from aiter.ops.triton.fused_moe import fused_moe as triton_score_fused_moe
from aiter.ops.triton.fused_moe import fused_experts as triton_fused_moe
from aiter.fused_moe import fused_topk, torch_moe, torch_moe_stage1, torch_moe_stage2
from aiter import ActivationType
from aiter.test_common import checkAllclose, perftest

import functools
import os
from typing import Any, Callable, Dict, List, Optional, Tuple

import triton
import triton.language as tl
import aiter
from aiter.ops.triton.fused_moe import triton_moe_sum
import vllm.envs as envs
from vllm import _custom_ops as ops
from vllm.logger import init_logger
from vllm.model_executor.layers.fused_moe.utils import (
    _resize_cache, moe_kernel_quantize_input, per_token_group_quant_fp8)
from vllm.model_executor.layers.fused_moe.config import (
    FusedMoEQuantConfig, get_config_quant_dtype)
from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
    moe_align_block_size)
from vllm.model_executor.layers.fused_moe import get_config_file_name

logger = init_logger(__name__)

def moe_sorting_ck(
    topk_ids,
    topk_weights,
    num_experts,
    model_dim,
    moebuf_dtype,
    block_size=16,
    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)

    aiter.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

#basic define
###################################################################################################################
def get_config_dtype_str(
        dtype: torch.dtype,
        use_int4_w4a16: Optional[bool] = False,
        use_int8_w8a16: Optional[bool] = False,
        use_fp8_w8a8: Optional[bool] = False,
        use_mxfp4_w4a4: Optional[bool] = False) -> Optional[str]:
    if use_fp8_w8a8:
        return "fp8_w8a8"
    elif use_int8_w8a16:
        return "int8_w8a16"
    elif use_int4_w4a16:
        return "int4_w4a16"
    elif use_mxfp4_w4a4:
        return "mxfp4_w4a4"
    elif dtype == torch.float:
        # avoiding cases where kernel fails when float32 MoE
        # use fp16/bfloat16 configs
        return "float32"
    return None

def try_get_optimal_moe_config(
    w1_shape: tuple[int, ...],
    w2_shape: tuple[int, ...],
    top_k: int,
    dtype: Optional[str],
    M: int,
    is_marlin: bool = False,
    block_shape: Optional[list[int]] = None,
) -> dict[str, int]:
    from vllm.model_executor.layers.fused_moe import get_config
    override_config = get_config()
    if override_config:
        config = override_config
    else:
        # First try to load optimal config from the file
        E, _, N = w2_shape
        if dtype == "int4_w4a16":
            N = N * 2
        block_n = block_shape[0] if block_shape else 0
        block_k = block_shape[1] if block_shape else 0
        configs = get_moe_configs(E, N, dtype, block_n, block_k)

        if configs:
            # If an optimal configuration map has been found, look up the
            # optimal config
            config = configs[min(configs.keys(), key=lambda x: abs(x - M))]
        else:
            # Else use the default config
            config = get_default_config(M, E, N, w1_shape[2], top_k, dtype,
                                        is_marlin, block_shape)
    return config

# Adapted from: https://github.com/sgl-project/sglang/pull/2628
@functools.lru_cache
def get_moe_configs(
    E: int,
    N: int,
    dtype: Optional[str],
    block_n: Optional[int] = None,
    block_k: Optional[int] = None,
) -> Optional[Dict[int, Any]]:
    """
    Return optimized configurations for the fused MoE kernel.

    The return value will be a dictionary that maps an irregular grid of
    batch sizes to configurations of the fused_moe kernel. To evaluate the
    kernel on a given batch size bs, the closest batch size in the grid should
    be picked and the associated configuration chosen to invoke the kernel.
    """

    # First look up if an optimized configuration is available in the configs
    # directory
    block_shape = [block_n, block_k] if block_n and block_k else None
    json_file_name = get_config_file_name(E, N, dtype, block_shape)

    config_file_path = os.path.join(
        os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name)
    if os.path.exists(config_file_path):
        with open(config_file_path) as f:
            logger.info("Using configuration from %s for MoE layer.",
                        config_file_path)
            # If a configuration has been found, return it
            return {int(key): val for key, val in json.load(f).items()}

    # If no optimized configuration is available, we will use the default
    # configuration
    logger.warning(
        ("Using default MoE config. Performance might be sub-optimal! "
         "Config file not found at %s"), config_file_path)
    return None

def get_moe_wna16_block_config(config: Dict[str,
                                            int], use_moe_wna16_cuda: bool,
                               num_valid_tokens: int, size_k: int, size_n: int,
                               num_experts: int, group_size: int,
                               real_top_k: int, block_size_m: int):
    if "BLOCK_SIZE_N" in config and "BLOCK_SIZE_K" in config:
        # optimal block config is set
        return {}
    if not use_moe_wna16_cuda:
        # triton moe wna16 kernel
        if num_valid_tokens // real_top_k == 1:
            # if bs=1, use a smaller BLOCK_SIZE_N
            return {"BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64}
        else:
            return {"BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}
    else:
        # cuda moe wna16 kernel
        # set default block_size 128, and increase them when num_blocks
        # is too large.
        block_size_n = 128
        block_size_k = 128
        if block_size_k <= group_size:
            block_size_k = group_size

        num_n_blocks = size_k // block_size_k
        num_k_blocks = size_n // block_size_k
        num_m_blocks = (num_valid_tokens + block_size_m - 1) / block_size_m + \
            num_experts
        if num_valid_tokens // real_top_k <= block_size_m:
            num_m_blocks = min(num_m_blocks, num_valid_tokens)
        num_blocks = num_m_blocks * num_n_blocks * num_k_blocks

        if size_k % 256 == 0 and num_blocks >= 256 and \
                block_size_k < 256:
            block_size_k = 256
            num_blocks = num_blocks // (256 // block_size_k)

        if num_m_blocks <= 16 and size_k % (block_size_k * 2) == 0 and \
                size_k % (block_size_k * 2) == 0 and block_size_k <= 512 and \
                num_blocks >= 512:
            block_size_k = block_size_k * 2
            num_blocks = num_blocks // 2

        if num_blocks > 1024:
            block_size_n = 256
            num_n_blocks = num_n_blocks // 2
            num_blocks = num_blocks // 2

        if size_n <= 1024 and num_blocks >= 1024:
            # The kernel performance got much better with BLOCK_SIZE_N=1024
            # when num_blocks is large, event when N is small.
            # Not sure why, maybe it force the CUDA SM process only one block
            # at the same time.
            block_size_n = 1024

        return {"BLOCK_SIZE_N": block_size_n, "BLOCK_SIZE_K": block_size_k}

def should_moe_wna16_use_cuda(num_valid_tokens: int, group_size: int,
                              num_experts: int, bit: int):
    # return bit == 4 and group_size in [32, 64, 128] and \
    #     num_valid_tokens / num_experts <= 6
    return False

def get_default_config(
    M: int,
    E: int,
    N: int,
    K: int,
    topk: int,
    dtype: Optional[str],
    is_marlin: bool,
    block_shape: Optional[List[int]] = None,
) -> Dict[str, int]:
    if dtype == "fp8_w8a8" and block_shape is not None:
        # Block-wise quant: BLOCK_SIZE_N must be divisible by block_shape[0]
        # BLOCK_SIZE_K must be divisible by block_shape[1]
        # num_stages=3 can cause triton.runtime.errors.OutOfResources
        # on ROCm, set it to 2 instead.
        config = {
            "BLOCK_SIZE_M": 64,
            "BLOCK_SIZE_N": block_shape[0],
            "BLOCK_SIZE_K": block_shape[1],
            "GROUP_SIZE_M": 32,
            "num_warps": 4,
            "num_stages": 3
        }
    elif dtype in ["int4_w4a16", "int8_w8a16"] and block_shape is not None:
        # moe wna16 kernels
        # only set BLOCK_SIZE_M
        # BLOCK_SIZE_N and BLOCK_SIZE_K would be set later
        bit = 4 if dtype == "int4_w4a16" else 8
        use_moe_wna16_cuda = should_moe_wna16_use_cuda(M * topk,
                                                       block_shape[1], E, bit)
        if use_moe_wna16_cuda:
            config = {"BLOCK_SIZE_M": min(16, M)}
        elif M <= 20:
            config = {"BLOCK_SIZE_M": 16, "GROUP_SIZE_M": 1}
        elif M <= 40:
            config = {"BLOCK_SIZE_M": 32, "GROUP_SIZE_M": 1}
        else:
            config = {"BLOCK_SIZE_M": 64, "GROUP_SIZE_M": 1}
    elif is_marlin:
        for block_size_m in [8, 16, 32, 48, 64]:
            if M * topk / E / block_size_m < 0.9:
                break
        return {"BLOCK_SIZE_M": block_size_m}
    elif M <= E:
        config = {
            "BLOCK_SIZE_M": 16,
            "BLOCK_SIZE_N": 32,
            "BLOCK_SIZE_K": 64,
            "GROUP_SIZE_M": 1,
        }
    else:
        config = {
            "BLOCK_SIZE_M": 64,
            "BLOCK_SIZE_N": 64,
            "BLOCK_SIZE_K": 32,
            "GROUP_SIZE_M": 8,
        }
    return config

def try_get_optimal_moe_config(
    w1_shape: Tuple[int, ...],
    w2_shape: Tuple[int, ...],
    top_k: int,
    dtype: Optional[str],
    M: int,
    is_marlin: bool = False,
    block_shape: Optional[List[int]] = None,
    is_bottom: bool = False,
):
    # from vllm.model_executor.layers.fused_moe import get_config           #vllm中解决get_config引入的rocm_aiter_asm_moe_tkw1多重注册问题后，使能此处
    # override_config = get_config()
    override_config=None
    
    if override_config:
        config = override_config
    else:
        # First try to load optimal config from the file
        E, _, N = w2_shape
        if dtype == "int4_w4a16":
            N = N * 2
        block_n = block_shape[0] if block_shape else 0
        block_k = block_shape[1] if block_shape else 0
        configs = get_moe_configs(E, N, dtype, block_n, block_k)

        if configs:
            # If an optimal configuration map has been found, look up the
            # optimal config
            config = configs[min(configs.keys(), key=lambda x: abs(x - M))]
        else:
            # Else use the default config
            config = get_default_config(M, E, N, w1_shape[2], top_k, dtype,
                                        is_marlin, block_shape)
    return config

#triton func
###################################################################################################################
@triton.jit
def write_zeros_to_output(c_ptr, stride_cm, stride_cn, pid_n, N, offs_token,
                          token_mask, BLOCK_SIZE_M, BLOCK_SIZE_N,
                          compute_type):
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=compute_type)
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[
        None, :]
    c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
    tl.store(c_ptrs, accumulator, mask=c_mask)

@triton.jit
def fused_moe_kernel_gptq_awq(
        # Pointers to matrices
        a_ptr,
        b_ptr,
        c_ptr,
        b_scale_ptr,
        b_zp_ptr,
        topk_weights_ptr,
        sorted_token_ids_ptr,
        expert_ids_ptr,
        num_tokens_post_padded_ptr,
        # Matrix dimensions
        N: tl.constexpr,
        K: tl.constexpr,
        EM,
        num_valid_tokens,
        # The stride variables represent how much to increase the ptr by when
        # moving by 1 element in a particular dimension. E.g. `stride_am` is
        # how much to increase `a_ptr` by to get the element one row down
        # (A has M rows).
        stride_am,
        stride_ak,
        stride_be,
        stride_bk,
        stride_bn,
        stride_cm,
        stride_cn,
        stride_bse,
        stride_bsk,
        stride_bsn,
        stride_bze,
        stride_bzk,
        stride_bzn,
        block_k_diviable: tl.constexpr,
        group_size: tl.constexpr,
        # Meta-parameters
        BLOCK_SIZE_M: tl.constexpr,
        BLOCK_SIZE_N: tl.constexpr,
        BLOCK_SIZE_K: tl.constexpr,
        GROUP_SIZE_M: tl.constexpr,
        MUL_ROUTED_WEIGHT: tl.constexpr,
        top_k: tl.constexpr,
        compute_type: tl.constexpr,
        has_zp: tl.constexpr,
        use_int4_w4a16: tl.constexpr,
        use_int8_w8a16: tl.constexpr):
    """
    Implements the fused computation for a Mixture of Experts (MOE) using
    token and expert matrices.

    Key Parameters:
    - A: The input tensor representing tokens with shape (*, K), where '*' can
        be any shape representing batches and K is the feature dimension of
        each token.
    - B: The stacked MOE weight tensor with shape (E, N, K), where E is
        the number of experts, K is the input feature dimension, and N is
        the output feature dimension.
    - C: The output cache tensor with shape (M, topk, N), where M is the
        total number of tokens post padding, topk is the number of times
        each token is repeated, and N is the output feature dimension.
    - sorted_token_ids: A tensor containing the sorted indices of tokens,
        repeated topk times and arranged by the expert index they are
        assigned to.
    - expert_ids: A tensor containing the indices of the expert for each
        block. It determines which expert matrix from B should be used for
        each block in A.
    This kernel performs the multiplication of a token by its corresponding
    expert matrix as determined by `expert_ids`. The sorting of
    `sorted_token_ids` by expert index and padding ensures divisibility by
    BLOCK_SIZE_M, which is necessary to maintain consistency in block matrix
    multiplication across different blocks processed by the same expert.
    """
    # -----------------------------------------------------------
    # Map program ids `pid` to the block of C it should compute.
    # This is done in a grouped ordering to promote L2 data reuse.
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    # ----------------------------------------------------------
    # Create pointers for the first blocks of A and B.
    # We will advance this pointer as we move in the K direction
    # and accumulate
    # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers
    # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers
    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
    if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
        return
    offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(
        tl.int64)
    offs_token = tl.load(sorted_token_ids_ptr + offs_token_id)
    token_mask = offs_token < num_valid_tokens

    off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
    if off_experts == -1:
        # -----------------------------------------------------------
        # Write back zeros to the output when the expert is not
        # in the current expert parallel rank.
        write_zeros_to_output(c_ptr, stride_cm, stride_cn, pid_n, N,
                              offs_token, token_mask, BLOCK_SIZE_M,
                              BLOCK_SIZE_N, compute_type)
        return

    offs_bn = (pid_n * BLOCK_SIZE_N +
               tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = a_ptr + (offs_token[:, None] // top_k * stride_am +
                      offs_k[None, :] * stride_ak)

    if use_int4_w4a16:
        b_ptrs = b_ptr + off_experts * stride_be + \
            (offs_k[:, None] // 2) * stride_bk + offs_bn[None, :] * \
                stride_bn
        b_shifter = (offs_k[:, None] % 2) * 4
    elif use_int8_w8a16:
        b_ptrs = b_ptr + off_experts * stride_be + \
            offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn

    if not has_zp and use_int4_w4a16:
        b_zp_num = 8
    if not has_zp and use_int8_w8a16:
        b_zp_num = 128
    elif has_zp and use_int4_w4a16:
        b_zp_shifter = (offs_bn[None, :] % 2) * 4

    # -----------------------------------------------------------
    # Iterate to compute a block of the C matrix.
    # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block
    # of fp32 values for higher accuracy.
    # `accumulator` will be converted back to fp16 after the loop.
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        # Load the next block of A and B, generate a mask by checking the
        # K dimension.

        if not block_k_diviable:
            k_mask = offs_k[:, None] < K - k * BLOCK_SIZE_K
            k_other = 0.0
        else:
            k_mask = None
            k_other = None

        a = tl.load(a_ptrs,
                    mask=token_mask[:, None] &
                    (offs_k[None, :] < K - k * BLOCK_SIZE_K),
                    other=0.0)
        b = tl.load(b_ptrs)
        if use_int4_w4a16:
            b = (b >> b_shifter) & 0xF

        b_scale_ptrs = b_scale_ptr + off_experts * stride_bse + \
            offs_bn[None, :] * stride_bsn + \
            ((offs_k[:, None] + BLOCK_SIZE_K * k) // group_size) * \
                stride_bsk
        b_scale = tl.load(b_scale_ptrs, mask=k_mask, other=k_other)
        b_scale = b_scale.to(tl.float32)

        if has_zp and use_int4_w4a16:
            offs_k_true = (offs_k[:, None] + BLOCK_SIZE_K * k) // group_size
            b_zp_ptrs = b_zp_ptr + off_experts * stride_bze + \
                (offs_bn[None, :] // 2) * stride_bzn + \
                offs_k_true * stride_bzk
            b_zp = tl.load(b_zp_ptrs, mask=k_mask, other=k_other)
            b_zp = ((b_zp >> b_zp_shifter) & 0xF)
            b_zp = b_zp.to(tl.float32)
        elif has_zp and use_int8_w8a16:
            offs_k_true = (offs_k[:, None] + BLOCK_SIZE_K * k) // group_size
            b_zp_ptrs = b_zp_ptr + off_experts * stride_bze + \
                offs_bn[None, :] * stride_bzn + \
                offs_k_true * stride_bzk
            b_zp = tl.load(b_zp_ptrs, mask=k_mask, other=k_other)
            b_zp = b_zp.to(tl.float32)

        # We accumulate along the K dimension.
        if has_zp:
            b = ((b.to(tl.float32) - b_zp) * b_scale).to(compute_type)
        else:
            b = ((b.to(tl.float32) - b_zp_num) * b_scale).to(compute_type)
        accumulator = tl.dot(a, b, acc=accumulator)

        # Advance the ptrs to the next K block.
        a_ptrs += BLOCK_SIZE_K * stride_ak
        if use_int4_w4a16:
            b_ptrs += (BLOCK_SIZE_K // 2) * stride_bk
        else:
            b_ptrs += BLOCK_SIZE_K * stride_bk

    if MUL_ROUTED_WEIGHT:
        moe_weight = tl.load(topk_weights_ptr + offs_token,
                             mask=token_mask,
                             other=0)
        accumulator = accumulator * moe_weight[:, None]

    accumulator = accumulator.to(compute_type)
    # -----------------------------------------------------------
    # Write back the block of the output
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[
        None, :]
    c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
    tl.store(c_ptrs, accumulator, mask=c_mask)

@triton.jit
def fused_moe_kernel(
    # Pointers to matrices
    a_ptr,
    b_ptr,
    c_ptr,
    a_scale_ptr,
    b_scale_ptr,
    topk_weights_ptr,
    sorted_token_ids_ptr,
    expert_ids_ptr,
    num_tokens_post_padded_ptr,
    # Matrix dimensions
    N,
    K,
    EM,
    num_valid_tokens,
    # The stride variables represent how much to increase the ptr by when
    # moving by 1 element in a particular dimension. E.g. `stride_am` is
    # how much to increase `a_ptr` by to get the element one row down
    # (A has M rows).
    stride_am,
    stride_ak,
    stride_be,
    stride_bk,
    stride_bn,
    stride_cm,
    stride_cn,
    stride_asm,
    stride_ask,
    stride_bse,
    stride_bsk,
    stride_bsn,
    # Block size for block-wise quantization
    group_n: tl.constexpr,
    group_k: tl.constexpr,
    # Meta-parameters
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
    MUL_ROUTED_WEIGHT: tl.constexpr,
    top_k: tl.constexpr,
    compute_type: tl.constexpr,
    use_fp8_w8a8: tl.constexpr,
    use_int8_w8a8: tl.constexpr,
    use_int8_w8a16: tl.constexpr,
    per_channel_quant: tl.constexpr,
):
    """
    Implements the fused computation for a Mixture of Experts (MOE) using
    token and expert matrices.

    Key Parameters:
    - A: The input tensor representing tokens with shape (*, K), where '*' can
        be any shape representing batches and K is the feature dimension of
        each token.
    - B: The stacked MOE weight tensor with shape (E, N, K), where E is
        the number of experts, K is the input feature dimension, and N is
        the output feature dimension.
    - C: The output cache tensor with shape (M, topk, N), where M is the
        total number of tokens post padding, topk is the number of times
        each token is repeated, and N is the output feature dimension.
    - sorted_token_ids: A tensor containing the sorted indices of tokens,
        repeated topk times and arranged by the expert index they are
        assigned to.
    - expert_ids: A tensor containing the indices of the expert for each
        block. It determines which expert matrix from B should be used for
        each block in A.
    This kernel performs the multiplication of a token by its corresponding
    expert matrix as determined by `expert_ids`. The sorting of
    `sorted_token_ids` by expert index and padding ensures divisibility by
    BLOCK_SIZE_M, which is necessary to maintain consistency in block matrix
    multiplication across different blocks processed by the same expert.
    """
    # -----------------------------------------------------------
    # Map program ids `pid` to the block of C it should compute.
    # This is done in a grouped ordering to promote L2 data reuse.
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    # ----------------------------------------------------------
    # Create pointers for the first blocks of A and B.
    # We will advance this pointer as we move in the K direction
    # and accumulate
    # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers
    # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers
    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
    if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
        return
    offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(
        tl.int64)
    offs_token = tl.load(sorted_token_ids_ptr + offs_token_id)
    token_mask = offs_token < num_valid_tokens

    off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
    if off_experts == -1:
        # -----------------------------------------------------------
        # Write back zeros to the output when the expert is not
        # in the current expert parallel rank.
        write_zeros_to_output(c_ptr, stride_cm, stride_cn, pid_n, N,
                              offs_token, token_mask, BLOCK_SIZE_M,
                              BLOCK_SIZE_N, compute_type)
        return

    offs_bn = (pid_n * BLOCK_SIZE_N +
               tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = a_ptr + (offs_token[:, None] // top_k * stride_am +
                      offs_k[None, :] * stride_ak)

    b_ptrs = b_ptr + off_experts * stride_be + (offs_k[:, None] * stride_bk +
                                                offs_bn[None, :] * stride_bn)
    if use_int8_w8a16:
        b_scale_ptrs = b_scale_ptr + off_experts * stride_bse + offs_bn[
            None, :] * stride_bsn
        b_scale = tl.load(b_scale_ptrs)

    if use_fp8_w8a8 or use_int8_w8a8:
        # block-wise
        if group_k > 0 and group_n > 0:
            a_scale_ptrs = a_scale_ptr + (offs_token // top_k) * stride_asm
            offs_bsn = offs_bn // group_n
            b_scale_ptrs = (b_scale_ptr + off_experts * stride_bse +
                            offs_bsn * stride_bsn)
        # channel-wise
        elif per_channel_quant:
            b_scale_ptrs = b_scale_ptr + off_experts * stride_bse + offs_bn[
                None, :] * stride_bsn
            b_scale = tl.load(b_scale_ptrs)
            # Load per-token scale for activations
            a_scale_ptrs = a_scale_ptr + (offs_token // top_k) * stride_asm
            a_scale = tl.load(a_scale_ptrs, mask=token_mask, other=0.0)[:,
                                                                        None]
        # tensor-wise
        else:
            a_scale = tl.load(a_scale_ptr)
            b_scale = tl.load(b_scale_ptr + off_experts)

    # -----------------------------------------------------------
    # Iterate to compute a block of the C matrix.
    # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block
    # of fp32 values for higher accuracy.
    # `accumulator` will be converted back to fp16 after the loop.
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        # Load the next block of A and B, generate a mask by checking the
        # K dimension.
        a = tl.load(a_ptrs,
                    mask=token_mask[:, None] &
                    (offs_k[None, :] < K - k * BLOCK_SIZE_K),
                    other=0.0)
        b = tl.load(b_ptrs,
                    mask=offs_k[:, None] < K - k * BLOCK_SIZE_K,
                    other=0.0)
        # We accumulate along the K dimension.
        if use_int8_w8a16:
            accumulator = tl.dot(a, b.to(compute_type), acc=accumulator)
        elif use_fp8_w8a8 or use_int8_w8a8:
            if group_k > 0 and group_n > 0:
                k_start = k * BLOCK_SIZE_K
                offs_ks = k_start // group_k
                a_scale = tl.load(a_scale_ptrs + offs_ks * stride_ask,
                                  mask=token_mask,
                                  other=0.0)
                b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk)

                accumulator += tl.dot(a, b) * a_scale[:,
                                                      None] * b_scale[None, :]
            else:
                if use_fp8_w8a8:
                    # acc used to enable fp8_fast_accum
                    accumulator = tl.dot(a, b, acc=accumulator)
                else:
                    accumulator += tl.dot(a, b)
        else:
            accumulator += tl.dot(a, b)
        # Advance the ptrs to the next K block.
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk

    if MUL_ROUTED_WEIGHT:
        moe_weight = tl.load(topk_weights_ptr + offs_token,
                             mask=token_mask,
                             other=0)
        accumulator = accumulator * moe_weight[:, None]
    if use_int8_w8a16:
        accumulator = (accumulator * b_scale).to(compute_type)
    elif use_fp8_w8a8 or use_int8_w8a8:
        if group_k > 0 and group_n > 0:
            accumulator = accumulator.to(compute_type)
        else:
            accumulator = (accumulator * a_scale * b_scale).to(compute_type)
    else:
        accumulator = accumulator.to(compute_type)
    # -----------------------------------------------------------
    # Write back the block of the output
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[
        None, :]
    c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
    tl.store(c_ptrs, accumulator, mask=c_mask)


###################################################################################################################
##########################################       This is the beginning       ######################################
###################################################################################################################

# @perftest(num_warmup=1, num_iters=2)
def torch_moe_test(
    hidden_states,
    w1,
    w2,
    topk_weight,
    topk_ids,
    # following for int8 quant
    w1_scale=None,  # [expert, inter_dim, 1]
    w2_scale=None,  # [expert, model_dim, 1]
    fc1_smooth_scale=None,  # [expert, 1, model_dim]
    fc2_smooth_scale=None,  # [expert, 1, inter_dim]
    activation=ActivationType.Silu,
):

    out_stage1 = torch_moe_stage1(
                                 hidden_states,
                                 w1,  # E, inter_dim*2, model_dim
                                 w2,  # E, model_dim, inter_dim
                                 topk_weight,
                                 topk_ids)

    out_stage2 = torch_moe_stage2(
                                 out_stage1,
                                 w1,  # E, inter_dim*2, model_dim
                                 w2,  # E, model_dim, inter_dim
                                 topk_weight,
                                 topk_ids)
   
    return out_stage1, out_stage2

#    return torch_moe(
#        hidden_states,
#        w1,
#        w2,
#        topk_weight,
#        topk_ids,
#        w1_scale,
#        w2_scale,
#        fc1_smooth_scale,
#        fc2_smooth_scale,
#        None,
#        activation,
#    )

@perftest(num_warmup=1, num_iters=10, testGraph=True)
def asm_moe_test(
    hidden_states,
    w1,
    w2,
    topk_weight,
    topk_ids,
    # following for int8 quant
    w1_scale=None,  # [expert, inter_dim, 1]
    w2_scale=None,  # [expert, model_dim, 1]
    w1_zp=None,  # [expert, 1, model_dim]
    w2_zp=None,  # [expert, 1, inter_dim]
    activation=ActivationType.Silu,
    use_fp8_w8a8: bool = False,
    use_int8_w8a8: bool = False,
    use_int8_w8a16: bool = False,
    use_int4_w4a16: bool = False,
):

    expert_mask = None
    E, model_dim, inter_dim = w2.shape
    global_E = E
    if expert_mask is not None:
        global_E = expert_mask.numel()
    M, topk = topk_ids.shape
    dtype = hidden_states.dtype
    device = topk_ids.device
    lastdim_mul = 8 if w1.dtype in {dtypes.i32, torch.uint32} else 1
    sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, tokens_positions_per_expert, moe_buf = (
        moe_sorting_ck(
    topk_ids, topk_weight, global_E, model_dim, dtype, 16, expert_mask
    )
    )

    moe_buf = torch.empty((hidden_states.size(0), w2.size(1)), dtype=torch.float16, device="cuda")

    d_w1_out = torch.empty((hidden_states.size(0) * topk, w1.size(1)), dtype=torch.float16, device="cuda")
    if use_int4_w4a16:
      d_silu = torch.empty((hidden_states.size(0) * topk, w2.size(2) * 2), dtype=torch.float16, device="cuda")
    else:
      d_silu = torch.empty((hidden_states.size(0) * topk, w2.size(2)), dtype=torch.float16, device="cuda")
    d_w2_out = torch.empty((hidden_states.size(0), topk, w2.size(1)), dtype=torch.float16, device="cuda")

    if use_int4_w4a16:
      #scales = w1_scale.view(-1, w1_scale.shape[2])
      #zeros = w1_zp.view(-1, w1_zp.shape[2])

#      print("zcf scales shape = ", w1_scale.shape)

#      int_view = w1_scale[0:64, 0:1, 0:12].view(torch.uint16)
#      int_view = int_view.reshape(-1, int_view.shape[2])
#      for y in range(64):
#          hex_str = [hex(x.item()) for x in int_view[y]]
#          print("zcf scales 00 = ", hex_str)
                 
#      int_view = w1_zp[0:64, 0:1, 0:12].view(torch.uint8)
#      int_view = int_view.reshape(-1, int_view.shape[2])
#      for y in range(64):
#          hex_str = [hex(x.item()) for x in int_view[y]]
#          print("zcf zeros 00 = ", hex_str)
      #fp16
      aiter.asm_fmoe_stage1(d_w1_out, hidden_states, w1, w2, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, topk, w1_scale, w1_scale, w1_zp, 2, 0, 16, 0)
      #bf16
      #aiter.asm_fmoe_stage1(d_w1_out, hidden_states, w1, w2, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, topk, w1_scale, w1_scale, w1_zp, 3, 0, 16)
    else:
      aiter.asm_fmoe_stage1(d_w1_out, hidden_states, w1, w2, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, topk)

    aiter.silu_and_mul(d_silu, d_w1_out)

    if use_int4_w4a16:
#      int_view = w2[0:64, 0:16, 0:12].view(torch.uint8)
#      int_view = int_view.reshape(-1, int_view.shape[2])
#      for y in range(1024):
#          hex_str = [hex(x.item()) for x in int_view[y]]
#          print("zcf w2 00 = ", hex_str)
# 
#      int_view = w2_scale[0:64, 0:16, 0:12].view(torch.uint16)
#      int_view = int_view.reshape(-1, int_view.shape[2])
#      for y in range(1024):
#          hex_str = [hex(x.item()) for x in int_view[y]]
#          print("zcf scales 00 = ", hex_str)
#                
#      int_view = w2_zp[0:64, 0:16, 0:12].view(torch.uint8)
#      int_view = int_view.reshape(-1, int_view.shape[2])
#      for y in range(1024):
#          hex_str = [hex(x.item()) for x in int_view[y]]
#          print("zcf zeros 00 = ", hex_str)
      #fp16
      aiter.asm_fmoe_stage2(d_w2_out, d_silu, w1, w2, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, topk, w2_scale, w2_scale, w2_zp, 2, 0, 16, 0)
      #bf16
      #aiter.asm_fmoe_stage2(d_w2_out, d_silu, w1, w2, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, topk, w2_scale, w2_scale, w2_zp, 3, 0, 16)
    else:
      aiter.asm_fmoe_stage2(d_w2_out, d_silu, w1, w2, sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, topk)

    triton_moe_sum(d_w2_out, moe_buf)

    return d_silu, moe_buf


# vllm define
###################################################################################################################
def invoke_fused_moe_kernel(A: torch.Tensor,
                            B: torch.Tensor,
                            C: torch.Tensor,
                            A_scale: Optional[torch.Tensor],
                            B_scale: Optional[torch.Tensor],
                            B_zp: Optional[torch.Tensor],
                            topk_weights: Optional[torch.Tensor],
                            sorted_token_ids: torch.Tensor,
                            expert_ids: torch.Tensor,
                            num_tokens_post_padded: torch.Tensor,
                            mul_routed_weight: bool,
                            top_k: int,
                            config: dict[str, Any],
                            compute_type: tl.dtype,
                            use_fp8_w8a8: bool,
                            use_int8_w8a8: bool,
                            use_int8_w8a16: bool,
                            use_int4_w4a16: bool,
                            per_channel_quant: bool,
                            block_shape: Optional[list[int]] = None) -> None:
    assert topk_weights is not None or not mul_routed_weight
    assert topk_weights is None or topk_weights.stride(1) == 1
    assert sorted_token_ids.stride(0) == 1

    if use_fp8_w8a8 or use_int8_w8a8:
        assert B_scale is not None
        assert (block_shape is None
                or triton.cdiv(B.size(-2), block_shape[0]) == B_scale.size(-2))
        assert (block_shape is None
                or triton.cdiv(B.size(-1), block_shape[1]) == B_scale.size(-1))

    elif use_int8_w8a16 or use_int4_w4a16:
        assert B_scale is not None
        assert block_shape is None or block_shape[0] == 0
    else:
        assert A_scale is None
        assert B_scale is None

    M = A.size(0)
    num_tokens = M * top_k
    triton_config = config
    triton_config.pop('blas_MT0', None)
    triton_config.pop('blas_MT1', None)
    triton_config.pop('blas_DepthU', None)
    triton_config.pop('MODE', None)
    EM = sorted_token_ids.size(0)
    if A.size(0) < triton_config["BLOCK_SIZE_M"]:
        # optimize for small batch_size.
        # We assume that top_ids of each token is unique, so
        # so num_valid_experts <= batch_size <= BLOCK_SIZE_M,
        # and we can skip some invalid blocks.
        EM = min(sorted_token_ids.size(0),
                 A.size(0) * top_k * triton_config['BLOCK_SIZE_M'])
    grid = lambda META: (triton.cdiv(EM, META['BLOCK_SIZE_M']) * triton.cdiv(
        B.size(1), META['BLOCK_SIZE_N']), )

    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    total_time = 0.0
    start_event.record()
    for _ in range(1):
        if (use_int8_w8a16 or use_int4_w4a16) and \
                block_shape is not None and block_shape[1] > 0:
            assert B_scale is not None and B_scale.ndim == 3
            assert B_zp is None or B_zp.ndim == 3
            use_moe_wna16_cuda = should_moe_wna16_use_cuda(
                num_valid_tokens=num_tokens,
                group_size=block_shape[1],
                num_experts=B.size(0),
                bit=4 if use_int4_w4a16 else 8)
            triton_config = triton_config.copy()
            triton_config.update(
                get_moe_wna16_block_config(config=triton_config,
                                           use_moe_wna16_cuda=use_moe_wna16_cuda,
                                           num_valid_tokens=num_tokens,
                                           size_k=A.size(1),
                                           size_n=B.size(1),
                                           num_experts=B.size(1),
                                           group_size=block_shape[1],
                                           real_top_k=top_k,
                                           block_size_m=triton_config["BLOCK_SIZE_M"]))
            if use_moe_wna16_cuda:
                bit = 4 if use_int4_w4a16 else 8
                ops.moe_wna16_gemm(A, C, B, B_scale, B_zp,
                                   topk_weights if mul_routed_weight else None,
                                   sorted_token_ids, expert_ids,
                                   num_tokens_post_padded, top_k,
                                   triton_config["BLOCK_SIZE_M"], triton_config["BLOCK_SIZE_N"],
                                   triton_config["BLOCK_SIZE_K"], bit)
            else:
                fused_moe_kernel_gptq_awq[grid](
                    A,
                    B,
                    C,
                    B_scale,
                    B_zp,
                    topk_weights,
                    sorted_token_ids,
                    expert_ids,
                    num_tokens_post_padded,
                    B.size(1),
                    A.size(1),
                    EM,
                    num_tokens,
                    A.stride(0),
                    A.stride(1),
                    B.stride(0),
                    B.stride(2),
                    B.stride(1),
                    C.stride(1),
                    C.stride(2),
                    B_scale.stride(0),
                    B_scale.stride(2),
                    B_scale.stride(1),
                    B_zp.stride(0) if B_zp is not None else 0,
                    B_zp.stride(2) if B_zp is not None else 0,
                    B_zp.stride(1) if B_zp is not None else 0,
                    block_k_diviable=A.size(1) % triton_config["BLOCK_SIZE_K"] == 0,
                    group_size=block_shape[1],
                    MUL_ROUTED_WEIGHT=mul_routed_weight,
                    top_k=top_k,
                    compute_type=compute_type,
                    has_zp=B_zp is not None,
                    use_int4_w4a16=use_int4_w4a16,
                    use_int8_w8a16=use_int8_w8a16,
                    **triton_config,
                )
        else:
            triton_config = triton_config.copy()
            BLOCK_SIZE_K = triton_config.pop("BLOCK_SIZE_K")
            if block_shape is not None:
                BLOCK_SIZE_K = min(BLOCK_SIZE_K, min(block_shape[0],
                                                     block_shape[1]))
            fused_moe_kernel[grid](
                A,
                B,
                C,
                A_scale,
                B_scale,
                topk_weights,
                sorted_token_ids,
                expert_ids,
                num_tokens_post_padded,
                B.size(1),
                B.size(2),
                EM,
                num_tokens,
                A.stride(0),
                A.stride(1),
                B.stride(0),
                B.stride(2),
                B.stride(1),
                C.stride(1),
                C.stride(2),
                A_scale.stride(0)
                if A_scale is not None and A_scale.ndim == 2 else 0,
                A_scale.stride(1)
                if A_scale is not None and A_scale.ndim == 2 else 0,
                B_scale.stride(0)
                if B_scale is not None and B_scale.ndim >= 2 else 0,
                B_scale.stride(2)
                if B_scale is not None and B_scale.ndim == 3 else 0,
                B_scale.stride(1)
                if B_scale is not None and B_scale.ndim >= 2 else 0,
                0 if block_shape is None else block_shape[0],
                0 if block_shape is None else block_shape[1],
                MUL_ROUTED_WEIGHT=mul_routed_weight,
                top_k=top_k,
                compute_type=compute_type,
                use_fp8_w8a8=use_fp8_w8a8,
                use_int8_w8a8=use_int8_w8a8,
                use_int8_w8a16=use_int8_w8a16,
                per_channel_quant=per_channel_quant,
                BLOCK_SIZE_K=BLOCK_SIZE_K,
                **triton_config,
            )

    end_event.record()
    end_event.synchronize()
    total_time = start_event.elapsed_time(end_event)
    print(f"triton平均耗时: {total_time / 1:.3f} ms")

    out_cuda = torch.zeros_like(C)
    
    total_time = 0.0
    for _ in range(1):
        start_event.record()
        #moe_gemm_w8a8(A, 
        #                B, 
        #                out_cuda,
        #                A_scale, 
        #                B_scale,
        #                topk_weights if mul_routed_weight else None,
        #                sorted_token_ids, 
        #                expert_ids,
        #                num_tokens_post_padded, 
        #                top_k,
        #                config)
        end_event.record()
        end_event.synchronize()
        total_time += start_event.elapsed_time(end_event)
    
    print(f"asm平均耗时: {total_time / 1:.3f} ms")

    torch.cuda.synchronize()
    stage = 0
    if not mul_routed_weight:
        stage = 1
    else:
        stage = 2
    if not (torch.allclose(C, out_cuda, rtol=1e-2, atol=1e-2)):
        #if not mul_routed_weight and stage == 1:
        #    gemm1_compare_failed_list.append(config['MODE'])
        #elif stage == 2:
        #    gemm2_compare_failed_list.append(config['MODE'])

        print(f"stage{stage} kernel 精度检查不合格!!!")
    else:
        print(f"kernel{stage} 精度检查合格.")

def fused_experts_impl(
    hidden_states: torch.Tensor,
    w1: torch.Tensor,
    w2: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    inplace: bool = False,
    activation: str = "silu",
    apply_router_weight_on_input: bool = False,
    use_fp8_w8a8: bool = False,
    use_int8_w8a8: bool = False,
    use_int8_w8a16: bool = False,
    use_int4_w4a16: bool = False,
    use_mxfp4_w4a4: 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,
) -> torch.Tensor:
    # Check constraints.
    if use_int4_w4a16:
        assert hidden_states.size(1) // 2 == w1.size(2), (
            "Hidden size mismatch")
    elif use_mxfp4_w4a4:
        # 16bit activation and fp4x2 packed weight
        assert hidden_states.size(1) // 2 == w1.size(2), "hidden size mismatch"
    else:
        assert hidden_states.size(1) == w1.size(2), (
            f"Hidden size mismatch {hidden_states.size(1)} != {w1.size(2)}")

    assert topk_weights.size() == topk_ids.size(), "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
    ]

    num_tokens = hidden_states.size(0)
    E, N, _ = w1.size()
    K = w2.size(1)
    if global_num_experts == -1:
        global_num_experts = E
    top_k_num = topk_ids.size(1)
    # We execute the fused_moe kernel in chunks to circumvent this issue:
    # https://github.com/vllm-project/vllm/issues/5938
    CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
    M = min(num_tokens, CHUNK_SIZE)
    config_dtype = get_config_dtype_str(use_fp8_w8a8=use_fp8_w8a8,
                                        use_int8_w8a16=use_int8_w8a16,
                                        use_int4_w4a16=use_int4_w4a16,
                                        use_mxfp4_w4a4=use_mxfp4_w4a4,
                                        dtype=hidden_states.dtype)

    qtype = get_config_quant_dtype(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_mxfp4_w4a4=use_mxfp4_w4a4)

    get_config_func = functools.partial(
        try_get_optimal_moe_config,
        w1.size(),
        w2.size(),
        top_k_num,
        config_dtype,
        block_shape=block_shape,
    )

    config = get_config_func(M)
    # We can reuse the memory between these because by the time we need
    # cache3, we're done with cache1
    cache13 = torch.empty(M * top_k_num * max(N, K),
                          device=hidden_states.device,
                          dtype=hidden_states.dtype)
    intermediate_cache1 = cache13[:M * top_k_num * N].view(M, top_k_num, N)
    intermediate_cache3 = cache13[:M * top_k_num * K].view(M, top_k_num, K)

    # This needs separate memory since it's used concurrently with cache1
    intermediate_cache2 = torch.empty((M * top_k_num, N // 2),
                                      device=hidden_states.device,
                                      dtype=hidden_states.dtype)

    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
    else:
        raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")

    if inplace:
        out_hidden_states = hidden_states
    else:
        out_hidden_states = torch.empty_like(hidden_states)

    if use_mxfp4_w4a4:
        # Weight has to be dequantized for mxfp4 emulation.
        w1 = dequant_mxfp4(w1, w1_scale, hidden_states.dtype)
        w1_scale = None
        w2 = dequant_mxfp4(w2, w2_scale, hidden_states.dtype)
        w2_scale = None

    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.size()

        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.
            intermediate_cache1 = intermediate_cache1[:tokens_in_chunk]
            intermediate_cache2 = intermediate_cache2[:tokens_in_chunk *
                                                      topk_ids.size(1)]
            intermediate_cache3 = intermediate_cache3[:tokens_in_chunk]
            config = get_config_func(tokens_in_chunk)

        curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
        curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
        qcurr_hidden_states, a1q_scale = moe_kernel_quantize_input(
            A=curr_hidden_states,
            A_scale=a1_scale,
            quant_dtype=qtype,
            per_act_token_quant=per_channel_quant,
            block_shape=block_shape)

        config1 = config
        config2 = config
        if M > 128:
            config1['blas_MT0']     = 32
            config1['blas_MT1']     = 32
            config1['blas_DepthU']  = 128
            config1['MODE']         = 0
            config2['blas_MT0']     = 32
            config2['blas_MT1']     = 1024
            config2['blas_DepthU']  = 16
            config2['MODE']         = 0
        else:
            config1['blas_MT0']     = 32
            config1['blas_MT1']     = 32
            config1['blas_DepthU']  = 128
            config1['MODE']         = 0
            config2['blas_MT0']     = 32
            config2['blas_MT1']     = 1024
            config2['blas_DepthU']  = 16
            config2['MODE']         = 0

        config['BLOCK_SIZE_M'] = config1['blas_MT0']
        sorted_token_ids, expert_ids, num_tokens_post_padded = (
            moe_align_block_size(curr_topk_ids, config1['BLOCK_SIZE_M'],
                                 global_num_experts, expert_map))

        invoke_fused_moe_kernel(qcurr_hidden_states,
                                w1,
                                intermediate_cache1,
                                a1q_scale,
                                w1_scale,
                                w1_zp,
                                curr_topk_weights,
                                sorted_token_ids,
                                expert_ids,
                                num_tokens_post_padded,
                                apply_router_weight_on_input,
                                top_k_num,
                                config1,
                                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,
                                per_channel_quant=per_channel_quant,
                                block_shape=block_shape)

        if activation == "silu":
            torch.ops._C.silu_and_mul(intermediate_cache2,
                                      intermediate_cache1.view(-1, N))
        elif activation == "gelu":
            torch.ops._C.gelu_and_mul(intermediate_cache2,
                                      intermediate_cache1.view(-1, N))
        else:
            raise ValueError(f"Unsupported FusedMoe activation: {activation}")

        qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
            A=intermediate_cache2,
            A_scale=a2_scale,
            quant_dtype=qtype,
            per_act_token_quant=per_channel_quant,
            block_shape=block_shape)

        invoke_fused_moe_kernel(qintermediate_cache2,
                                w2,
                                intermediate_cache3,
                                a2q_scale,
                                w2_scale,
                                w2_zp,
                                curr_topk_weights,
                                sorted_token_ids,
                                expert_ids,
                                num_tokens_post_padded,
                                not apply_router_weight_on_input,
                                1,
                                config2,
                                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,
                                per_channel_quant=per_channel_quant,
                                block_shape=block_shape)

        ops.moe_sum(intermediate_cache3.view(*intermediate_cache3.size()),
                    out_hidden_states[begin_chunk_idx:end_chunk_idx])

    return out_hidden_states


# the above tests is for correctness, the following is for performance
###################################################################################################################
def perftest():
    M = [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,64,128,256,512]
    N = [256]
    K = [7168]
    E = [256]
    TOPK = [8]
    EP_SIZE = [1]
    DTYPE = [torch.float16]
    GROUP_SIZE = [64]
    HAS_ZP = [True]
    WEIGHT_BITS = [4]
    
    for m,n,k,e,topk,ep_size,dtype,group_size,has_zp,weight_bits in  itertools.product(
        M, N, K, E, TOPK, EP_SIZE, DTYPE, GROUP_SIZE, HAS_ZP, WEIGHT_BITS):
        
        input = torch.randn((m, k), device="cuda", dtype=dtype) / 10
        w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
        w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
        score = torch.randn((m, e), device="cuda", dtype=dtype)

        #input = torch.randint(0, 2, (m, k), device="cuda", dtype=dtype)
        #w1 = torch.randint(0, 2, (e, 2 * n, k), device="cuda", dtype=dtype)
        #w2 = torch.randint(0, 2, (e, k, n), device="cuda", dtype=dtype)
        #score = torch.randint(0, 2, (m, e), device="cuda", dtype=dtype)

        if weight_bits == 4:
            pack_factor = 2
            quant_type = scalar_types.uint4 if has_zp else scalar_types.uint4b8
        elif weight_bits == 8:
            pack_factor = 1
            quant_type = scalar_types.uint8 if has_zp else scalar_types.uint8b128

        w1_ref = w1.clone()
        w2_ref = w2.clone()
        w1_qweight = torch.empty((e, 2 * n, k // pack_factor),
                                device="cuda",
                                dtype=torch.uint8)
        w2_qweight = torch.empty((e, k, n // pack_factor),
                                device="cuda",
                                dtype=torch.uint8)
        w1_scales = torch.empty((e, 2 * n, k // group_size),
                                device="cuda",
                                dtype=dtype)
        w2_scales = torch.empty((e, k, n // group_size),
                                device="cuda",
                                dtype=dtype)
        #asm
        w1_qzeros = torch.empty((e, 2 * n, k // group_size // pack_factor),
                                device="cuda",
                                dtype=torch.uint8)
        w2_qzeros = torch.empty((e, k, math.ceil(n // group_size / pack_factor)),
                                device="cuda",
                                dtype=torch.uint8)
        #triton
        #w1_qzeros = torch.empty((e, 2 * n // pack_factor, k // group_size),
        #                        device="cuda",
        #                        dtype=torch.uint8)
        #w2_qzeros = torch.empty((e, k // pack_factor, n // group_size),
        #                        device="cuda",
        #                        dtype=torch.uint8)

        for i in range(e * 2):
            expert_id = i % e
            if i // e == 0:
                w, w_ref, w_qweight, w_scales, w_qzeros = \
                    w1, w1_ref, w1_qweight, w1_scales, w1_qzeros
            else:
                w, w_ref, w_qweight, w_scales, w_qzeros = \
                    w2, w2_ref, w2_qweight, w2_scales, w2_qzeros
            weight, qweight, scales, qzeros = quantize_weights(
                w[expert_id].T, quant_type, group_size, has_zp, False)
            weight = weight.T
            qweight = qweight.T.contiguous().to(torch.uint8)
            scales = scales.T
            if has_zp:
                qzeros = qzeros.T.contiguous().to(torch.uint8)
            if weight_bits == 4:
                qweight = qweight[:, 1::2] * 16 + qweight[:, ::2]   # 偶数列存储低4位，奇数列存储高4位
                if has_zp:
                    #asm qzeros
                    qzeros = qzeros[:, 1::2] * 16 + qzeros[:, ::2]
                    #triton qzeros
                    #qzeros = qzeros[1::2, :] * 16 + qzeros[::2, :]

            w_ref[expert_id] = weight
            w_qweight[expert_id] = qweight
            w_scales[expert_id] = scales
            if has_zp:
                w_qzeros[expert_id] = qzeros
        

        if ep_size > 1:
            local_e = e // ep_size
            e_ids = torch.randint(0,
                                e, (local_e, ),
                                device="cuda",
                                dtype=torch.int32)
            e_map = torch.full((e, ), -1, device="cuda", dtype=torch.int32)
            e_map[e_ids] = torch.arange(local_e, device="cuda", dtype=torch.int32)
            w1_ref = w1_ref[e_ids]
            w2_ref = w2_ref[e_ids]
            w1_qweight = w1_qweight[e_ids]
            w2_qweight = w2_qweight[e_ids]
            w1_scales = w1_scales[e_ids]
            w2_scales = w2_scales[e_ids]
            w1_qzeros = w1_qzeros[e_ids]
            w2_qzeros = w2_qzeros[e_ids]
        else:
            e_map = None
    
        ## without token topk score calc
        topk_weights, topk_ids = fused_topk(input, score, topk, True)
        
        print("###### topk_weights dtype = {}, shape = {}".format(topk_weights.dtype, topk_weights.shape))
        print("###### topk_ids dtype = {}, shape = {}".format(topk_ids.dtype, topk_ids.shape))
        print("###### w1_qweight dtype = {}, shape = {}".format(w1_qweight.dtype, w1_qweight.shape))
        print("###### w2_qweight dtype = {}, shape = {}".format(w2_qweight.dtype, w2_qweight.shape))
        print("###### w1_scales dtype = {}, shape = {}".format(w1_scales.dtype, w1_scales.shape))
        print("###### w2_scales dtype = {}, shape = {}".format(w2_scales.dtype, w2_scales.shape))
        print("###### w1_qzeros dtype = {}, shape = {}".format(w1_qzeros.dtype if has_zp else None, w1_qzeros.shape if has_zp else None))
        print("###### w2_qzeros dtype = {}, shape = {}".format(w2_qzeros.dtype if has_zp else None, w2_qzeros.shape if has_zp else None))
        print("###### w1_ref dtype = {}, shape = {}".format(w1_ref.dtype, w1_ref.shape))
        print("###### w2_ref dtype = {}, shape = {}".format(w2_ref.dtype, w2_ref.shape))

        torch_stage1, torch_output = torch_moe_test(input, w1_ref, w2_ref, topk_weights, topk_ids) 

        if has_zp:
          (asm_stage1, asm_output), ave_t  = asm_moe_test(input, 
                                            w1_qweight,
                                            w2_qweight,
                                            topk_weights,
                                            topk_ids,
                                            w1_scale=w1_scales,
                                            w2_scale=w2_scales,
                                            w1_zp=w1_qzeros if has_zp else None,
                                            w2_zp=w2_qzeros if has_zp else None,
                                            use_int4_w4a16 = weight_bits == 4
                                            ) 
        else:
          asm_stage1, asm_output = asm_moe_test(input, w1_ref, w2_ref, topk_weights, topk_ids) 


#        triton_output = fused_experts_impl(input,
#                                        w1_qweight,
#                                        w2_qweight,
#                                        topk_weights,
#                                        topk_ids,
#                                        use_int8_w8a16 = weight_bits == 8,
#                                        use_int4_w4a16 = weight_bits == 4,
#                                        global_num_experts=e,
#                                        expert_map=e_map,
#                                        w1_scale=w1_scales,
#                                        w2_scale=w2_scales,
#                                        w1_zp=w1_qzeros if has_zp else None,
#                                        w2_zp=w2_qzeros if has_zp else None,
#                                        block_shape=[0, group_size])

        # torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0)

        #rel_diff = (torch.mean(
        #    torch.abs(triton_output.to(torch.float32) - torch_output.to(torch.float32))) /
        #            torch.mean(torch.abs(torch_output.to(torch.float32))))
        #print("###### triton and torch diff = ", rel_diff)
        #checkAllclose(triton_output, torch_output, rtol=0.01, atol=0.01, msg="triton check")

#        print("zcf asm_stage1 = ", asm_stage1.shape);
#        print("zcf torch_stage1 = ", torch_stage1.shape);
        torch_stage1 = torch_stage1.reshape(-1,torch_stage1.shape[2])

#        int_view = asm_output[0:2, 0:32].view(torch.uint16)
#        for y in range(2):
#            hex_str = [hex(x.item()) for x in int_view[y]]
#            print("zcf asm_output 00 = ", hex_str)
#
#        int_view = torch_output[0:2, 0:32].view(torch.uint16)
#        for y in range(2):
#            hex_str = [hex(x.item()) for x in int_view[y]]
#            print("zcf torch_output 00 = ", hex_str)
        #or i in range(4096):
          #int_view = asm_stage1[2048+i:2049+i, 0:1].view(torch.uint16)
          #for y in range(1):
          #    hex_str = [x.item() for x in int_view[y]]
          #    print("zcf index 00 = ", hex_str)
        # asm_stage = asm_stage1[i:1+i, 0:1]
        # torch_stage = torch_stage1[i:1+i, 0:1]
        checkAllclose(asm_stage1, torch_stage1, rtol=0.01, atol=0.01, msg="asm stage1 check")
        checkAllclose(asm_output, torch_output, rtol=0.01, atol=0.01, msg="asm stage2 check")
        print("zcf time ", ave_t)

if __name__ == "__main__":
    perftest()
