# 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 aiter.fused_moe_c import moe_c_fused_experts,moe_c_fused_experts_bench,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 aiter.ops.shuffle import w4a8_moe_layout_shuffle_gemm1,w4a8_moe_layout_shuffle_gemm2
import triton.language as tl
# import _custom_ops as ops
import os
from aiter.test_common import checkAllclose
import aiter

def compare_tensors(
    tensor1: torch.Tensor,
    tensor2: torch.Tensor,
    atol: float = 1e-2,
    rtol: float = 1e-2
) -> None:
    """
    比较两个任意维度的PyTorch张量的差异，支持绝对误差和相对误差阈值。
    直接输出详细比较结果，包括每个元素在原始张量中的多维坐标。
    无返回值。

    参数:
        tensor1: 第一个张量（如Triton模型输出）
        tensor2: 第二个张量（如PyTorch模型输出）
        atol: 绝对误差阈值，默认1e-5
        rtol: 相对误差阈值，默认1e-8
    """
    # -------------------------- 1. 输入合法性校验 --------------------------
    if not isinstance(tensor1, torch.Tensor) or not isinstance(tensor2, torch.Tensor):
        raise TypeError("输入必须是PyTorch张量（torch.Tensor）")
    
    if tensor1.shape != tensor2.shape:
        raise ValueError(f"张量形状不匹配！ tensor1形状: {tensor1.shape}, tensor2形状: {tensor2.shape}")
    
    if tensor1.device != tensor2.device:
        tensor2 = tensor2.to(tensor1.device)
        print(f"警告：张量设备不一致，已将tensor2转移到{tensor1.device}")

    # -------------------------- 2. 核心差异计算 --------------------------
    abs_diff = torch.abs(tensor1 - tensor2)
    denom = torch.maximum(torch.abs(tensor1), torch.abs(tensor2))
    rel_diff = abs_diff / (denom + 1e-12)
    match_mask = (abs_diff <= atol) | (rel_diff <= rtol)

    # -------------------------- 3. 展平张量 --------------------------
    tensor1_flat = tensor1.flatten()
    tensor2_flat = tensor2.flatten()
    abs_diff_flat = abs_diff.flatten()
    match_mask_flat = match_mask.flatten()

    # -------------------------- 4. 收集一维索引 --------------------------
    def get_indices_1d(mask: torch.Tensor) -> list:
        indices = torch.nonzero(mask).squeeze(dim=1)
        return indices.tolist() if indices.numel() > 0 else []

    match_indices_1d = get_indices_1d(match_mask_flat)
    mismatch_indices_1d = get_indices_1d(~match_mask_flat)

    # -------------------------- 5. 总体统计信息 --------------------------
    total = tensor1_flat.numel()
    matched = len(match_indices_1d)
    mismatched = len(mismatch_indices_1d)
    match_rate = matched / total if total > 0 else 0.0
    max_abs_diff = abs_diff.max().item() if total > 0 else 0.0
    avg_abs_diff = abs_diff.mean().item() if total > 0 else 0.0

    # -------------------------- 6. 格式化输出 --------------------------
    print("=" * 60)
    print("张量比较结果汇总")
    print("=" * 60)
    print(f"张量形状: {tensor1.shape} | 总元素数: {total}")
    print(f"阈值设置: 绝对误差(atol)={atol:.2e}, 相对误差(rtol)={rtol:.2e}")
    print("-" * 60)
    print(f"匹配元素数: {matched} ({match_rate:.2%})")
    print(f"不匹配元素数: {mismatched} ({1 - match_rate:.2%})")
    print(f"最大绝对差异: {max_abs_diff:.6f}")
    print(f"平均绝对差异: {avg_abs_diff:.6f}")
    print("=" * 60)

    # -------------------------- 7. 输出匹配/不匹配示例 --------------------------
    def print_sample(name: str, indices_1d: list, max_samples: int = 3, elem_per_sample: int = 10) -> None:
        if not indices_1d:
            print(f"\n【{name}样本】无数据")
            return
        
        print(f"\n【{name}样本】（最多展示{max_samples}组，每组{elem_per_sample}个元素）")
        print("-" * 50)
        
        num_samples = min(max_samples, (len(indices_1d) + elem_per_sample - 1) // elem_per_sample)
        
        for i in range(num_samples):
            start = i * elem_per_sample
            end = start + elem_per_sample
            sample_indices_1d = indices_1d[start:end]
            
            # <<< 关键改动: 使用 torch.unravel_index 转换为多维坐标 >>>
            # 此函数能处理任意维度
            sample_coords = torch.unravel_index(torch.tensor(sample_indices_1d), tensor1.shape)
            # 将结果从张量元组转换为坐标元组列表
            sample_coords_list = list(zip(*[coord.tolist() for coord in sample_coords]))
            
            print(f"\n第{i+1}组:")
            print(f"  原始多维坐标: {sample_coords_list}")
            print(f"  tensor1: {[round(tensor1_flat[idx].item(), 6) for idx in sample_indices_1d]}")
            print(f"  tensor2: {[round(tensor2_flat[idx].item(), 6) for idx in sample_indices_1d]}")
            print(f"  绝对差异: {[round(abs_diff_flat[idx].item(), 6) for idx in sample_indices_1d]}")

    print_sample("匹配", match_indices_1d, max_samples=2)
    print_sample("不匹配", mismatch_indices_1d, max_samples=3)
    print("\n" + "=" * 60)

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 pack_int4_on_k(weight):
    # 假设 weight 形状为 (E, n2, k2)
    E, n2, k2 = weight.shape
    assert k2 % 2 == 0, "K 维度必须是偶数才能成对打包"
    
    # 1. 将数据钳制在 int4 有符号范围内 [-8, 7]
    # 如果你的数据已经是量化好的，可以跳过 clamp
    # w_clamped = torch.clamp(weight, -8, 7)
    
    # 2. 转换为无符号字节处理位运算，防止符号位干扰
    # 我们只需要每个数字的后 4 位 (0x0F)
    w_unsigned = weight.to(torch.uint8) & 0x0F
    

    original_shape = w_unsigned.shape
    temp_view = w_unsigned.view(*original_shape[:-1], -1, 8)
    selected_low = temp_view[..., :4]
    w_low = selected_low.reshape(*original_shape[:-1], -1)
    selected_high = temp_view[..., 4:]
    w_high = selected_high.reshape(*original_shape[:-1], -1)



    # 3. 分离偶数位置和奇数位置的 K
    # low_k 占据 0-3 位, high_k 准备占据 4-7 位
    # low_k = w_unsigned[:, :, 0::2]       # 形状 (E, n2, k2/2)
    # high_k = w_unsigned[:, :, 1::2]      # 形状 (E, n2, k2/2)
    
    # 4. 打包：将 high_k 左移 4 位，然后与 low_k 合并
    packed_weight = (w_low << 4) | w_high
    
    # 5. 转回 int8 存储
    return packed_weight.to(torch.int8)

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.rand((m, k), device=device) ).to(dtype=out_dtype)  / 10000
    # input_data = torch.ones_like(input_data) 
   

    # weight1_ori = torch.ones((E, n, k ), device=device, dtype=torch.int8) * 1
    weight1_ori = torch.randint(0,15,(E, n, k ), device=device, dtype=torch.int8)
    result1 = torch.where(weight1_ori > 7, weight1_ori - 16, weight1_ori)
    # result1 = torch.ones((E, n, k ), device=device, dtype=torch.int8) * 1
    weight1 = pack_int4_on_k(weight1_ori)

   
    
    # weight2_ori = torch.ones((E, n2, k2 ), device=device, dtype=torch.int8) * 1
    weight2_ori = torch.randint(0,15,(E, n2, k2 ), device=device, dtype=torch.int8)
    result = torch.where(weight2_ori > 7, weight2_ori - 16, weight2_ori)
    # result = torch.ones((E, n2, k2 ), device=device, dtype=torch.int8) * 1
    weight2 = pack_int4_on_k(weight2_ori)
    
    

    if per_channel_quant:
        # weight1_scale = torch.randn((E, n, 1), device=device, dtype=torch.float32)
        weight1_scale = torch.randn((E, n, 1), dtype=torch.float32, device=device)

        # weight1_scale = torch.ones_like(weight1_scale)
        # weight2_scale = torch.randn((E, n2, 1), device=device, dtype=torch.float32)
        weight2_scale = torch.randn((E, n2, 1), dtype=torch.float32, device=device)

        # weight2_scale = torch.ones_like(weight2_scale) 

    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) / 10000
    # # 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)
    #     # weight2_scale = torch.randn((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.ones((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_list = []
    for i in range(weight1.shape[0]):
        w1_marlin_in = w4a8_moe_layout_shuffle_gemm2(weight1[i])
        w1_marlin_list.append(w1_marlin_in)
    w1_marlin = torch.stack(w1_marlin_list, dim=0)
    


    w2_marlin_list = []
    for i in range(weight2.shape[0]):
        # w2_marlin_in = w8a8_marlin_weight_ours(weight2[i])
        w2_marlin_in = w4a8_moe_layout_shuffle_gemm2(weight2[i])
        w2_marlin_list.append(w2_marlin_in)
    w2_marlin = torch.stack(w2_marlin_list, dim=0)


    # storage = w2_marlin.untyped_storage()

    # # 打印物理内存中的前 10 个元素（按其原始 Dtype）
    # # 注意：这会完全忽略张量的 shape 和 stride，只按地址顺序读
    # print(f"Physical sequence (first 10): {storage[:256]}")


    #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,result1,result



@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
    )
    
    # print("*******************************************************output_triton")
    # print(output_triton)
    

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


    # input_gemm2 = torch.ones_like(input_gemm2)
    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
        )
    
    # print("golden 量化后：")
    # print(qinput2)
    
    
    # print("**********************************************qinput2")
    # print(qinput2)
    
    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,3,4,5,6,7,8,9,10,11,12,13,14,15,16,32,64,128,256,512,1024,2048,4096,8192,16384,32768]
    # M = [256]
    M = [4096]
    N = [2048]
    K = [7168]
    # 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 ,weight1_ori,weight2_ori= _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():
            for i in range(10):
                out=  moe_c_fused_experts(input_data, w1_marlin, w2_marlin, topk_weights, topk_ids,MODE1,  MODE2,BM,
                            inplace = False, activation = "silu", use_fp8_w8a8 = False, use_int8_w8a8 = False,use_int8_w4a8 = 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
            # )
            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_ori,weight2_ori,weight1_scale,weight2_scale,topk_weights, topk_ids,sorted_token_ids, expert_ids, num_tokens_post_padded,dtype,True)
            # print("**********************************ms")
            # print(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) 
        checkAllclose(ref_out, out, rtol=0.01, atol=100)

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