# SPDX-License-Identifier: Apache-2.0

# Adapted from https://github.com/sgl-project/sglang/pull/3730
import itertools
import unittest
import triton
import torch

# from vllm.model_executor.layers.activation import SiluAndMul
from aiter.fused_moe_c  import moe_c_fused_experts,moe_c_fused_experts_bench,fused_experts_impl_channelwise_w8a8,moe_align_block_size,moe_kernel_prepare_input
from typing import Any, Callable, Dict, List, Optional, Tuple, Type
import numpy as np
import torch.nn.functional as F
# from w8a8_quant import w8a8_2_marlin_weight,w8a8_marlin_weight_ours
from aiter.ops.shuffle import moe_layout_shuffle_gemm1,moe_layout_shuffle_gemm2
import triton.language as tl
# import _custom_ops as ops
import aiter
from aiter.test_common import  perftest

@perftest(num_warmup=5, num_iters=10,testGraph=True)
def moe_impl(input_data, weight1, weight2, w1_marlin,w2_marlin, topk_weights, topk_ids,BM,BN,BK,kloops,nloops,BN2,
                            BK2,kloops2,nloops2,
                            inplace, activation, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16,
                            use_int4_w4a16, use_int4_w4a16_base, global_num_experts, expert_map,
                            w1_scale, w2_scale, w1_zp, w2_zp, a1_scale,
                            a2_scale, block_shape):
    return moe_c_fused_experts(input_data, weight1, weight2, w1_marlin,w2_marlin, topk_weights, topk_ids,BM,BN,BK,kloops,nloops,BN2,
                            BK2,kloops2,nloops2,
                            inplace, activation, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16,
                            use_int4_w4a16, use_int4_w4a16_base, global_num_experts, expert_map,
                            w1_scale, w2_scale, w1_zp, w2_zp, a1_scale,
                            a2_scale, block_shape)

def generate_unique_int_tensor(m, topk, low=0, high=256):
    """生成[m, topk]的整数张量，每行元素不重复"""
    if topk > (high - low):
        raise ValueError(f"topk ({topk}) 不能超过范围 ({high - low})")
    
    tensor = torch.zeros((m, topk), dtype=torch.int)
    for i in range(m):
        row = torch.randperm(high - low)[:topk] + low
        tensor[i] = row
    return tensor

