import pytest
import torch
import itertools
import json
from itertools import zip_longest

from aiter import dtypes
from aiter.ops.triton.fused_moe import triton_moe_sum
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
from aiter import ActivationType
from aiter.test_common import checkAllclose, perftest
from aiter.fused_moe import moe_sorting

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

import ctypes
import aiter
import triton
import triton.language as tl
import vllm.envs as envs
from vllm import _custom_ops as ops
from vllm.logger import init_logger
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
    per_token_group_quant_fp8)
from vllm.model_executor.layers.quantization.utils.int8_utils import (
    per_token_group_quant_int8, per_token_quant_int8)
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_int4_w4a8: Optional[bool] = False,
                         use_int8_w8a16: Optional[bool] = False,
                         use_fp8_w8a8: Optional[bool] = False,
                         use_int8_w8a8: Optional[bool] = False,):
    if use_fp8_w8a8:
        return "fp8_w8a8"
    elif use_int8_w8a8:
        return "int8_w8a8"
    elif use_int8_w8a16:
        return "int8_w8a16"
    elif use_int4_w4a8:
        return "int4_w4a8"
    elif use_int4_w4a16:
        return "int4_w4a16"
    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,
    is_bottom: bool = False,
):
    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" or dtype == "int4_w4a8":
            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, is_bottom)

        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, is_bottom)
    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,
    is_bottom: bool = False,
) -> 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, is_bottom)

    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()}
    elif is_bottom:
        # if config with is_bottom json file not found, try to fallback use config without bottom json.
        fallback_json_file_name = get_config_file_name(E, N, dtype, block_shape)
        fallback_config_file_path = os.path.join(
            os.path.dirname(os.path.realpath(__file__)), "configs", fallback_json_file_name)

        if os.path.exists(fallback_config_file_path):
            with open(fallback_config_file_path) as f:
                logger.info("Using fallback configuration from %s for MoE layer.",
                            fallback_config_file_path)
                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
    #暂时为False
    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,
    is_bottom: bool = False,
) -> 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]
        config = {
            "BLOCK_SIZE_M": 64,
            "BLOCK_SIZE_N": block_shape[0],
            "BLOCK_SIZE_K": block_shape[1],
            "GROUP_SIZE_M": 32,
            "COMBINE_SCALE_LOAD": False,
            "num_warps": 4,
            "num_stages": 3,
        }
    elif dtype == "int4_w4a8" and block_shape is not None:
        config = {
            "BLOCK_SIZE_M": 16,
            "BLOCK_SIZE_N": 64,
            "BLOCK_SIZE_K": 64,
            "GROUP_SIZE_M": 1,
            "COMBINE_SCALE_LOAD": False,
            "num_warps": 4,
            "num_stages": 1,
        }
    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, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 1, "COMBINE_SCALE_LOAD": False}
        elif M <= 40:
            config = {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 1, "COMBINE_SCALE_LOAD": False}
        else:
            config = {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 1, "COMBINE_SCALE_LOAD": False}
    else:
        config = {
            "BLOCK_SIZE_M": 64,
            "BLOCK_SIZE_N": 64,
            "BLOCK_SIZE_K": 32,
            "GROUP_SIZE_M": 8,
            "COMBINE_SCALE_LOAD": False,
        }
        # A heuristic: fused marlin works faster with this config for small M
        if M <= E or (is_marlin and M <= 32):
            config = {
                "BLOCK_SIZE_M": 16,
                "BLOCK_SIZE_N": 32,
                "BLOCK_SIZE_K": 64,
                "GROUP_SIZE_M": 1,
                "COMBINE_SCALE_LOAD": False,
            }
    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
    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" or dtype == "int4_w4a8":
            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, is_bottom)

        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, is_bottom)
    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, :]).to(tl.int64)
    c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
    tl.store(c_ptrs, accumulator, mask=c_mask)



