# 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,moe_align_block_size,moe_kernel_prepare_input,triton_moe_sum,invoke_fused_moe_kernel_marlin
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 w8a16_marlin_weight_1,w8a16_marlin_weight_2
import triton.language as tl
# import _custom_ops as ops
import aiter
from aiter.test_common import  perftest


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 _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.float16,
    device: str = "cuda",
    per_channel_quant: bool = True
):
    # torch.manual_seed(0) 
    # torch.cuda.manual_seed_all(0)
    # """生成测试数据"""
    input_data = (torch.randn((m, k), device=device) ).to(dtype=out_dtype) /100000
    
    # input_data = torch.ones_like(input_data) /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.randn(-128, 127, (E, n, k), device=device, dtype=torch.float16)
    # weight2 = torch.randn(-128, 127, (E, k2, n2), device=device, dtype=torch.float16)
    # weight1 = torch.empty((E, n, k), device="cuda", dtype=torch.uint8)
    # weight2 = torch.empty((E, k, n2), device="cuda", dtype=torch.uint8)
    # weight1_fp = torch.randn((E, n, k), device=device, dtype=torch.float16)
    # weight2_fp = torch.randn((E, k2, n2), device=device, dtype=torch.float16)
    # weight1_fp = weight1_fp.clamp(-128, 127)
    # weight2_fp = weight2_fp.clamp(-128, 127)
    # weight1 = weight1_fp.to(torch.int8)
    # weight2 = weight2_fp.to(torch.int8)

    weight1 = torch.randint(-127, 127, (E, n, k), device=device, dtype=torch.int8)
    weight2 = torch.randint(-127, 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)
    # input_data = torch.ones_like(input_data)
    # weight1 = torch.ones_like(weight1)
    # weight2 = torch.ones_like(weight2)
    # # -----------------------------------------------------------
    # # 按元素逐个循环：[1,2,3,4,5,6,7,8,1,2,3,...]
    # # -----------------------------------------------------------

    # # 1. 生成 1~8 的循环序列 (长度为 weight1 的总元素数)
    # total_elems_1 = weight1.numel()
    # vals_1 = (torch.arange(total_elems_1, device=weight1.device, dtype=weight1.dtype) % 127)
    # weight1 = vals_1.view(weight1.shape).contiguous()

    # # 2. 对 weight2 做同样的事
    # total_elems_2 = weight2.numel()
    # vals_2 = (torch.arange(total_elems_2, device=weight2.device, dtype=weight2.dtype) % 8) + 1
    # weight2 = vals_2.view(weight2.shape).contiguous()
    #     # 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.ones((E, n, 1), device=device, dtype=torch.float16)
        weight1_scale = torch.randn((E, n, 1), device=device, dtype=torch.float16)
        # weight1_scale = torch.ones_like(weight1_scale)
        #weight2_scale = torch.ones((E, n2, 1), device=device, dtype=torch.float16)
        weight2_scale = torch.randn((E, n2, 1), device=device, dtype=torch.float16)
        # weight2_scale = torch.ones_like(weight2_scale)

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

    # 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_qweight_uint32 = weight1.view(-1).view(torch.uint32)
    # # new_shape = (e, 2 * n, k // 128, 16)  # uint32张量的形状
    # new_shape = (E, n // 16, 16, k // 32, 8)  # uint32张量的形状
    # w1_qweight_uint32_reshaped = w1_qweight_uint32.view(new_shape)
    # w1_qweight_uint32_transposed = w1_qweight_uint32_reshaped.transpose(2, 3).contiguous()
    # new_shape = (E, n // 16, k // 128, 4, 16, 8)
    # w1_new_trans = w1_qweight_uint32_transposed.view(new_shape)
    # w1_marlin = w1_new_trans.transpose(1, 2).contiguous()

    # w2_qweight_uint32 = weight2.view(-1).view(torch.uint32)
    # # new_shape = (e, 2 * n, k // 128, 16)  # uint32张量的形状
    # new_shape = (E, n2 // 16, 16, k2 // 32, 8)  # uint32张量的形状
    # w2_qweight_uint32_reshaped = w2_qweight_uint32.view(new_shape)
    # w2_qweight_uint32_transposed = w2_qweight_uint32_reshaped.transpose(2, 3).contiguous()
    # new_shape = (E, n2 // 16, k2 // 128, 4, 16, 8)
    # w2_new_trans = w2_qweight_uint32_transposed.view(new_shape)
    # w2_marlin = w2_new_trans.transpose(1, 2).contiguous()
    # w1_marlin = w1_marlin.view(-1).view(torch.uint8).view(*weight1.shape)
    # w2_marlin = w2_marlin.view(-1).view(torch.uint8).view(*weight2.shape)
    # # w1_marlin_list = []
    w1_marlin = w8a16_marlin_weight_1(weight1)
    w2_marlin = w8a16_marlin_weight_2(weight2)
    # for i in range(weight1.shape[0]):
    #     w1_marlin_in = W8A16_2_marlin_weight(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 = W8A16_marlin_weight_ours(weight2[i])
    #     w2_marlin_in = W8A16_marlin_weight_ours(weight2[i])


    #     w2_marlin_list.append(w2_marlin_in)
    # w2_marlin = torch.stack(w2_marlin_list, dim=0)
    
    #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": 32,
        "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
        )