def _fused_moe_helper(
    m: int, n: int, k: int, n2: int, k2: int, topk: int, E: int, ep_size: int = 1,
    out_dtype:  torch.dtype = torch.bfloat16,
    device: str = "cuda",
    per_channel_quant: bool = True
):
    """生成测试数据"""
    
    input_data = (torch.randn((m, k), device=device) ).to(dtype=out_dtype) / 100000
    
    # weight1 = to_int8(torch.randn((E, n, k), device=device) * 5)
    # weight2 = to_int8(torch.randn((E, n2, k2), device=device) * 5)
    weight1 = torch.randint(-128, 127, (E,n ,k), device=device, dtype=torch.int8)
    weight2 = torch.randint(-128, 127, (E,n2 ,k2), device=device, dtype=torch.int8)
    # input_data = torch.ones_like(input_data)
    # weight1 = torch.ones_like(weight1)
    # weight2 = torch.ones_like(weight2)
    # B, H, W = weight1.shape  # 从已有张量获取维度

    # # # 1. 生成列与列不同的 1D 向量（长度=W，每个元素代表一列的值）
    # # # 这里用 0~W-1 举例，可替换为任意你需要的不同值（如随机数、等差序列等）
    # col_values = torch.arange(W, dtype=torch.int32, device='cpu').clamp(max=127) -60
    # # 2. 转为int8类型（此时所有值≤127，无溢出）
    # col_values = col_values.to(dtype=torch.int8)

    # # # # 2. 扩展为 (1, 1, W)，利用广播机制让每个二维切片的列值相同
    # col_values_expanded = col_values[None, None, :]  # shape: (1, 1, W)

    # # # # 3. 赋值：广播到 (B, H, W)，满足「每列内相同、列与列不同、每个二维张量相同」
    # weight1 = col_values_expanded.expand(B, H, W).clone().contiguous().to('cuda')
    

    if per_channel_quant:
        weight1_scale = torch.randn((E, n, 1), device=device, dtype=torch.float32)
        weight2_scale = torch.randn((E, n2, 1), device=device, dtype=torch.float32)
    else:
        # 这里可以添加block-wise量化的支持
        weight1_scale = torch.randn((E, n, 1), device=device, dtype=torch.float32)
        weight2_scale = torch.randn((E, n2, 1), device=device, dtype=torch.float32)
    

    # weight1_scale = torch.ones_like(weight1_scale)
    # weight2_scale = torch.ones_like(weight2_scale)

    # topk_weights [m, topk]
    tensor = torch.rand((m, topk), device=device)
    topk_weights = F.softmax(tensor, dim=1)
    # topk_weights = torch.ones_like(topk_weights)
    
    # topk_ids [m, topk]
    topk_ids = generate_unique_int_tensor(m=m, topk=topk, low=0, high=E).cuda()

    # """生成测试数据"""
    # input_data = (torch.ones((m, k), device=device) ).to(dtype=out_dtype)
    # # weight1 = to_int8(torch.randn((E, n, k), device=device) * 5)
    # # weight2 = to_int8(torch.randn((E, n2, k2), device=device) * 5)
    # # weight1 = torch.randint(-128, 127, (E,n ,k), device=device, dtype=torch.int8)
    # # weight2 = torch.randint(-128, 127, (E,n2 ,k2), device=device, dtype=torch.int8)
    # weight1 = torch.ones((E,n ,k), device=device, dtype=torch.int8)
    # weight2 = torch.ones((E,n2 ,k2), device=device, dtype=torch.int8)

    # #weight1 = torch.ones_like(weight1)

    # if per_channel_quant:
    #     weight1_scale = torch.ones((E, n, 1), device=device, dtype=torch.float32)
    #     weight2_scale = torch.ones((E, n2, 1), device=device, dtype=torch.float32)
    # else:
    #     # 这里可以添加block-wise量化的支持
    #     weight1_scale = torch.ones((E, n, 1), device=device, dtype=torch.float32)
    #     weight2_scale = torch.ones((E, n2, 1), device=device, dtype=torch.float32)
    
    # # topk_weights [m, topk]
    # tensor = torch.rand((m, topk), device=device)
    # topk_weights = F.softmax(tensor, dim=1)
    
    # # topk_ids [m, topk]
    # topk_ids = generate_unique_int_tensor(m=m, topk=topk, low=0, high=E).cuda()

    # weight重排
  

    w1_marlin = moe_layout_shuffle_gemm1(weight1)
    w2_marlin = moe_layout_shuffle_gemm2(weight2)
    #return input_data, weight1, weight2, weight1_scale, weight2_scale, topk_weights, topk_ids
    return input_data, w1_marlin, w2_marlin, weight1, weight2, weight1_scale, weight2_scale, topk_weights, topk_ids



@triton.jit
def fused_moe_kernel(
    # Pointers to matrices
    a_ptr,
    b_ptr,
    c_ptr,
    a_scale_ptr,
    b_scale_ptr,
    topk_weights_ptr,
    sorted_token_ids_ptr,
    expert_ids_ptr,
    num_tokens_post_padded_ptr,
    # Matrix dimensions
    N,
    K,
    EM,
    num_valid_tokens,
    # The stride variables represent how much to increase the ptr by when
    # moving by 1 element in a particular dimension. E.g. `stride_am` is
    # how much to increase `a_ptr` by to get the element one row down
    # (A has M rows).
    stride_am,
    stride_ak,
    stride_be,
    stride_bk,
    stride_bn,
    stride_cm,
    stride_cn,
    stride_asm,
    stride_ask,
    stride_bse,
    stride_bsk,
    stride_bsn,
    # Block size for block-wise quantization
    group_n: tl.constexpr,
    group_k: tl.constexpr,
    # Meta-parameters
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
    MUL_ROUTED_WEIGHT: tl.constexpr,
    top_k: tl.constexpr,
    compute_type: tl.constexpr,
    use_fp8_w8a8: tl.constexpr,
    use_int8_w8a8: tl.constexpr,
    use_int8_w8a16: tl.constexpr,
    per_channel_quant: tl.constexpr,
    delta: tl.constexpr,
):
    """
    Implements the fused computation for a Mixture of Experts (MOE) using
    token and expert matrices.
    """
    # Map program ids `pid` to the block of C it should compute.
    pid = tl.program_id(axis=0)
    if GROUP_SIZE_M ==1:
        num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
        pid_m = pid // num_pid_n
        pid_n = pid % num_pid_n
    else:
        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.
    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
    if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
        return
    offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(
        tl.int64)
    offs_token = tl.load(sorted_token_ids_ptr + offs_token_id)
    token_mask = offs_token < num_valid_tokens

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

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

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

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

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

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