@triton.heuristics(values={
    'block_k_diviable': lambda nargs: nargs['K'] % (nargs['BLOCK_SIZE_K']) == 0,
    'block_n_diviable': lambda nargs: nargs['N'] % nargs['BLOCK_SIZE_N'] == 0,
    "num_groups": lambda nargs: triton.cdiv(nargs["BLOCK_SIZE_K"], nargs["group_size"]),
    'group_size_divisible': lambda nargs: nargs['BLOCK_SIZE_K'] % nargs['group_size'] == 0,
})
@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,
        block_n_diviable: tl.constexpr,
        group_size: tl.constexpr,
        group_size_divisible: tl.constexpr,
        num_groups: tl.constexpr,
        # Meta-parameters
        BLOCK_SIZE_M: tl.constexpr,
        BLOCK_SIZE_N: tl.constexpr,
        BLOCK_SIZE_K: tl.constexpr,
        GROUP_SIZE_M: tl.constexpr,
        COMBINE_SCALE_LOAD: tl.constexpr,
        MUL_ROUTED_WEIGHT: tl.constexpr,
        USE_ADDR_OFFSET_INT64_A: tl.constexpr,
        USE_ADDR_OFFSET_INT64_C: 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.
    """

    tl.assume(stride_am >= 0)
    tl.assume(stride_ak >= 0)
    tl.assume(stride_be >= 0)
    tl.assume(stride_bk >= 0)
    tl.assume(stride_bn >= 0)
    tl.assume(stride_cm >= 0)
    tl.assume(stride_cn >= 0)
    tl.assume(stride_bse >= 0)
    tl.assume(stride_bsk >= 0)
    tl.assume(stride_bsn >= 0)
    tl.assume(stride_bze >= 0)
    tl.assume(stride_bzk >= 0)
    tl.assume(stride_bzn >= 0)

    # to notify the compiler that sorted_token_ids_ptr is a pointer to the memory,
    # and all value in the memory is non-negative.
    tl.assume(sorted_token_ids_ptr.to(tl.int64) >= 0)

    tl.static_assert(COMBINE_SCALE_LOAD == False, "COMBINE_SCALE_LOAD not support for awq!")
    # -----------------------------------------------------------
    # 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.int32)
    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)
    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

    tl.assume(off_experts >= 0)

    offs_bn = (pid_n * BLOCK_SIZE_N +
               tl.arange(0, BLOCK_SIZE_N)).to(tl.int32) % N
    offs_k = tl.arange(0, BLOCK_SIZE_K).to(tl.int32)


    if USE_ADDR_OFFSET_INT64_A:
        a_ptrs = a_ptr + (offs_token[:, None] // top_k * stride_am +
                        offs_k[None, :] * stride_ak).to(tl.int64)
    else:
        a_ptrs = a_ptr + (offs_token[:, None] // top_k * stride_am +
                        offs_k[None, :] * stride_ak).to(tl.int32)

    if use_int4_w4a16:
        if group_size_divisible and has_zp:
            offs_k_continue = tl.arange(0, BLOCK_SIZE_K // 2).to(tl.int32)
            b_ptrs = b_ptr + (off_experts * stride_be + \
                offs_bn[:, None] * stride_bn + offs_k_continue[None, :] * \
                    stride_bk).to(tl.int32)
        else:
            b_ptrs = b_ptr + (off_experts * stride_be + \
                (offs_k[:, None] // 2) * stride_bk + offs_bn[None, :] * \
                    stride_bn).to(tl.int32)
        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).to(tl.int32)

    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

    # w4a16 deepseek case
    if use_int4_w4a16 and (group_size_divisible and has_zp):
        # -----------------------------------------------------------
        # 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.
            offs_szk = k * num_groups + tl.arange(0, num_groups)

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

                masks_szk = offs_szk < K // group_size
            else:
                k_mask = None
                k_other = None

                masks_szk = 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 = tl.interleave(b, b)
                b = tl.trans(b)

                b = (b >> b_shifter) & 0xF

            b_scale_ptrs = b_scale_ptr + (off_experts * stride_bse + \
                offs_bn[None, :] * stride_bsn + \
                offs_szk[:, None] * stride_bsk).to(tl.int32)
            b_scale = tl.load(b_scale_ptrs, mask=masks_szk, other=k_other)
            b_scale = b_scale.to(tl.float32)

            b_zp_ptrs = b_zp_ptr + (off_experts * stride_bze + \
                (offs_bn[None, :]//2) * stride_bzn + \
                offs_szk[:, None] * stride_bzk).to(tl.int32)
            b_zp = tl.load(b_zp_ptrs, mask=masks_szk, other=k_other)
            b_zp = ((b_zp >> b_zp_shifter) & 0xF)
            b_zp = b_zp.to(tl.float32)


            if num_groups == 1:
                # Original efficient implementation for single group
                b_scale = tl.broadcast_to(b_scale, (BLOCK_SIZE_K, BLOCK_SIZE_N))
                b_zp = tl.broadcast_to(b_zp, (BLOCK_SIZE_K, BLOCK_SIZE_N))
            else:
                # Reshape to (num_groups, 1, N) then broadcast to (num_groups, group_size_in_block, N)
                b_scale = tl.broadcast_to(b_scale[:, None, :], (num_groups, group_size, BLOCK_SIZE_N))
                b_scale = tl.reshape(b_scale, (BLOCK_SIZE_K, BLOCK_SIZE_N))
                b_zp = tl.broadcast_to(b_zp[:, None, :], (num_groups, group_size, BLOCK_SIZE_N))
                b_zp = tl.reshape(b_zp, (BLOCK_SIZE_K, BLOCK_SIZE_N))

            # 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
    else:
        # -----------------------------------------------------------
        # 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).to(tl.int32)
            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).to(tl.int32)
                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).to(tl.int32)
                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)

    if USE_ADDR_OFFSET_INT64_C:
        c_ptrs = c_ptr + (stride_cm * offs_token[:, None] + stride_cn * offs_cn[
            None, :]).to(tl.int64)
    else:
        c_ptrs = c_ptr + (stride_cm * offs_token[:, None] + stride_cn * offs_cn[
            None, :]).to(tl.int32)

    if block_n_diviable:
        c_mask = token_mask[:, None]
    else:
        c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
    tl.store(c_ptrs, accumulator, mask=c_mask)


@triton.heuristics(values={
    'block_k_diviable': lambda nargs: nargs['K'] % (nargs['BLOCK_SIZE_K']) == 0,
    'block_n_diviable': lambda nargs: nargs['N'] % nargs['BLOCK_SIZE_N'] == 0,
    "num_groups": lambda nargs: triton.cdiv(nargs["BLOCK_SIZE_K"], nargs["group_size"]),
    'group_size_divisible': lambda nargs: nargs['BLOCK_SIZE_K'] % nargs['group_size'] == 0,
})
@triton.jit
def fused_moe_kernel_gptq_awq_w4a8(
        # Pointers to matrices
        a_ptr,
        b_ptr,
        c_ptr,
        a_scale_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_asm,
        stride_ask,
        stride_bse,
        stride_bsk,
        stride_bsn,
        stride_bze,
        stride_bzk,
        stride_bzn,
        group_k: tl.constexpr,                 # a quant group size: ie. 128
        block_k_diviable: tl.constexpr,
        block_n_diviable: tl.constexpr,
        group_size: tl.constexpr,              # b w4 group size: ie. 64
        group_size_divisible: tl.constexpr,
        num_groups: tl.constexpr,
        # Meta-parameters
        BLOCK_SIZE_M: tl.constexpr,
        BLOCK_SIZE_N: tl.constexpr,
        BLOCK_SIZE_K: tl.constexpr,
        GROUP_SIZE_M: tl.constexpr,
        COMBINE_SCALE_LOAD: tl.constexpr,
        MUL_ROUTED_WEIGHT: tl.constexpr,
        USE_ADDR_OFFSET_INT64_A: tl.constexpr,
        USE_ADDR_OFFSET_INT64_C: tl.constexpr,
        top_k: tl.constexpr,
        compute_type: tl.constexpr,
        has_zp: tl.constexpr,
        use_int4_w4a16: tl.constexpr,
        use_int4_w4a8: 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.
    """

    tl.assume(stride_am >= 0)
    tl.assume(stride_ak >= 0)
    tl.assume(stride_be >= 0)
    tl.assume(stride_bk >= 0)
    tl.assume(stride_bn >= 0)
    tl.assume(stride_cm >= 0)
    tl.assume(stride_cn >= 0)
    tl.assume(stride_asm >= 0)
    tl.assume(stride_ask >= 0)
    tl.assume(stride_bse >= 0)
    tl.assume(stride_bsk >= 0)
    tl.assume(stride_bsn >= 0)
    tl.assume(stride_bze >= 0)
    tl.assume(stride_bzk >= 0)
    tl.assume(stride_bzn >= 0)

    # to notify the compiler that sorted_token_ids_ptr is a pointer to the memory,
    # and all value in the memory is non-negative.
    tl.assume(sorted_token_ids_ptr.to(tl.int64) >= 0)

    tl.static_assert(use_int4_w4a8 == True, "Must use int4_w4a8")
    tl.static_assert(has_zp == True, "only for deepseek w4a8 case.")
    tl.static_assert(group_k > 0, "group_k must be greater than 0 for deepseek w4a8 case.")
    tl.static_assert(block_k_diviable == True, "block_k_diviable must be True for deepseek w4a8 case.")
    tl.static_assert(BLOCK_SIZE_K <= group_k and group_k % BLOCK_SIZE_K == 0,
        "BLOCK_SIZE_K must be divisible by GROUP_SIZE_K")
    tl.static_assert(group_size_divisible == True, "group_size_divisible must be True for deepseek w4a8 case.")
    tl.static_assert(BLOCK_SIZE_K == group_size, "BLOCK_SIZE_K must be equal to group_size for deepseek w4a8 case.")
    tl.static_assert(group_size == group_k, "group_size must be equal to group_k for deepseek w4a8 case.")

    # -----------------------------------------------------------
    # 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.int32)
    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)
    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

    tl.assume(off_experts >= 0)

    offs_bn = (pid_n * BLOCK_SIZE_N +
               tl.arange(0, BLOCK_SIZE_N)).to(tl.int32) % N
    offs_k = tl.arange(0, BLOCK_SIZE_K).to(tl.int32)


    if USE_ADDR_OFFSET_INT64_A:
        a_ptrs = a_ptr + (offs_token[:, None] // top_k * stride_am +
                        offs_k[None, :] * stride_ak).to(tl.int64)
    else:
        a_ptrs = a_ptr + (offs_token[:, None] // top_k * stride_am +
                        offs_k[None, :] * stride_ak).to(tl.int32)

    offs_k_continue = tl.arange(0, BLOCK_SIZE_K // 2).to(tl.int32)
    b_ptrs = b_ptr + (off_experts * stride_be + \
        offs_bn[:, None] * stride_bn + offs_k_continue[None, :] * \
            stride_bk).to(tl.int32)   #[N, K//2]
    b_shifter = (offs_k[:, None] % 2) * 4  #[K]

    b_zp_shifter = (offs_bn[None, :] % 2) * 4

    if COMBINE_SCALE_LOAD:
        a_scale_ptrs = a_scale_ptr + (offs_token[:, None] // top_k) * stride_asm
    else:
        a_scale_ptrs = a_scale_ptr + (offs_token // top_k) * stride_asm

    # w4a8 deepseek case

    # -----------------------------------------------------------
    # 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)
    if COMBINE_SCALE_LOAD:
        for k in range(0, tl.cdiv(K, BLOCK_SIZE_K), 2):
            # Load the next block of A and B, generate a mask by checking the
            # K dimension.
            offs_szk = k + tl.arange(0, 2)

            a0 = tl.load(a_ptrs, mask=token_mask[:, None], other=0.0)
            b0 = tl.load(b_ptrs)

            b0 = tl.interleave(b0, b0)  #[N, K]
            b0 = tl.trans(b0)  #[K, N]
            b0 = (b0 >> b_shifter) & 0xF
            b0 = b0.to(tl.int32)

            offs_ks = k + tl.arange(0, 2)
            a_scale = tl.load(a_scale_ptrs + offs_ks[None, :] * stride_ask,
                              mask=token_mask[:, None],
                              other=0.0) # [M, 2]
            a_scale0, a_scale1 = tl.split(a_scale) # [M]

            # b_scale shape: [N, K] = [N, 2]
            b_scale_ptrs = b_scale_ptr + (off_experts * stride_bse + \
                offs_bn[:, None] * stride_bsn + \
                offs_szk[None, :] * stride_bsk).to(tl.int32)
            b_scale = tl.load(b_scale_ptrs)
            b_scale = b_scale.to(tl.float32)
            b_scale0, b_scale1 = tl.split(b_scale) # [N]


            b_zp_ptrs = b_zp_ptr + (off_experts * stride_bze + \
                (offs_bn[:, None]//2) * stride_bzn + \
                offs_szk[None, :] * stride_bzk).to(tl.int32)   #[N, 2]
            b_zp = tl.load(b_zp_ptrs)
            b_zp0, b_zp1 = tl.split(b_zp) # [N]

            b_zp0 = ((b_zp0 >> b_zp_shifter) & 0xF)
            b_zp1 = ((b_zp1 >> b_zp_shifter) & 0xF)
            b_zp0 = b_zp0.to(tl.int32)
            b_zp1 = b_zp1.to(tl.int32)

            a1 = tl.load(a_ptrs + BLOCK_SIZE_K * stride_ak, mask=token_mask[:, None], other=0.0)
            b1 = tl.load(b_ptrs + (BLOCK_SIZE_K // 2) * stride_bk)
            b1 = tl.interleave(b1, b1)
            b1 = tl.trans(b1)
            b1 = (b1 >> b_shifter) & 0xF
            b1 = b1.to(tl.int32)


            b_scale0 = tl.reshape(b_scale0, (1, BLOCK_SIZE_N))
            b_scale1 = tl.reshape(b_scale1, (1, BLOCK_SIZE_N))
            b_scale0 = tl.broadcast_to(b_scale0, (BLOCK_SIZE_M, BLOCK_SIZE_N))
            b_scale1 = tl.broadcast_to(b_scale1, (BLOCK_SIZE_M, BLOCK_SIZE_N))

            b_zp0 = tl.reshape(b_zp0, (1, BLOCK_SIZE_N))
            b_zp1 = tl.reshape(b_zp1, (1, BLOCK_SIZE_N))
            b_zp0 = tl.broadcast_to(b_zp0, (BLOCK_SIZE_K, BLOCK_SIZE_N))
            b_zp1 = tl.broadcast_to(b_zp1, (BLOCK_SIZE_K, BLOCK_SIZE_N))


            # We accumulate along the K dimension.
            b0 = (b0 - b_zp0).to(tl.int8)
            # accumulator += tl.dot(a0, b0) * a_scale0[:, None] * b_scale0[None, :]
            accumulator += tl.dot(a0, b0) * a_scale0[:, None] * b_scale0

            # We accumulate along the K dimension.
            b1 = (b1 - b_zp1).to(tl.int8)
            accumulator += tl.dot(a1, b1) * a_scale1[:, None] * b_scale1

            # Advance the ptrs to the next K block.
            a_ptrs += BLOCK_SIZE_K * stride_ak * 2
            b_ptrs += (BLOCK_SIZE_K // 2) * stride_bk * 2
    else:
        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.
            tl.static_assert(num_groups == 1, "num_groups must be 1")
            offs_szk = k * num_groups

            a = tl.load(a_ptrs, mask=token_mask[:, None], other=0.0)
            b = tl.load(b_ptrs)

            b = tl.interleave(b, b)
            b = tl.trans(b)
            b = (b >> b_shifter) & 0xF
            b = b.to(tl.int32)

            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_ptrs = b_scale_ptr + (off_experts * stride_bse + \
                offs_bn[None, :] * stride_bsn + \
                offs_szk * stride_bsk).to(tl.int32)
            b_scale = tl.load(b_scale_ptrs)
            b_scale = b_scale.to(tl.float32)

            b_zp_ptrs = b_zp_ptr + (off_experts * stride_bze + \
                (offs_bn[None, :]//2) * stride_bzn + \
                offs_szk[:, None] * stride_bzk).to(tl.int32)
            b_zp = tl.load(b_zp_ptrs)
            b_zp = ((b_zp >> b_zp_shifter) & 0xF)
            b_zp = b_zp.to(tl.int32)


            b_scale = tl.broadcast_to(b_scale, (BLOCK_SIZE_M, BLOCK_SIZE_N))
            b_zp = tl.broadcast_to(b_zp, (BLOCK_SIZE_K, BLOCK_SIZE_N))

            # We accumulate along the K dimension.
            b = (b - b_zp).to(tl.int8)
            accumulator += tl.dot(a, b) * a_scale[:, None] * b_scale

            # Advance the ptrs to the next K block.
            a_ptrs += BLOCK_SIZE_K * stride_ak
            b_ptrs += (BLOCK_SIZE_K // 2) * 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)

    if USE_ADDR_OFFSET_INT64_C:
        c_ptrs = c_ptr + (stride_cm * offs_token[:, None] + stride_cn * offs_cn[
            None, :]).to(tl.int64)
    else:
        c_ptrs = c_ptr + (stride_cm * offs_token[:, None] + stride_cn * offs_cn[
            None, :]).to(tl.int32)

    if block_n_diviable:
        c_mask = token_mask[:, None]
    else:
        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,
    block_k_diviable: tl.constexpr,
    block_n_diviable: tl.constexpr,
    # Meta-parameters
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
    COMBINE_SCALE_LOAD: tl.constexpr,
    MUL_ROUTED_WEIGHT: tl.constexpr,
    USE_ADDR_OFFSET_INT64_A: tl.constexpr,
    USE_ADDR_OFFSET_INT64_C: 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.
    """

    tl.assume(stride_am >= 0)
    tl.assume(stride_ak >= 0)
    tl.assume(stride_be >= 0)
    tl.assume(stride_bk >= 0)
    tl.assume(stride_bn >= 0)
    tl.assume(stride_cm >= 0)
    tl.assume(stride_cn >= 0)
    tl.assume(stride_bse >= 0)
    tl.assume(stride_bsk >= 0)
    tl.assume(stride_bsn >= 0)

    # to notify the compiler that sorted_token_ids_ptr is a pointer to the memory,
    # and all value in the memory is non-negative.
    tl.assume(sorted_token_ids_ptr.to(tl.int64) >= 0)

    if group_k > 0:
        tl.static_assert(BLOCK_SIZE_K <= group_k and group_k % BLOCK_SIZE_K == 0,
            "BLOCK_SIZE_K must be divisible by GROUP_SIZE_K")
    if COMBINE_SCALE_LOAD: # used for use_int8_w8a8
        tl.static_assert(stride_ask == 1,
            "COMBINE_SCALE_LOAD implictly stride_ask == 1!")
        tl.static_assert(MUL_ROUTED_WEIGHT == False,
            "COMBINE_SCALE_LOAD and MUL_ROUTED_WEIGHT cannot be both true due to w1_scale and w2_scale diff layout!")
        tl.static_assert(block_k_diviable == True and BLOCK_SIZE_K == group_k,
            "COMBINE_SCALE_LOAD only add and verify on block_k_diviable!")
        tl.static_assert(use_int8_w8a8 == True,
            "COMBINE_SCALE_LOAD only add and verify on use_int8_w8a8!")
    # -----------------------------------------------------------
    # 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.int32)
    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.int32)
    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.int32)) % N

    offs_k = tl.arange(0, BLOCK_SIZE_K)
    if USE_ADDR_OFFSET_INT64_A:
        a_ptrs = a_ptr + (offs_token[:, None] // top_k * stride_am +
                offs_k[None, :] * stride_ak).to(tl.int64)
    else:
        a_ptrs = a_ptr + (offs_token[:, None] // top_k * stride_am +
                offs_k[None, :] * stride_ak).to(tl.int32)
    b_ptrs = b_ptr + off_experts * stride_be + (offs_k[:, None] * stride_bk +
                                                offs_bn[None, :] * stride_bn).to(tl.int32)

    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:
            if COMBINE_SCALE_LOAD:
                a_scale_ptrs = a_scale_ptr + (offs_token[:, None] // top_k) * stride_asm
                offs_bsn = offs_bn // group_n
                b_scale_ptrs = (b_scale_ptr + off_experts * stride_bse +
                                offs_bsn[:, None] * stride_bsn)
            else:
                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)
    if COMBINE_SCALE_LOAD:
        for k in range(0, tl.cdiv(K, BLOCK_SIZE_K), 2):
            # Load the next block of A and B, generate a mask by checking the
            # K dimension.
            if not block_k_diviable:
                a0 = tl.load(a_ptrs,
                            mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K),
                            other=0.0)
                b0 = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
            else:
                a0 = tl.load(a_ptrs, mask=token_mask[:, None], other=0.0)
                b0 = tl.load(b_ptrs)

            # We accumulate along the K dimension.
            if use_int8_w8a16:
                accumulator = tl.dot(a0, b0.to(compute_type), acc=accumulator)
                tl.static_assert(False, "Not implemented")
            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 + tl.arange(0, 2)
                    a_scale = tl.load(a_scale_ptrs + offs_ks[None, :] * stride_ask,
                                    mask=token_mask[:, None],
                                    other=0.0)
                    b_scale = tl.load(b_scale_ptrs + offs_ks[None, :] * stride_bsk)
                    a_scale_0, a_scale_1 = tl.split(a_scale)
                    b_scale_0, b_scale_1 = tl.split(b_scale)

                    accumulator += tl.dot(a0, b0) * a_scale_0[:,
                                                        None] * b_scale_0[None, :]

                    if not block_k_diviable:
                        a1 = tl.load(a_ptrs + BLOCK_SIZE_K * stride_ak,
                                    mask=token_mask[:, None] & (offs_k[None, :] < K - (k + 1) * BLOCK_SIZE_K),
                                    other=0.0)
                        b1 = tl.load(b_ptrs + BLOCK_SIZE_K * stride_bk,
                                    mask=offs_k[:, None] < K - (k + 1) * BLOCK_SIZE_K, other=0.0)
                    else:
                        a1 = tl.load(a_ptrs + BLOCK_SIZE_K * stride_ak, mask=token_mask[:, None], other=0.0)
                        b1 = tl.load(b_ptrs + BLOCK_SIZE_K * stride_bk)
                    accumulator += tl.dot(a1, b1) * a_scale_1[:,
                                                        None] * b_scale_1[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)
                    tl.static_assert(False, "Not implemented")
            else:
                accumulator += tl.dot(a, b)
                tl.static_assert(False, "Not implemented")

            # Advance the ptrs to the next K block.
            a_ptrs += BLOCK_SIZE_K * stride_ak * 2
            b_ptrs += BLOCK_SIZE_K * stride_bk * 2

    else: # non-COMBINE_SCALE_LOAD

        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:
                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)
            else:
                a = tl.load(a_ptrs, mask=token_mask[:, None], other=0.0)
                b = tl.load(b_ptrs)

            # 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)
    if USE_ADDR_OFFSET_INT64_C:
        c_ptrs = c_ptr + (stride_cm * offs_token[:, None] + stride_cn * offs_cn[
            None, :]).to(tl.int64)
    else:
        c_ptrs = c_ptr + (stride_cm * offs_token[:, None] + stride_cn * offs_cn[
            None, :]).to(tl.int32)
    if not block_n_diviable:
        c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
    else:
        c_mask = token_mask[:, None]

    tl.store(c_ptrs, accumulator, mask=c_mask)



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 envs.VLLM_ENABLE_MOE_ALIGN_BLOCK_SIZE_TRITON or 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
            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)
    if expert_map is not None:
        expert_ids = expert_map[expert_ids]

    return sorted_ids, expert_ids, num_tokens_post_pad


###################################################################################################################
##########################################       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,
):
    return torch_moe(
        hidden_states,
        w1,
        w2,
        topk_weight,
        topk_ids,
        w1_scale,
        w2_scale,
        fc1_smooth_scale,
        fc2_smooth_scale,
        None,
        activation,
    )


# 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],
                            B_zp_asm: Optional[torch.Tensor],
                            topk_weights: torch.Tensor,
                            topk_ids: 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],
                            sorted_ids: torch.Tensor,
                            sorted_expert_ids: torch.Tensor,
                            top_k_num: int,
                            compute_type: tl.dtype,
                            use_fp8_w8a8: bool,
                            use_int8_w8a8: bool,
                            use_int8_w8a16: bool,
                            use_int4_w4a16: bool,
                            use_int4_w4a8: bool,
                            per_channel_quant: bool,
                            block_shape: Optional[List[int]] = None) -> None:
    assert topk_weights.stride(1) == 1
    assert sorted_token_ids.stride(0) == 1
    odtype = 0
    if A.dtype == torch.bfloat16:
        odtype = 1
    use_triton_kernel_int4_w4a8 = True

    if use_int4_w4a8:
        assert B_scale is not None
        assert block_shape is None or block_shape[0] == 0

        # if not use_triton_kernel_int4_w4a8:
        #     # step 1. dequantize B
        #     # B_dq = custom_int4_dequant(B, B_scale, B_zp,
        #     #                 group_size=block_shape[1], pack_factor=2,
        #     #                 output_dtype=torch.float32)
        #     # B_dq = triton_int4_dequant(B, B_scale, B_zp, block_shape[1])
        #     # # step 2. quantize B_q
        #     # # B_q, B_s = custom_block_quant(B_dq, [128, 128], eps=1e-10, dtype=torch.int8)
        #     # B_q, B_s = custom_block_quant_triton_3d(B_dq, [128, 128], eps=1e-10, dtype=torch.int8)

        #     B_q, B_s = fused_int4_to_block_int8(B, B_scale, B_zp, block_shape[1], [128, 128])

        #     # step 3. rename or reassign value.
        #     B = B_q
        #     B_scale = B_s
        #     B_zp = None
        #     block_shape = [128, 128]
        #     use_int4_w4a16 = False
        #     use_int8_w8a8  = True
        #     use_int4_w4a8 = False

    if use_fp8_w8a8:
        assert B_scale is not None
        if block_shape is None:
            A, A_scale = ops.scaled_fp8_quant(A, A_scale)
        else:
            assert len(block_shape) == 2
            block_n, block_k = block_shape[0], block_shape[1]
            A, A_scale = per_token_group_quant_fp8(A, block_k)
            assert triton.cdiv(A.shape[-1], block_k) == A_scale.shape[-1]
            assert triton.cdiv(B.shape[-2], block_n) == B_scale.shape[-2]
            assert triton.cdiv(B.shape[-1], block_k) == B_scale.shape[-1]
    elif use_int8_w8a8:
        assert B_scale is not None
        if block_shape is None:
            A, A_scale = per_token_quant_int8(A)
        else:
            assert len(block_shape) == 2
            block_n, block_k = block_shape[0], block_shape[1]
            A, A_scale = per_token_group_quant_int8(A, block_k)
            assert triton.cdiv(A.shape[-1], block_k) == A_scale.shape[-1]
            assert triton.cdiv(B.shape[-2], block_n) == B_scale.shape[-2]
            assert triton.cdiv(B.shape[-1], block_k) == B_scale.shape[-1]
    elif use_int4_w4a8:
        assert B_scale is not None
        assert block_shape is None or block_shape[0] == 0
        block_n, block_k = block_shape[0], block_shape[1]
        A, A_scale = per_token_group_quant_int8(A, block_k)
        assert triton.cdiv(A.shape[-1], block_k) == A_scale.shape[-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

    MODE = config["MODE"]
    config.pop('MODE', None)
    EM = sorted_token_ids.shape[0]
    #AA = torch.ones_like(A)
    #A =  AA
    #BB = torch.ones_like(B)
    #B =  BB
    #AA_scale = torch.ones_like(A_scale)
    #A_scale = AA_scale
    #BB_scale = torch.ones_like(B_scale)
    #B_scale = BB_scale
    #print("###### A_scale dtype = {}, shape = {}".format(A_scale.dtype, A_scale.shape))
    #print("###### B_scale dtype = {}, shape = {}".format(B_scale.dtype, B_scale.shape))
    #opk_weightss  = torch.ones_like(topk_weights) / 100
    #opk_weights = topk_weightss
    if A.shape[0] < 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.shape[0],
                 A.shape[0] * top_k * config['BLOCK_SIZE_M'])
    grid = lambda META: (triton.cdiv(EM, META['BLOCK_SIZE_M']) * triton.cdiv(
        B.shape[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 or use_int4_w4a8) 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=topk_ids.numel(),
                group_size=block_shape[1],
                num_experts=B.shape[0],
                bit=4 if use_int4_w4a16 else 8)
            config = config.copy()
            config.update(
                get_moe_wna16_block_config(config=config,
                                           use_moe_wna16_cuda=use_moe_wna16_cuda,
                                           num_valid_tokens=topk_ids.numel(),
                                           size_k=A.shape[1],
                                           size_n=B.shape[1],
                                           num_experts=B.shape[1],
                                           group_size=block_shape[1],
                                           real_top_k=topk_ids.shape[1],
                                           block_size_m=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,
                                   config["BLOCK_SIZE_M"], config["BLOCK_SIZE_N"],
                                   config["BLOCK_SIZE_K"], bit)
                return

            offset_max = 2**31 - 1
            use_addr_offset_int64_a = A.numel() * A.element_size() >= offset_max
            use_addr_offset_int64_c = C.numel() * C.element_size() >= offset_max
            if (A.numel() * A.element_size() >= offset_max or
                B.numel() * B.element_size() >= offset_max or
                C.numel() * C.element_size() >= offset_max):
                logger.warning(
                    ("A,B,C numel：%ld, %ld, %ld has out of range for fused_moe_kernel kernel!"
                    "Use int64 for address offset."), A.numel(), B.numel(), C.numel())

            if use_int4_w4a8:
                fused_moe_kernel_gptq_awq_w4a8[grid](
                    A,
                    B,
                    C,
                    A_scale,
                    B_scale,
                    B_zp,
                    topk_weights,
                    sorted_token_ids,
                    expert_ids,
                    num_tokens_post_padded,
                    B.shape[1],
                    A.shape[1],
                    EM,
                    topk_ids.numel(),
                    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),
                    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,
                    group_k=block_shape[1],
                    group_size=block_shape[1],
                    MUL_ROUTED_WEIGHT=mul_routed_weight,
                    USE_ADDR_OFFSET_INT64_A=use_addr_offset_int64_a,
                    USE_ADDR_OFFSET_INT64_C=use_addr_offset_int64_c,
                    top_k=top_k,
                    compute_type=compute_type,
                    has_zp=B_zp is not None,
                    use_int4_w4a16=use_int4_w4a16,
                    use_int4_w4a8=use_int4_w4a8,
                    use_int8_w8a16=use_int8_w8a16,
                    **config
                )
            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.shape[1],
                    A.shape[1],
                    EM,
                    topk_ids.numel(),
                    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,
                    group_size=block_shape[1],
                    MUL_ROUTED_WEIGHT=mul_routed_weight,
                    USE_ADDR_OFFSET_INT64_A=use_addr_offset_int64_a,
                    USE_ADDR_OFFSET_INT64_C=use_addr_offset_int64_c,
                    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,
                    **config,
                )
        else:
            # simple check out of range, not accurate.
            offset_max = 2**31 - 1
            use_addr_offset_int64_a = A.numel() * A.element_size() >= offset_max
            use_addr_offset_int64_c = C.numel() * C.element_size() >= offset_max
            if (A.numel() * A.element_size() >= offset_max or
                B.numel() * B.element_size() >= offset_max or
                C.numel() * C.element_size() >= offset_max):
                logger.warning(
                    ("A,B,C numel：%ld, %ld, %ld has out of range for fused_moe_kernel kernel!"
                    "Use int64 for address offset."), A.numel(), B.numel(), C.numel())


            config = config.copy()
            BLOCK_SIZE_K = 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.shape[1],
                B.shape[2],
                EM,
                topk_ids.numel(),
                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],
                block_k_diviable=A.shape[1] % BLOCK_SIZE_K == 0,
                block_n_diviable=B.shape[1] % config["BLOCK_SIZE_N"] == 0,
                MUL_ROUTED_WEIGHT=mul_routed_weight,
                USE_ADDR_OFFSET_INT64_A=use_addr_offset_int64_a,
                USE_ADDR_OFFSET_INT64_C=use_addr_offset_int64_c,
                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,
                BLOCK_SIZE_K=BLOCK_SIZE_K,
                per_channel_quant=per_channel_quant,
                COMBINE_SCALE_LOAD=config.pop("COMBINE_SCALE_LOAD", None),
                **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)
    #print("out_cuda",out_cuda.shape)
    #print("sorted_token_ids",sorted_token_ids.shape)
    #print("topk_weights",topk_weights.shape)
    #print("topk_weights",topk_weights)
    #print("expert_ids",expert_ids.shape)

    #print("out_cuda ",hex(out_cuda.data_ptr()))
    #print("A ",hex(A.data_ptr()))
    #print("B ",hex(B.data_ptr()))
    #print("A_scale" ,hex(A_scale.data_ptr()))
    #print("B_scale ",hex(B_scale.data_ptr()))
    #print("B_zp_asm ",hex(B_zp_asm.data_ptr()))
    #print("sorted_token_ids ",hex(sorted_token_ids.data_ptr()))
    #print("topk_weights ",hex(topk_weights.data_ptr()))
    #print("expert_ids ",hex(expert_ids.data_ptr()))
    #print("num_tokens_post_padded ",hex(num_tokens_post_padded.data_ptr()))

    #print("A shape",A.shape)
    #print("A ",A)
    #print("A_scale shape",A_scale.shape)
    #print("A_scale",A_scale)
    #print("B shape",B.shape)
    #print("B",B)
    #print("B_scale shape",B_scale.shape)
    #print("B_scale",B_scale)
    #print("B_zp_asm",B_zp_asm.shape)
    #print("B_zp_asm",B_zp_asm)
    #torch.set_printoptions(threshold=40_000)
    #print("sorted_token_ids ",sorted_token_ids)
    total_time = 0.0
    if not mul_routed_weight:
        stage = 10
    else:
        stage = 11     
    start_event.record()
    cnt = 5
    for _ in range(cnt):
        aiter.asm_fmoe_a8(out_cuda, 
                          A,
                          B,
                          B,
                          sorted_ids,
                          topk_weights,
                          sorted_expert_ids,
                          num_tokens_post_padded,
                          top_k_num,
                          scale_a=A_scale,
                          scale_b=B_scale,
                          zero_points=B_zp_asm,
                          mode = stage,
                          solidx = MODE,
                          out_type = odtype,
                          persist_groups = 64)
    end_event.record()
    end_event.synchronize()
    total_time += start_event.elapsed_time(end_event)
    stage = 0
    if not mul_routed_weight:
        stage = 1
        print(f"asms{stage} m = {A.size(0)} BLOCK_SIZE_M = {config['BLOCK_SIZE_M']} MODE = {MODE} 平均耗时: {total_time / cnt:.3f} ms")
    else:
        stage = 2      
        print(f"asms{stage} m = {A.size(0) // 8} BLOCK_SIZE_M = {config['BLOCK_SIZE_M']} MODE = {MODE} 平均耗时: {total_time / cnt:.3f} ms")

    if not (torch.allclose(C, out_cuda, rtol=1e-2, atol=1e-2)):
        print(f"stage{stage} MODE = {MODE} kernel type={odtype} 精度检查不合格!!!")
        checkAllclose(C, out_cuda, rtol=0.01, atol=0.01, msg=f"asm check mode{MODE}")
        #torch.set_printoptions(threshold=40_0000)
        #torch.set_printoptions(linewidth=110)
        #torch.set_printoptions(precision=4, sci_mode=True)
        print("C  ",C)
        #print("C  ",C.shape)
        #torch.set_printoptions(linewidth=110)
        #torch.set_printoptions(precision=4, sci_mode=True)
        print("D  ",out_cuda)          
    else:
        print(f"kernel{stage} 精度检查合格.")

#@perftest(num_warmup=1, num_iters=2, testGraph=True)
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",
                       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,
                       w1_zp_asm: Optional[torch.Tensor] = None,
                       w2_zp_asm: Optional[torch.Tensor] = None,                       
                       a1_scale: Optional[torch.Tensor] = None,
                       a2_scale: Optional[torch.Tensor] = None,
                       block_shape: Optional[List[int]] = None,
                       Test: dict[str, Any] = None):
    # 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
    ]

    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 = 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_w8a8=use_int8_w8a8,
                                        use_int8_w8a16=use_int8_w8a16,
                                        use_int4_w4a16=use_int4_w4a16,
                                        use_int4_w4a8=use_int4_w4a8,
                                        dtype=hidden_states.dtype)

    get_config_func = functools.partial(
        try_get_optimal_moe_config,
        w1.shape,
        w2.shape,
        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, w2.shape[1]),
                          device=hidden_states.device,
                          dtype=hidden_states.dtype)
    intermediate_cache1 = cache13[:M * top_k_num * N].view(
        (M, topk_ids.shape[1], N))
    intermediate_cache3 = cache13[:M * top_k_num * w2.shape[1]].view(
        (M, topk_ids.shape[1], w2.shape[1]))

    # 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)

    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.
            intermediate_cache1 = intermediate_cache1[:tokens_in_chunk]
            intermediate_cache2 = intermediate_cache2[:tokens_in_chunk *
                                                      topk_ids.shape[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]

        config['BLOCK_SIZE_M']     = Test['BLOCK_SIZE_M']
        config['num_stages']     = 1
        config['MODE']         = 10000 + Test['MODE1']
        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))

        sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, tokens_positions_per_expert, moe_buf = (
            moe_sorting_ck(
        curr_topk_ids, curr_topk_weights, global_num_experts, w2.shape[1], hidden_states.dtype, config['BLOCK_SIZE_M'], expert_map
        ))
        invoke_fused_moe_kernel(curr_hidden_states,
                                w1,
                                intermediate_cache1,
                                a1_scale,
                                w1_scale,
                                w1_zp,
                                w1_zp_asm,
                                curr_topk_weights,
                                curr_topk_ids,
                                sorted_token_ids,
                                expert_ids,
                                num_tokens_post_padded,
                                False,
                                top_k_num,
                                config,
                                sorted_ids,
                                sorted_expert_ids,
                                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)

        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}")

        if tokens_in_chunk < CHUNK_SIZE and chunk > 0:
            config = get_config_func(tokens_in_chunk, is_bottom=True)
        else:
            config = get_config_func(M, is_bottom=True)

        config['BLOCK_SIZE_M']     = Test['BLOCK_SIZE_M']
        #config['num_stages']     = 1
        
        config['MODE']         = 20000 + Test['MODE2']
        config['num_stages'] = 1
        invoke_fused_moe_kernel(intermediate_cache2,
                                w2,
                                intermediate_cache3,
                                a2_scale,
                                w2_scale,
                                w2_zp,
                                w2_zp_asm,
                                curr_topk_weights,
                                curr_topk_ids,
                                sorted_token_ids,
                                expert_ids,
                                num_tokens_post_padded,
                                True,
                                1,
                                config,
                                sorted_ids,
                                sorted_expert_ids,
                                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)

        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])
        else:
            ops.moe_sum(intermediate_cache3.view(*intermediate_cache3.shape),
                        out_hidden_states[begin_chunk_idx:end_chunk_idx])
    #print("silu in",intermediate_cache2.shape)
    #print("silu out",qintermediate_cache2.shape)
    #print("w1 out",w2.shape)
    #print("w2 out",w2.shape)
    #print("intermediate_cache3 sum",intermediate_cache3.shape)
    #print("out_hidden_states",out_hidden_states.shape)
    return out_hidden_states


# the above tests is for correctness, the following is for performance
###################################################################################################################
def perftest():
    M = [1, 2, 4, 8, 12, 16, 24, 32, 48, 64, 96, 128, 160, 256, 320, 512, 768, 896, 1024, 1280, 1536, 2048, 2304, 2560, 2816, 3072, 4096, 4352, 4608, 4864, 5120, 5376, 5504, 5760, 6144, 7168, 8192, 10240, 12288, 16384, 24576, 32768]
    #M = [2048, 2304, 2560, 2816, 3072]
    N = [256]
    K = [7168]
    E = [256]
    TOPK = [8]
    EP_SIZE = [1]
    DTYPE = [torch.float16, torch.bfloat16]
    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):
        #torch.manual_seed(0)
        #torch.set_printoptions(precision=4, sci_mode=True)
        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, n // group_size // pack_factor),
                                device="cuda",
                                dtype=torch.uint8)
        #triton
        w1_qzeros_triton = torch.empty((e, 2 * n // pack_factor, k // group_size),
                                device="cuda",
                                dtype=torch.uint8)
        w2_qzeros_triton = 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, w_qzeros_triton = \
                    w1, w1_ref, w1_qweight, w1_scales, w1_qzeros, w1_qzeros_triton
            else:
                w, w_ref, w_qweight, w_scales, w_qzeros, w_qzeros_triton = \
                    w2, w2_ref, w2_qweight, w2_scales, w2_qzeros, w2_qzeros_triton
            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_asm = qzeros[:, 1::2] * 16 + qzeros[:, ::2]
                    #triton qzeros
                    qzeros_triton = 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_asm
                w_qzeros_triton[expert_id] = qzeros_triton
        

        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]
            w1_qzeros_triton = w1_qzeros_triton[e_ids]
            w2_qzeros_triton = w2_qzeros_triton[e_ids]            
        else:
            e_map = None

        ## without token topk score calc
        topk_weights, topk_ids = fused_topk(input, score, topk, True)
        
        ## 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_qzeros_triton dtype = {}, shape = {}".format(w1_qzeros_triton.dtype if has_zp else None, w1_qzeros_triton.shape if has_zp else None))
        #print("###### w2_qzeros_triton dtype = {}, shape = {}".format(w2_qzeros_triton.dtype if has_zp else None, w2_qzeros_triton.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))

        #w2_scale = torch.ones_like(w2_scales)
        #topk_weight = torch.ones_like(topk_weights)
        #torch_output = torch_moe_test(input, w1_ref, w2_ref, topk_weights, topk_ids) 

        config = [
            (16, [0, 1, 2, 3], [ 0, 100, 101]),
            (32, [1000, 1001, 1002, 1009], [ 1000, 1100, 1101]),
            (64, [2000], [2000, 2100, 2101]),
            (128, [3000], [3000, 3100, 3101])
        ]
        for mt, modes1, modes2 in config:
            fillvalue = modes2[-1] if len(modes1) > len(modes2) else modes1[-1] 
            print(" fillvalue  ",fillvalue)           
            for mode1, mode2 in zip_longest(modes1, modes2, fillvalue=fillvalue):     
                Test = {}
                Test['BLOCK_SIZE_M']     = mt
                Test['MODE1']         =   mode1
                Test['MODE2']         =   mode2
                triton_output  = fused_experts_impl(input,
                                                w1_qweight,
                                                w2_qweight,
                                                topk_weights,
                                                topk_ids,
                                                #use_int8_w8a8=True,
                                                use_int4_w4a8=True,
                                                #use_int8_w8a16 = weight_bits == 8,
                                                #use_int4_w4a16 = weight_bits == 4,
                                                global_num_experts=e,
                                                expert_map=None,
                                                w1_scale=w1_scales,
                                                w2_scale=w2_scales,
                                                w1_zp=w1_qzeros_triton if has_zp else None,
                                                w2_zp=w2_qzeros_triton if has_zp else None,
                                                w1_zp_asm=w1_qzeros if has_zp else None,
                                                w2_zp_asm=w2_qzeros if has_zp else None,
                                                #per_channel_quant=True,
                                                #block_shape=None)
                                                block_shape=[0, group_size],
                                                Test=Test)

        #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=100, msg=msg)

if __name__ == "__main__":
    perftest()