# @perftest(num_warmup=5, num_iters=10,testGraph=True)
def perchannel_w8a16_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=False,
            use_int8_w8a16=True,
            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
        False,   # use_int8_w8a8
        True,  # use_int8_w8a16
        False,  # use_int4_w4a16
        per_channel_quant,
        None,   # block_shape
        False   # use_nn_moe
    )
    # if 1:
    #     return output_triton
    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=False,
            use_int8_w8a16=True,
            use_int4_w4a16=False,
            per_channel_quant=per_channel_quant,
            block_shape=None
        )
    # if 1:
    #     return input_gemm2
    # qinput2 = torch.ones_like(qinput2)
    # weight2 = torch.ones_like(weight2)
    # weight_scale2 = torch.ones_like(weight_scale2)


    # print("topk_weights_trit:", topk_weights)
    # print("curr_topk_ids_trit:", topk_ids)
    # print("sorted_token_ids_trit:", sorted_token_ids)
    # print("expert_ids_trit:", expert_ids)


    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
        False,   # use_int8_w8a8
        True,  # use_int8_w8a16
        False,  # use_int4_w4a16
        per_channel_quant,
        None,   # block_shape
        False   # use_nn_moe
    )

    # if 1:
    #     return output_triton2
    
    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 TestW8A16BlockINT8FusedMoE(unittest.TestCase):
    DTYPES = [torch.float16]
    # M = [1]
    M = [1,2,8,16,32,64,128,256,512,1024,8096]
    # M = [4096]
    # M = [1,2,3]
    N = [2048]
    K = [6144]
    # K = [512]
    BM = [16]
    # MODE1 = [183]
    # MODE2 = [54]
    
    E = [128]
    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 _W8A16_block_int8_fused_moe(self, M, N, K, MODE1,MODE2, BM,E, topk, block_size, dtype, seed):
    def _W8A16_block_int8_fused_moe(self, M, N, K, 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, N, K, topk, E, out_dtype = dtype, per_channel_quant=True
        # )
        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
        )
        # 打印形状信息
        # print(f"input_data: {input_data}")
          # 打印形状信息
        # print(f"input_data: {input_data.shape}")
        # print(f"w1_marlin: {w1_marlin.shape}\n")   
        # print(f"w2_marlin: {w2_marlin.shape}")
        # print(f"weight1: {weight1.shape}\n")   
        # print(f"weight2: {weight2.shape}")
        # print(f"weight1_scale: {weight1_scale.shape}\n")  
        # print(f"weight2_scale: {weight2_scale.shape}\n")        
        # topk_weights= torch.ones_like(topk_weights) * 5
        
        with torch.inference_mode():
            sorted_token_ids, expert_ids, num_tokens_post_padded = (
                moe_align_block_size(topk_ids, 16, E, None)
            )
            triton_output = perchannel_w8a16_triton(input_data,weight1,weight2,weight1_scale,weight2_scale,topk_weights, topk_ids,sorted_token_ids, expert_ids, num_tokens_post_padded,dtype,True)            
            moe_output =  moe_c_fused_experts(
                         input_data, 
                         w1_marlin,
                         w2_marlin,
                         topk_weights,
                         topk_ids,
                         MODE1=1,
                         MODE2=1,
                         BM=1,BN=1,BK=1,
                         kloops=1,nloops=1,
                         BN2=1,
                         BK2=1,
                         kloops2=1,
                         nloops2=1,
                         inplace=True,
                         activation="silu",
                         use_fp8_w8a8=False,
                         use_int8_w8a8=False,
                         use_int8_w4a8 = False,
                         use_int8_w8a16=True,
                         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,
                         routed_scaling_factor = 1.0)              
            # 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_W8A16_block_int8_moe(
            #     a, w1, w2, w1_s_cuda, w2_s_cuda, score, topk, block_size
            # )




        # 打印形状信息
        print(f"------M=: ",M)
        print(f"MOE_OUTPUT输出形状: {moe_output.shape}")
        print(f"TRITION输出形状: {triton_output.shape}\n")
        # print("aiter_ms:",aiter_ms)
        # print("triton_ms:",triton_ms)
        # 计算差异张量和掩码
        diff = torch.abs(moe_output - triton_output)
        atol_threshold = 2e-1  # 与断言保持一致的阈值
        match_mask = diff <= atol_threshold  # 匹配的位置（True）和不匹配的位置（False）

        # 展平所有张量以便切片操作
        triton_flat = moe_output.flatten()
        torch_flat = triton_output.flatten()
        diff_flat = diff.flatten()
        match_mask_flat = match_mask.flatten()

        # 分别收集匹配和不匹配的索引
        match_indices = torch.nonzero(match_mask_flat).squeeze().tolist()
        mismatch_indices = torch.nonzero(~match_mask_flat).squeeze().tolist()

        # 确保索引是列表格式（处理只有一个元素的情况）
        if not isinstance(match_indices, list):
            match_indices = [match_indices] if match_indices is not None else []
        if not isinstance(mismatch_indices, list):
            mismatch_indices = [mismatch_indices] if mismatch_indices is not None else []

        # 打印总体统计
        print(f"===== 总体匹配统计 =====")
        print(f"总元素数: {triton_flat.numel()}")
        print(f"匹配元素数: {len(match_indices)} ({len(match_indices)/triton_flat.numel():.2%})")
        print(f"不匹配元素数: {len(mismatch_indices)} ({len(mismatch_indices)/triton_flat.numel():.2%})")
        print(f"最大差异: {diff.max().item():.6f}")
        print(f"平均差异: {diff.mean().item():.6f}\n")

        # 定义打印指定索引数据的函数
        def print_samples(name, indices, max_samples=3, elements_per_sample=20):
            if not indices:
                print(f"没有{name}的数据")
                return
                
            print(f"===== {name}数据示例 =====")
            # 控制最大展示样本数
            num_samples = min(max_samples, len(indices) // elements_per_sample + 1)
            
            for i in range(num_samples):
                start_idx = i * elements_per_sample
                end_idx = start_idx + elements_per_sample
                sample_indices = indices[start_idx:end_idx]
                
                if not sample_indices:
                    break
                    
                print(f"\n样本 {i+1}（索引 {sample_indices[0]}-{sample_indices[-1]}）:")
                print(f"MOE_OUTPUT: {[triton_flat[idx].item() for idx in sample_indices]}")
                print(f"TRITON: {[torch_flat[idx].item() for idx in sample_indices]}")
                print(f"差异:    {[diff_flat[idx].item() for idx in sample_indices]}")

        # # 打印匹配的数据示例
        # print_samples("匹配", match_indices, max_samples=5)

        # # 打印不匹配的数据示例
        # print_samples("不匹配", mismatch_indices, max_samples=3)

        # 保留断言检查
        torch.testing.assert_close(moe_output, triton_output, atol=atol_threshold, rtol=0)            

            
            
        
        # 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(
        TestW8A16BlockINT8FusedMoE.M,
        TestW8A16BlockINT8FusedMoE.N,
        TestW8A16BlockINT8FusedMoE.K,
        # TestW8A16BlockINT8FusedMoE.MODE1,
        # TestW8A16BlockINT8FusedMoE.MODE2,
        TestW8A16BlockINT8FusedMoE.BM,
        TestW8A16BlockINT8FusedMoE.E,
        TestW8A16BlockINT8FusedMoE.TOP_KS,
        TestW8A16BlockINT8FusedMoE.BLOCK_SIZE,
        TestW8A16BlockINT8FusedMoE.DTYPES,
        TestW8A16BlockINT8FusedMoE.SEEDS,
    ))

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


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


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