def invoke_fused_moe_kernel_int8(A: torch.Tensor,
                            B: torch.Tensor,
                            C: torch.Tensor,
                            A_scale: Optional[torch.Tensor],
                            B_scale: Optional[torch.Tensor],
                            B_zp: Optional[torch.Tensor],
                            topk_weights: Optional[torch.Tensor],
                            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,
                            compute_type: tl.dtype,
                            use_fp8_w8a8: bool,
                            use_int8_w8a8: bool,
                            use_int8_w8a16: bool,
                            use_int4_w4a16: bool,
                            per_channel_quant: bool,
                            block_shape: Optional[List[int]] = None,
                            use_nn_moe: Optional[bool]=False) -> None:
    assert topk_weights is not None or not mul_routed_weight
    assert topk_weights is None or topk_weights.stride(1) == 1
    assert sorted_token_ids.stride(0) == 1
    assert use_int8_w8a8 is True
    
    M = A.shape[0]
    num_tokens = M * top_k

    triton_config = {   "BLOCK_SIZE_M": 64,
        "BLOCK_SIZE_N": 64,
        "BLOCK_SIZE_K": 64,
        "GROUP_SIZE_M": 1,
        "kpack": 1,
        # "num_stages": 0, #triton后端更新 =0时错误改为等于2
        "num_stages": 2, #triton后端更新 =0时错误改为等于2
        "num_warps": 2
    }
    triton_config["BLOCK_SIZE_M"] = 16
    
    EM = sorted_token_ids.shape[0]
    if A.shape[0] < triton_config["BLOCK_SIZE_M"]:
        # optimize for small batch_size.
        # We assume that top_ids of each token is unique, so
        # so num_valid_experts <= batch_size <= BLOCK_SIZE_M,
        # and we can skip some invalid blocks.
        EM = min(sorted_token_ids.shape[0],
                 A.shape[0] * top_k * triton_config['BLOCK_SIZE_M'])
    grid = lambda META: (triton.cdiv(EM, META['BLOCK_SIZE_M']) * triton.cdiv(
        B.shape[1] if not use_nn_moe else B.shape[2], META['BLOCK_SIZE_N']), )

    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] if not use_nn_moe else B.shape[2],
            A.shape[1],
            EM,
            num_tokens,
            A.stride(0),
            A.stride(1),
            B.stride(0),
            B.stride(2) if not use_nn_moe else B.stride(1),
            B.stride(1) if not use_nn_moe else B.stride(2),
            C.stride(1),
            C.stride(2),
            A_scale.stride(0)
            if A_scale is not None and A_scale.ndim == 2 else 0,
            A_scale.stride(1)
            if A_scale is not None and A_scale.ndim == 2 else 0,
            B_scale.stride(0)
            if B_scale is not None and B_scale.ndim >= 2 else 0,
            B_scale.stride(2)
            if B_scale is not None and B_scale.ndim == 3 else 0,
            B_scale.stride(1)
            if B_scale is not None and B_scale.ndim >= 2 else 0,
            0 if block_shape is None else block_shape[0],
            0 if block_shape is None else block_shape[1],
            MUL_ROUTED_WEIGHT=mul_routed_weight,
            top_k=top_k,
            compute_type=compute_type,
            use_fp8_w8a8=use_fp8_w8a8,
            use_int8_w8a8=use_int8_w8a8,
            use_int8_w8a16=use_int8_w8a16,
            per_channel_quant=per_channel_quant,
            delta=1,
            **triton_config
        )

def perchannel_w8a8_triton(input,
        weight1,
        weight2,
        weight_scale1,
        weight_scale2,
        topk_weights,
        topk_ids,  
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        out_dtype,
        per_channel_quant
        ):

    if out_dtype == torch.bfloat16:
        compute_type = tl.bfloat16
    elif out_dtype == torch.float16:
        compute_type = tl.float16
    elif out_dtype == torch.float32:
        compute_type = tl.float32
    else:
        raise ValueError(f"Unsupported compute_type: {out_dtype}")

    n1 = weight1.shape[1]
    top_k = topk_weights.shape[1]
    m = input.shape[0]


    output_shape = (m, top_k, n1)
    output_triton = torch.zeros(output_shape, device=input.device, dtype=input.dtype)
    output_shape2 = (m  , top_k, weight2.shape[1])
    output_triton2 = torch.zeros(output_shape2, device=input.device, dtype=input.dtype)
    input_gemm2 = torch.empty((m * top_k, n1 // 2), device=input.device, dtype=input.dtype)


    qinput1, qa_scale1 = moe_kernel_prepare_input(
            A=input,
            B=weight1,
            A_scale=None,
            B_scale=weight_scale1,
            use_fp8_w8a8=False,
            use_int8_w8a8=True,
            use_int8_w8a16=False,
            use_int4_w4a16=False,
            per_channel_quant=per_channel_quant,
            block_shape=None
        )
    
    invoke_fused_moe_kernel_int8(
        qinput1,
        weight1,
        output_triton,
        qa_scale1,
        weight_scale1,
        None,  # B_zp
        topk_weights,
        topk_ids,  # topk_ids (不需要)
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        False,  # mul_routed_weight
        top_k ,
        compute_type,
        False,  # use_fp8_w8a8
        True,   # use_int8_w8a8
        False,  # use_int8_w8a16
        False,  # use_int4_w4a16
        per_channel_quant,
        None,   # block_shape
        False   # use_nn_moe
    )



    aiter.moe_c_silu_and_mul(input_gemm2, output_triton.view(-1, n1))

    qinput2, qa_scale2 = moe_kernel_prepare_input(
            A=input_gemm2,
            B=weight2,
            A_scale=None,
            B_scale=weight_scale2,
            use_fp8_w8a8=False,
            use_int8_w8a8=True,
            use_int8_w8a16=False,
            use_int4_w4a16=False,
            per_channel_quant=per_channel_quant,
            block_shape=None
        )

    
    invoke_fused_moe_kernel_int8(
        qinput2,
        weight2,
        output_triton2,
        qa_scale2,
        weight_scale2,
        None,  # B_zp
        topk_weights,
        topk_ids,  # topk_ids (不需要)
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        True,  # mul_routed_weight
        1,
        compute_type,
        False,  # use_fp8_w8a8
        True,   # use_int8_w8a8
        False,  # use_int8_w8a16
        False,  # use_int4_w4a16
        per_channel_quant,
        None,   # block_shape
        False   # use_nn_moe
    )


    out_hidden_states = torch.empty_like(input)
    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
        
    mode_use_triton_moe_sum = False

    if mode_use_triton_moe_sum:
        triton_moe_sum(output_triton2.view(*output_triton2.shape),
                        out_hidden_states)
    else:
        aiter.moe_c_moe_sum(output_triton2.view(*output_triton2.shape),
                    out_hidden_states,topk_ids)
    # print("**************************************triton")
    # print(out_hidden_states)
    return out_hidden_states



class TestW8A8BlockINT8FusedMoE(unittest.TestCase):
    DTYPES = [torch.bfloat16]
    M = [1,2,4,5,6,7,8,9,10,11,12,13,14,15,16,32,64,128,256,512,1024,1024,2048,4096,8192,16384,32768]
    M = [1024,2048,4096,8192,16384,32768]
    M = [1]
    N = [128]
    K = [6144]
    # K = [512]
    BM = [16]
    MODE1 = [183]
    MODE2 = [54]
    
    E = [256]
    TOP_KS = [8]
    # TOP_KS = [1]

    BLOCK_SIZE = [[]]
    SEEDS = [0]

    @classmethod
    def setUpClass(cls):
        if not torch.cuda.is_available():
            raise unittest.SkipTest("CUDA is not available")
        torch.set_default_device("cuda")

    def _w8a8_block_int8_fused_moe(self, M, N, K, MODE1,MODE2, BM,E, topk, block_size, dtype, seed):
        
        torch.manual_seed(seed)
        factor_for_scale = 1e-2

        finfo = torch.finfo(torch.float8_e4m3fn)
        fp8_min = finfo.min
        fp8_max = finfo.max
        int8_min = -127
        int8_max = 127
        


        input_data, w1_marlin, w2_marlin, weight1, weight2, weight1_scale, weight2_scale, topk_weights, topk_ids = _fused_moe_helper(
        M, 2*N, K, K, N, topk, E, out_dtype = dtype, per_channel_quant=True
    )



        w1_marlin = w1_marlin.view(*weight1.shape)
        w2_marlin = w2_marlin.view(*weight2.shape)

        with torch.inference_mode():
            sorted_token_ids, expert_ids, num_tokens_post_padded = (
                moe_align_block_size(topk_ids, 16, E, None)
            )
            ref_out = perchannel_w8a8_triton(input_data,weight1,weight2,weight1_scale,weight2_scale,topk_weights, topk_ids,sorted_token_ids, expert_ids, num_tokens_post_padded,dtype,True)
            
            out=  moe_c_fused_experts(input_data, w1_marlin,w2_marlin, topk_weights, topk_ids,MODE1,  MODE2,BM,
                              inplace = True, activation = "silu", use_fp8_w8a8 = False, use_int8_w8a8 = True, use_int8_w8a16 = False,
                              use_int4_w4a16 = False, use_int4_w4a16_base = False, global_num_experts = E, expert_map = None,
                              w1_scale = weight1_scale,w2_scale =  weight2_scale, w1_zp = None, w2_zp = None, a1_scale = None,
                              a2_scale = None, block_shape = None )
            
            # ref_out = fused_moe(
            #     a, w1, w2, score, topk,
            #     renormalize=False,
            #     use_int8_w8a16=True,
            #     global_num_experts=E,
            #     w1_scale=w1_s_cuda,
            #     w2_scale=w2_s_cuda,
            #     block_shape=block_size,
            #     BM = BM,BN = BN,BK =BK,mloops= mloops, nloops = nloops,
            #     BN2 = BN2,BK2 =BK2,mloops2= mloops2, nloops2 = nloops2
            # )
            # ref_out = torch_w8a8_block_int8_moe(
            #     a, w1, w2, w1_s_cuda, w2_s_cuda, score, topk, block_size
            # )
            
            # print("********************************ms",ms)
            
            
        
        print(f"M={M}: Triton has NaN={torch.isnan(out.cpu()).any().item()}\n PyTorch has NaN={torch.isnan(ref_out.cpu()).any().item()}")
        # print(out)
        out_flat = out.flatten()[:100]
        ref_out_flat = ref_out.flatten()[:100]
        for i in range(100):
            print(f"out_flat[{i}]={out_flat[i]} vs ref_out_flat[{i}]={ref_out_flat[i]}")
        # for i in range(100):
        #     print(f"out_flat[{i}]={out_flat[i]} ")
        # torch.testing.assert_close(out, ref_out, atol=2e-1, rtol=0) 
        

        self.assertTrue(
            torch.mean(torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)))
            / torch.mean(torch.abs(ref_out.to(torch.float32)))
            # < 0.02
            < 0.02
        )
        # compare_tensors(out,ref_out)
        
        
      
            

# 动态为每个参数组合生成独立的测试方法
def generate_test_methods():
    # 将参数组合转换为列表（可索引访问）
    params_list = list(itertools.product(
        TestW8A8BlockINT8FusedMoE.M,
        TestW8A8BlockINT8FusedMoE.N,
        TestW8A8BlockINT8FusedMoE.K,
        TestW8A8BlockINT8FusedMoE.MODE1,
        TestW8A8BlockINT8FusedMoE.MODE2,
        TestW8A8BlockINT8FusedMoE.BM,
        TestW8A8BlockINT8FusedMoE.E,
        TestW8A8BlockINT8FusedMoE.TOP_KS,
        TestW8A8BlockINT8FusedMoE.BLOCK_SIZE,
        TestW8A8BlockINT8FusedMoE.DTYPES,
        TestW8A8BlockINT8FusedMoE.SEEDS,
    ))

    
    for idx, params in enumerate(params_list):
        M = params[0]  # 用M值作为方法名的一部分
        
        # 定义测试方法（闭包，绑定当前参数）
        def test_method(self, params=params):
            self._w8a8_block_int8_fused_moe(*params)
        
        # 方法名必须以test_开头，unittest才会识别为测试方法
        method_name = f"test_moe_M_{M}_case_{idx}"
        
        # 将方法绑定到测试类
        setattr(TestW8A8BlockINT8FusedMoE, method_name, test_method)


# 调用函数生成所有测试方法
generate_test_methods()
    
        
        
            
        


if __name__ == "__main__":
    unittest.main(verbosity=0)
