# Note: will remove this file after BU, just for verify. !!!!!!!!!!!!!!!

# SPDX-License-Identifier: MIT
 
import torch
import pytest
from typing import Dict, Optional
import aiter
from aiter.ops.triton.moe_op import (
    fused_moe as triton_moe,
    moe_set_use_persistent_kernel as triton_moe_set_use_persistent_kernel,
)
from aiter import per_token_quant_hip, per_block_quant_wrapper

# from aiter.ops.triton.moe_op_e2e import (
#     e2e_moe as triton_e2e_moe,
#     moe_set_use_persistent_kernel as triton_e2e_moe_set_use_persistent_kernel,
# )
# from aiter.ops.triton.moe_op_silu_fused import (
#     fused_moe_silu as triton_moe_silu,
#     moe_set_use_persistent_kernel as triton_moe_silu_set_use_persistent_kernel,
# )
# from aiter.ops.triton.moe_op_gelu import (
#     fused_moe_gelu as triton_moe_gelu,
#     moe_set_use_persistent_kernel as triton_moe_gelu_set_use_persistent_kernel,
# )

from aiter.ops.triton.utils.moe_config_utils import get_optimal_moe_config_func
from aiter.ops.triton.utils.types import torch_to_triton_dtype

DEBUG_MODE = False


def torch_silu_and_mul_ref(input):
    """
    Performs the SiLU activation on the first half of the input tensor and
    multiplies it element-wise with the second half.
    Args:
        input (torch.Tensor): Input tensor of shape [..., 2 * d].
        param (float): Parameter for the SiLU activation function.
    Returns:
        torch.Tensor: Output tensor of shape [..., d].
    """
    dtype = input.dtype
    d = input.size(-1) // 2
    A, B = input[:, :d], input[:, d:]

    silu_A = A / (1.0 + torch.exp(-A.float()))

    output = silu_A * B

    return output.to(dtype)

def native_w8a8_block_matmul(A: torch.Tensor, B: torch.Tensor,
                             As: torch.Tensor, Bs: torch.Tensor, block_size,
                             output_dtype):
    """This function performs matrix multiplication with block-wise
    quantization using native torch.
    It is agnostic to the input data type and can be used for both int8 and
    fp8 data types.

    It takes two input tensors `A` and `B` (int8) with scales `As` and
    `Bs` (float32).
    The output is returned in the specified `output_dtype`.
    """
    A = A.to(torch.float32)
    B = B.to(torch.float32)
    assert A.shape[-1] == B.shape[-1]
    assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2
    assert len(block_size) == 2
    block_n, block_k = block_size[0], block_size[1]
    assert (A.shape[-1] + block_k - 1) // block_k == As.shape[-1]
    assert A.shape[:-1] == As.shape[:-1]

    M = A.numel() // A.shape[-1]
    N, K = B.shape
    origin_C_shape = A.shape[:-1] + (N, )
    A = A.reshape(M, A.shape[-1])
    As = As.reshape(M, As.shape[-1])
    n_tiles = (N + block_n - 1) // block_n
    k_tiles = (K + block_k - 1) // block_k
    assert n_tiles == Bs.shape[0]
    assert k_tiles == Bs.shape[1]

    C_shape = (M, N)
    C = torch.zeros(C_shape, dtype=torch.float32, device=A.device)

    A_tiles = [
        A[:, i * block_k:min((i + 1) * block_k, K)] for i in range(k_tiles)
    ]
    B_tiles = [[
        B[
            j * block_n:min((j + 1) * block_n, N),
            i * block_k:min((i + 1) * block_k, K),
        ] for i in range(k_tiles)
    ] for j in range(n_tiles)]
    C_tiles = [
        C[:, j * block_n:min((j + 1) * block_n, N)] for j in range(n_tiles)
    ]
    As_tiles = [As[:, i:i + 1] for i in range(k_tiles)]

    for i in range(k_tiles):
        for j in range(n_tiles):
            a = A_tiles[i]
            b = B_tiles[j][i]
            c = C_tiles[j]
            s = As_tiles[i] * Bs[j][i]
            c[:, :] += torch.matmul(a, b.t()) * s

    C = C.reshape(origin_C_shape).to(output_dtype)
    return C

def native_per_token_group_quant_int8(x,
                                      group_size,
                                      eps=1e-10,
                                      dtype=torch.bfloat16):
    """Function to perform per-token-group quantization on an input tensor
    `x` using native torch.

    It converts the tensor values into int8 values and returns the
    quantized tensor along with the scaling factor used for quantization.
    """
    assert (x.shape[-1] % group_size == 0
            ), "the last dimension of `x` cannot be divisible by `group_size`"
    assert x.is_contiguous(), "`x` is not contiguous"

    iinfo = torch.iinfo(torch.int8)
    int8_min = iinfo.min
    int8_max = iinfo.max

    x_ = x.reshape(x.numel() // group_size, group_size)
    # Use float32 for scale calculation for stability
    amax = x_.abs().max(dim=-1,
                        keepdim=True)[0].clamp(min=eps).to(dtype)
    x_s = amax / int8_max
    x_q = (x_.to(torch.float32) / x_s).round().clamp(
        min=int8_min, max=int8_max).to(torch.int8)  # Round before clamping
    x_q = x_q.reshape(x.shape)
    x_s = x_s.reshape(x.shape[:-1] + (x.shape[-1] // group_size, ))

    return x_q, x_s

def native_per_token_group_quant_fp8(x,
                                     group_size,
                                     eps=1e-10,
                                     dtype=torch.float8_e4m3fn):
    """Function to perform per-token-group quantization on an input tensor
    `x` using native torch."""
    assert x.shape[-1] % group_size == 0, ("the last dimension of `x` cannot "
                                           "be divisible by `group_size`")
    assert x.is_contiguous(), "`x` is not contiguous"

    finfo = torch.finfo(dtype)
    fp8_min = finfo.min
    fp8_max = finfo.max

    x_ = x.reshape(x.numel() // group_size, group_size)
    amax = x_.abs().max(dim=-1,
                        keepdim=True)[0].clamp(min=eps).to(torch.float32)
    x_s = amax / fp8_max
    x_q = (x_ / x_s).clamp(min=fp8_min, max=fp8_max).to(dtype)
    x_q = x_q.reshape(x.shape)
    x_s = x_s.reshape(x.shape[:-1] + (x.shape[-1] // group_size, ))

    return x_q, x_s

def torch_moe_ref(
    a,
    b,
    c,
    a_scale,
    b_scale,
    b_zp,
    group_size,
    topk_ids,
    topk_weights,
    routed_weight,
    sorted_token_ids,
    expert_ids,
    num_tokens_post_padded,
    dtype,
    fp8_w8a8,
    int8_w8a16,
    int4_w4a16,
    gelu=False,
):
    if fp8_w8a8:
        a, _, a_scale = quantize_fp8(a)

    M, top_k, N = c.shape
    _, K = a.shape

    if int4_w4a16:
        b = torch.repeat_interleave(b, repeats=2, dim=2)  # Expand to (E, N, K)
        b_shifter = ((torch.arange(0, K, device=b.device) % 2) * 4)[None, None, :]
        b = (b >> b_shifter) & 0xF
        b_scale = torch.repeat_interleave(
            b_scale, repeats=group_size, dim=2
        )  # (E, N, K)
        if b_zp is not None:
            b_zp = torch.repeat_interleave(
                b_zp, repeats=2, dim=1
            )  # (E,N//2,K//group_size) -> (E, N, K // group_size)
            b_zp = torch.repeat_interleave(
                b_zp, repeats=group_size, dim=2
            )  # (E,N,K//group_size) -> (E, N, K)
            b_zp_shifter = ((torch.arange(0, N, device=b.device) % 2) * 4)[
                None, :, None
            ]
            b_zp = (b_zp >> b_zp_shifter) & 0xF
            b = (b - b_zp) * b_scale
        else:
            b = (b - 8) * b_scale

    # Repeat a -> (M, top_k, K)
    a_expanded = a.unsqueeze(1).repeat(1, top_k, 1)
    # (M, top_k, N, K)
    if fp8_w8a8:
        b_indexed = b.half()[topk_ids]
    else:
        b_indexed = b[topk_ids]

    c = torch.einsum("mek,menk->men", a_expanded.to(dtype), b_indexed.to(dtype))

    if routed_weight:
        c *= topk_weights.unsqueeze(-1)

    if not routed_weight and gelu:
        c = 0.5 * c * (1.0 + torch.tanh(0.7978845608 * (c + 0.044715 * c * c * c)))

    if fp8_w8a8:
        c = c * b_scale[topk_ids].unsqueeze(-1)
        c = c * a_scale
        c = c.to(dtype)

    if int8_w8a16:
        c = c * b_scale[topk_ids].unsqueeze(-1)
        c = c.to(dtype)

    return c

def torch_moe_ref_block_w8a8(
    a,
    b,
    c,
    a_scale,
    b_scale,
    b_zp,
    block_shape,
    topk_ids,
    topk_weights,
    routed_weight,
    sorted_token_ids,
    expert_ids,
    num_tokens_post_padded,
    dtype,
    fp8_w8a8,
    int8_w8a8,
):
    assert (int8_w8a8 == True) and block_shape is not None
    M, top_k, N = c.shape
    _, K = a.shape
    E, _, _ = b.shape
    topk_ids = topk_ids.view(-1)
    a_rp = a.view(M, -1, K).repeat(1, top_k, 1).reshape(-1, K)
    as_rp = a_scale.view(M, -1, K // block_shape[1]).repeat(1, top_k, 1).reshape(-1, K // block_shape[1])
    c = c.view(M * top_k, N)

    for i in range(E):
        mask = topk_ids == i
        if mask.sum():
            c[mask] = native_w8a8_block_matmul(a_rp[mask],
                                                 b[i],
                                                 as_rp[mask],
                                                 b_scale[i],
                                                 block_shape,
                                                 output_dtype=dtype)
    c = c.view(M, top_k, N)
    if routed_weight:
       c = c * topk_weights.unsqueeze(-1)
    return c

def torch_moe_ref_block_w8a8_fp8(
    a,
    b,
    c,
    a_scale,
    b_scale,
    b_zp,
    block_shape,
    topk_ids,
    topk_weights,
    routed_weight,
    sorted_token_ids,
    expert_ids,
    num_tokens_post_padded,
    dtype,
    fp8_w8a8,
    int8_w8a8,
):
    assert (int8_w8a8 == True or fp8_w8a8 == True) and block_shape is not None
    M, top_k, N = c.shape
    _, K = a.shape
    E, _, _ = b.shape
    topk_ids = topk_ids.view(-1)
    a_rp = a.view(M, -1, K).repeat(1, top_k, 1).reshape(-1, K)
    as_rp = a_scale.view(M, -1, K // block_shape[1]).repeat(1, top_k, 1).reshape(-1, K // block_shape[1])
    c = c.view(M * top_k, N)

    for i in range(E):
        mask = topk_ids == i
        if mask.sum():
            # Convert FP8 to half before indexing to avoid CPU indexing issue
            c[mask] = native_w8a8_block_matmul(a_rp.half()[mask],
                                                 b[i],
                                                 as_rp[mask],
                                                 b_scale[i],
                                                 block_shape,
                                                 output_dtype=dtype)
    c = c.view(M, top_k, N)
    if routed_weight:
       c = c * topk_weights.unsqueeze(-1)
    return c

def _moe_align_block_size(
    topk_ids: torch.Tensor,
    num_experts: int,
    top_k: int,
    block_size: int,
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
) -> None:
    M, top_k = topk_ids.shape

    expert_to_tokens = [[] for _ in range(num_experts)]
    # For each token, for each selected expert, we append (token_id, expert)
    for token_id in range(M):
        for j in range(top_k):
            e_id = topk_ids[token_id, j].item()
            expert_to_tokens[e_id].append(token_id * top_k + j)

    # Reorder tokens block by block, padding if needed
    reordered_token_ids = []
    reordered_expert_ids = []

    for e_id in range(num_experts):
        tokens_for_expert = expert_to_tokens[e_id]
        num_tokens = len(tokens_for_expert)

        n_blocks = (num_tokens + block_size - 1) // block_size
        # If not a multiple of block_size, pad up to the next multiple
        padded_size = n_blocks * block_size

        # Reorder all actual tokens for expert e_id
        reordered_token_ids.extend(tokens_for_expert)
        # reordered_expert_ids.extend([e_id]*num_tokens)
        reordered_expert_ids.extend([e_id] * n_blocks)

        # Pad with dummy token_id = topk_ids.numel()
        if padded_size > num_tokens:
            pad_count = padded_size - num_tokens
            reordered_token_ids.extend([topk_ids.numel()] * pad_count)

    token_length = len(reordered_token_ids)
    expert_length = len(reordered_expert_ids)

    sorted_token_ids[:token_length] = torch.tensor(
        reordered_token_ids,
        dtype=sorted_token_ids.dtype,
        device=sorted_token_ids.device,
    )
    expert_ids[:expert_length] = torch.tensor(
        reordered_expert_ids, dtype=expert_ids.dtype, device=expert_ids.device
    )

    # Fill remainder with topk_ids.numel() if these arrays are bigger than total_length
    if token_length < sorted_token_ids.numel():
        sorted_token_ids[token_length:] = topk_ids.numel()
    if expert_length < expert_ids.numel():
        expert_ids[expert_length:] = topk_ids.numel()

    num_tokens_post_pad.fill_(token_length)


def torch_moe_align_block_size_ref(
    topk_ids: torch.Tensor, block_size: int, num_experts: int
) -> 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.

    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.
    """
    top_k = topk_ids.shape[1]
    sorted_ids = torch.empty(
        (topk_ids.numel() + num_experts * (block_size - 1),),
        dtype=torch.int32,
        device=topk_ids.device,
    )
    expert_ids = torch.empty(
        (topk_ids.numel() + num_experts,), dtype=torch.int32, device=topk_ids.device
    )
    sorted_ids.fill_(topk_ids.numel())
    num_tokens_post_pad = torch.empty((1), dtype=torch.int32, device=topk_ids.device)
    _moe_align_block_size(
        topk_ids,
        num_experts,
        top_k,
        block_size,
        sorted_ids,
        expert_ids,
        num_tokens_post_pad,
    )

    return sorted_ids, expert_ids, num_tokens_post_pad


def torch_e2e_moe(
    a,
    w1,
    w2,
    c,
    a_scale,
    w1_scale,
    w2_scale,
    topk_ids,
    topk_weights,
    routed_weight,
    dtype,
    fp8_w8a8,
    int8_w8a16,
):
    if fp8_w8a8:
        a, _, a_scale = quantize_fp8(a)

    M, top_k, _ = c.shape
    E, N, _ = w1.shape

    # Repeat a -> (M, top_k, K)
    a_expanded = a.unsqueeze(1).repeat(1, top_k, 1)
    # (M, top_k, N, K)
    if fp8_w8a8:
        w1_indexed = w1.half()[topk_ids]
    else:
        w1_indexed = w1[topk_ids]

    intermidiate = torch.einsum(
        "mek,menk->men", a_expanded.to(dtype), w1_indexed.to(dtype)
    )

    if fp8_w8a8:
        intermidiate = intermidiate * w1_scale[topk_ids].unsqueeze(-1)
        intermidiate = intermidiate * a_scale
        intermidiate = intermidiate.to(dtype)

    if int8_w8a16:
        intermidiate = intermidiate * w1_scale[topk_ids].unsqueeze(-1)
        intermidiate = intermidiate.to(dtype)

    if fp8_w8a8:
        w2_indexed = w2.half()[topk_ids]
    else:
        w2_indexed = w2[topk_ids]

    print(intermidiate.shape)

    silu_out = torch.zeros([M * top_k, N // 2], dtype=a.dtype, device=a.device)
    silu_out = torch_silu_and_mul_ref(intermidiate.view(-1, N))

    silu_out = silu_out.view(M, top_k, N // 2)

    if fp8_w8a8:
        silu_out, _, silu_out_scale = quantize_fp8(silu_out)

    c = torch.einsum("mek,menk->men", silu_out.to(dtype), w2_indexed.to(dtype))

    if fp8_w8a8:
        c = c * w2_scale[topk_ids].unsqueeze(-1)
        c = c * silu_out_scale
        c = c.to(dtype)

    if int8_w8a16:
        c = c * w2_scale[topk_ids].unsqueeze(-1)
        c = c.to(dtype)

    if routed_weight:
        c *= topk_weights.unsqueeze(-1)
    return c


def get_default_config() -> Dict[str, int]:
    config = {
        "BLOCK_SIZE_M": 64,
        "BLOCK_SIZE_N": 64,
        "BLOCK_SIZE_K": 32,
        "GROUP_SIZE_M": 8,
    }
    return config


def get_default_config_moe_e2e(persistent: bool) -> Dict[str, int]:
    if persistent:
        return {
            "BLOCK_SIZE_M": 64,
            "BLOCK_SIZE_N1": 128,
            "BLOCK_SIZE_N2": 64,
            "BLOCK_SIZE_K1": 64,
            "BLOCK_SIZE_K2": 64,
        }
    return {
        "BLOCK_SIZE_M": 64,
        "BLOCK_SIZE_N": 128,
        "BLOCK_SIZE_K1": 64,
        "BLOCK_SIZE_K2": 64,
        "GROUP_SIZE_M": 2,
    }  # TODO setting GROUP_SIZE_M = 1 gives set fault, why?


def quantize_fp8(
    tensor: torch.Tensor, dim=()
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    fp8_type = torch.float8_e4m3fn

    quantize_dim = [i for i in range(tensor.dim()) if i not in dim]
    max_vals = tensor.abs().amax(dim=quantize_dim, keepdim=True)
    max_repr_val = torch.finfo(fp8_type).max
    max_vals[max_vals == 0] = 1e-8  # Avoid division by zero

    # Compute scale factors for each channel
    scale: torch.Tensor = max_repr_val / max_vals.to(torch.float32)

    # Quantize the tensor
    tensor = tensor * scale
    tensor.clamp_(-max_repr_val, max_repr_val)
    tensor_quantized = tensor.to(fp8_type)

    scale = scale.squeeze(dim=quantize_dim)

    return tensor_quantized, scale, 1 / scale


def quantize_int8(
    tensor: torch.Tensor, dim=()
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    quantize_dim = [i for i in range(tensor.dim()) if i not in dim]
    max_vals = tensor.abs().amax(dim=quantize_dim, keepdim=True)
    max_repr_val = torch.iinfo(torch.int8).max
    max_vals[max_vals == 0] = 1e-8  # Avoid division by zero

    # Compute scale factors for each channel
    scale: torch.Tensor = max_repr_val / max_vals.to(torch.float32)

    # Quantize the tensor
    tensor = tensor * scale
    tensor.clamp_(-max_repr_val, max_repr_val)
    tensor = tensor.round_()
    tensor_quantized = tensor.to(torch.int8)

    scale = scale.squeeze(dim=quantize_dim)

    return tensor_quantized, scale, 1 / scale


def quantize_int4(
    tensor: torch.Tensor, group_size: int, has_zp: bool
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

    # reshape tensor
    k, n = tensor.shape
    tensor = tensor.reshape(-1, group_size, n)
    tensor = tensor.permute(1, 0, 2)

    max_val = torch.max(tensor, 0, keepdim=True).values
    min_val = torch.min(tensor, 0, keepdim=True).values

    # Asymmetric quantization
    zp = None
    if has_zp:
        max_q_val = 15
        min_q_val = 0  # Min maps to 0
        scale = (max_val - min_val).clamp(min=1e-5) / (max_q_val)
        zp = torch.round(torch.abs(min_val / scale)).clamp(min_q_val, max_q_val).int()
    # Symmetric quantization
    else:
        max_q_val = 7
        min_q_val = -7
        abs_max_val = torch.maximum(torch.abs(max_val), torch.abs(min_val))
        scale = abs_max_val.clamp(min=1e-5) / max_q_val

    # quantize and clamp
    tensor_q = torch.round(tensor / scale).int() + (zp if has_zp else 0)
    tensor_q = torch.clamp(tensor_q, min_q_val, max_q_val)

    # restore shapes
    tensor_q = tensor_q.reshape((group_size, -1, n))
    tensor_q = tensor_q.permute(1, 0, 2)
    tensor_q = tensor_q.reshape((k, n)).contiguous()

    # scale
    scale = scale.reshape((-1, n)).contiguous()

    # zp
    if zp is not None:
        zp = zp.reshape((-1, n)).contiguous()
        zp = zp.to(device=tensor.device)

    return tensor_q, scale, zp

def input_helper(
    M: int,
    N: int,
    K: int,
    top_k: int,
    E: int,
    routed_weight: bool,
    dtype,
    int8_w8a16: bool,
    int8_w8a8: bool,
    fp8_w8a8: bool,
    block_shape: Optional[list[int]],
):
    assert not (fp8_w8a8 and int8_w8a16)

    oN = N
    iK = K
    if not routed_weight:
      # w1 weight
      oN = 2 * N
      a = torch.randn((M, K), dtype=dtype, device="cpu") / 10
      b = torch.rand((E, 2 * N, K), dtype=dtype, device="cpu") / 10
    else:
      oN = K
      iK = N
      a = torch.randn((M, N), dtype=dtype, device="cpu") / 10
      b = torch.rand((E, K, N), dtype=dtype, device="cpu") / 10

    a_scale = None
    b_scale = None
    b_zp = False

    if fp8_w8a8:
        b, _, b_scale = quantize_fp8(b, dim=(0,))
    if int8_w8a16:
        b, _, b_scale = quantize_int8(b, dim=(0,))

    if int8_w8a8:
        assert (block_shape is not None) and len(block_shape) == 2 and block_shape[0] > 0 and block_shape[1] > 0
        factor_for_scale = 1e-2
        int8_info = torch.iinfo(torch.int8)
        int8_max, int8_min = int8_info.max, int8_info.min
        bf = (torch.rand((E, oN, iK), dtype=dtype, device="cpu") - 0.5) * 2 * int8_max
        b = bf.clamp(min=int8_min, max=int8_max).to(torch.int8)

        block_n, block_k = block_shape[0], block_shape[1]
        n_tiles = (oN + block_n - 1) // block_n
        k_tiles = (iK + block_k - 1) // block_k
        b_scale = (torch.rand(
            (E, n_tiles, k_tiles), dtype=torch.float32, device="cpu") * factor_for_scale)
        a, a_scale = native_per_token_group_quant_int8(a, block_k)
    if fp8_w8a8:
        # reference test_block_fp8.py
        assert (block_shape is not None) and len(block_shape) == 2 and block_shape[0] > 0 and block_shape[1] > 0
        factor_for_scale = 1e-2
        fp8_info = torch.finfo(torch.float8_e4m3fn)
        fp8_max, fp8_min = fp8_info.max, fp8_info.min
        bf = (torch.rand((E, oN, iK), dtype=dtype, device="cpu") - 0.5) * 2 * fp8_max
        b = bf.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)

        block_n, block_k = block_shape[0], block_shape[1]
        n_tiles = (oN + block_n - 1) // block_n
        k_tiles = (iK + block_k - 1) // block_k
        b_scale = (torch.rand(
            (E, n_tiles, k_tiles), dtype=torch.float32, device="cpu") * factor_for_scale)
        a, a_scale = native_per_token_group_quant_fp8(a, block_k)


    c = torch.zeros((M, top_k, oN), dtype=dtype, device="cpu")
    c_silu = torch.zeros((M * top_k, oN // 2), dtype=dtype, device="cpu")

    values = torch.randn(M, E, dtype=dtype, device="cpu")

    softmax_vals = torch.softmax(values, dim=1)
    topk_weights, topk_ids = torch.topk(softmax_vals, k=top_k, dim=1)

    moe_config_func = get_optimal_moe_config_func(
        a, b, topk_ids, use_int8_w8a16=int8_w8a16, use_int8_w8a8=int8_w8a8,
        use_fp8_w8a8=fp8_w8a8, is_bottom=routed_weight, block_shape=block_shape)
    config = moe_config_func(M)

    sorted_ids, expert_ids, num_tokens_post_pad = (
        torch_moe_align_block_size_ref(topk_ids, config["BLOCK_SIZE_M"], E)
    )
    if DEBUG_MODE:
        print(f"M={M}, N={N}, K={K}, top_K={top_k}, E={E}")
        print(f"config={config}")
        print(f"topk_ids={topk_ids}")
        print(f"sorted_ids.shape={sorted_ids.shape}")
        print(f"sorted_ids={sorted_ids}")
        print(f"expert_ids.shape={expert_ids.shape}")
        print(f"expert_ids={expert_ids}")
        print(f"num_tokens_post_padded={num_tokens_post_pad}")

    return (
        a,
        b,
        c,
        c_silu,
        b_zp,
        a_scale,
        b_scale,
        topk_weights,
        topk_ids,
        sorted_ids,
        expert_ids,
        num_tokens_post_pad,
        config,
    )


def input_helper_int4_w4a16(
    M: int,
    N: int,
    K: int,
    top_k: int,
    E: int,
    routed_weight: bool,
    dtype: torch.dtype,
    group_size: int,
    has_zp: bool,
    use_int4_w4a8: bool = False,
):

    oN = N
    iK = K
    pack_factor = 2
    if not routed_weight:
      # w1 weight
      oN = 2 * N
      a = torch.randn((M, K), dtype=dtype, device="cuda") / 10
      b = torch.rand((E, 2 * N, K), dtype=dtype, device="cuda") / 10
    else:
      oN = K
      iK = N
      a = torch.randn((M, N), dtype=dtype, device="cuda") / 10
      b = torch.rand((E, K, N), dtype=dtype, device="cuda") / 10

    b_q = torch.empty((E, oN, iK // pack_factor), dtype=torch.uint8, device="cuda")
    b_scale = torch.empty((E, oN, iK // group_size), dtype=dtype, device="cuda")
    if has_zp:
        b_zp = torch.empty(
            (E, oN // pack_factor, iK // group_size), dtype=torch.uint8, device="cuda"
        )
    else:
        b_zp = None

    for e in range(E):
        q, scale, zp = quantize_int4(b[e].T, group_size=group_size, has_zp=has_zp)
        q = q.T
        q = (
            q[:, 1::2] * 16 + q[:, ::2]
        )  # Note, 2<<4=16. For bf16, etc, torch doesn't have shift.
        b_q[e] = q
        b_scale[e] = scale.T
        if has_zp:
            zp = zp.T.contiguous().to(torch.uint8)
            zp = (
                zp[1::2, :] << 4 | zp[::2, :]
            )  # Note, 2<<4=16. For bf16, etc, torch doesn't have shift.
            b_zp[e] = zp

    b = b_q

    c = torch.zeros((M, top_k, oN), dtype=dtype, device="cuda")
    c_silu = torch.zeros((M * top_k, oN // 2), dtype=dtype, device="cuda")

    values = torch.randn(M, E, dtype=dtype, device="cuda")

    softmax_vals = torch.softmax(values, dim=1)
    topk_weights, topk_ids = torch.topk(softmax_vals, k=top_k, dim=1)
    topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)

    use_int4_w4a16 = True
    if use_int4_w4a8:
        use_int4_w4a16 = False
    moe_config_func = get_optimal_moe_config_func(a, b, topk_ids, use_int4_w4a16=use_int4_w4a16,
                                                  use_int4_w4a8=use_int4_w4a8,
                                                  is_bottom=routed_weight,
                                                  block_shape=(0, group_size))
    config = moe_config_func(M)

    sorted_ids, expert_ids, num_tokens_post_pad = (
        torch_moe_align_block_size_ref(topk_ids, config["BLOCK_SIZE_M"], E)
    )
    if DEBUG_MODE:
        print(f"M={M}, N={N}, K={K}, top_K={top_k}, E={E}")
        print(f"config={config}")
        print(f"topk_ids={topk_ids}")
        print(f"sorted_ids.shape={sorted_ids.shape}")
        print(f"sorted_ids={sorted_ids}")
        print(f"expert_ids.shape={expert_ids.shape}")
        print(f"expert_ids={expert_ids}")
        print(f"num_tokens_post_padded={num_tokens_post_pad}")

    return (
        a,
        b,
        c,
        c_silu,
        b_zp,
        b_scale,
        topk_weights,
        topk_ids,
        sorted_ids,
        expert_ids,
        num_tokens_post_pad,
        config,
    )


def input_helper_e2e(
    M: int,
    N: int,
    K: int,
    top_k: int,
    E: int,
    routed_weight: bool,
    dtype,
    fp8_w8a8: bool,
    int8_w8a16: bool,
    persistent: bool,
):
    assert not (fp8_w8a8 and int8_w8a16)

    a = torch.randn((M, K), dtype=dtype, device="cuda")
    w1 = torch.rand((E, N, K), dtype=dtype, device="cuda")
    w2 = torch.rand((E, K, N // 2), dtype=dtype, device="cuda")
    a_scale = None
    w1_scale = None
    w2_scale = None

    if fp8_w8a8:
        w1, _, w1_scale = quantize_fp8(w1, dim=(0,))
        w2, _, w2_scale = quantize_fp8(w2, dim=(0,))

    if int8_w8a16:
        w1, _, w1_scale = quantize_int8(w1, dim=(0,))
        w2, _, w2_scale = quantize_int8(w2, dim=(0,))

    c = torch.zeros((M, top_k, K), dtype=dtype, device="cuda")

    values = torch.randn(M, E, dtype=dtype, device="cuda")

    softmax_vals = torch.softmax(values, dim=1)
    topk_weights, topk_ids = torch.topk(softmax_vals, k=top_k, dim=1)

    config = get_default_config_moe_e2e(persistent)
    sorted_token_ids, expert_ids, num_tokens_post_padded = (
        torch_moe_align_block_size_ref(topk_ids, config["BLOCK_SIZE_M"], E)
    )

    return (
        a,
        w1,
        w2,
        c,
        a_scale,
        w1_scale,
        w2_scale,
        topk_weights,
        topk_ids,
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        config,
    )


# Note: TODO These 2 result in accuracy issues (64, 14336, 4096, 2, 8), (1, 1024, 16384, 1, 2)
'''
@pytest.mark.parametrize(
    "M, N, K, top_k, E",
    [
        (1, 256, 7168, 8, 32),
        #(1, 256, 7168, 8, 256),
        #(2, 256, 7168, 8, 256),
        #(4, 256, 7168, 8, 256),
        #(8, 256, 7168, 8, 256),
        #(16, 256, 7168, 8, 256),
        #(64, 256, 7168, 8, 256),
        #(96, 256, 7168, 8, 256),
        #(128, 256, 7168, 8, 256),
        #(256, 256, 7168, 8, 256),
    ],
)
@pytest.mark.parametrize("routed_weight", [False, True])
@pytest.mark.parametrize("fp8_w8a8, int8_w8a16, int8_w8a8, block_shape",
                         [
                          #(False, False, False, None),
                          (False, False, True, [128, 128])
                         ])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
# @pytest.mark.parametrize("dtype", [torch.bfloat16])
@pytest.mark.parametrize("persistent", [False])
@pytest.mark.parametrize("silu_fused", [False])
'''

def test_fused_moe(
    M: int,
    N: int,
    K: int,
    top_k: int,
    E: int,
    routed_weight: bool,
    fp8_w8a8: bool,
    int8_w8a16: bool,
    int8_w8a8: bool,
    persistent: bool,
    silu_fused: bool,
    dtype,
    block_shape,
):
    #torch.manual_seed(20)
    #torch.set_printoptions(threshold=10000)
    # if persistent:
    #     (
    #         triton_moe_silu_set_use_persistent_kernel(True)
    #         if silu_fused
    #         else triton_moe_set_use_persistent_kernel(True)
    #     )
    # else:
    #     (
    #         triton_moe_silu_set_use_persistent_kernel(False)
    #         if silu_fused
    #         else triton_moe_set_use_persistent_kernel(False)
    #     )

    (
        a,
        b,
        triton_out,
        triton_out_silu,
        b_zp,
        a_scale,
        b_scale,
        topk_weights,
        topk_ids,
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        config,
    ) = input_helper(
        M,
        N,
        K,
        top_k,
        E,
        routed_weight=routed_weight,
        dtype=dtype,
        fp8_w8a8=fp8_w8a8,
        int8_w8a16=int8_w8a16,
        int8_w8a8=int8_w8a8,
        block_shape=block_shape,
    )
    # config will be auto selected in triton_moe
    # config = None
    _triton_moe = triton_moe_silu if silu_fused else triton_moe
    triton_out_device = triton_out.to("cuda")
    triton_out_silu_device = triton_out_silu.to("cuda")

    _triton_moe(
        a.to("cuda"),
        b.to("cuda"),
        triton_out_silu_device if silu_fused else triton_out_device,
        a_scale.to("cuda"),
        b_scale.to("cuda"),
        b_zp,
        topk_weights.to("cuda"),
        topk_ids,
        sorted_token_ids.to("cuda"),
        expert_ids.to("cuda"),
        num_tokens_post_padded.to("cuda"),
        routed_weight,
        top_k,
        torch_to_triton_dtype[dtype],
        use_fp8_w8a8=fp8_w8a8,
        use_int8_w8a8=int8_w8a8,
        use_int8_w8a16=int8_w8a16,
        block_shape=block_shape,
        config=config,
    )
    # print(f"a:{a}, b:{b}")
    triton_out_cpu = triton_out_device.to("cpu")
    torch_out = torch.empty_like(triton_out_cpu)
    if (int8_w8a8 or fp8_w8a8) and block_shape is not None:
        if int8_w8a8:
            torch_out = torch_moe_ref_block_w8a8(
                a,
                b,
                torch_out,
                a_scale,
                b_scale,
                None,
                block_shape,
                topk_ids,
                topk_weights,
                routed_weight,
                sorted_token_ids,
                expert_ids,
                num_tokens_post_padded,
                dtype,
                fp8_w8a8,
                int8_w8a8,
            )
        elif fp8_w8a8:
            torch_out = torch_moe_ref_block_w8a8_fp8(
                a,
                b,
                torch_out,
                a_scale,
                b_scale,
                None,
                block_shape,
                topk_ids,
                topk_weights,
                routed_weight,
                sorted_token_ids,
                expert_ids,
                num_tokens_post_padded,
                dtype,
                fp8_w8a8,
                int8_w8a8,
            )
    else:
        torch_out = torch_moe_ref(
            a,
            b,
            torch_out,
            a_scale,
            b_scale,
            None,
            0,
            topk_ids,
            topk_weights,
            routed_weight,
            sorted_token_ids,
            expert_ids,
            num_tokens_post_padded,
            dtype,
            fp8_w8a8,
            int8_w8a16,
            False,
        )
    if silu_fused:
        torch_out_silu = torch_silu_and_mul_ref(torch_out.view(-1, N))

    if DEBUG_MODE:
        print(f"triton_out={triton_out}")
        print(f"torch_out={torch_out}")
    # Validate correctness
    if silu_fused:
        torch.testing.assert_close(
            triton_out_silu, torch_out_silu, atol=1e-1, rtol=1e-1
        )
    else:
        torch.testing.assert_close(triton_out_cpu, torch_out, atol=2e-2, rtol=2e-2)


@pytest.mark.parametrize(
    "M, N, K, top_k, E",
    [
        (1, 256, 7168, 8, 256),
        (2, 256, 7168, 8, 256),
        (4, 256, 7168, 8, 256),
        (8, 256, 7168, 8, 256),
        (16, 256, 7168, 8, 256),
        (64, 256, 7168, 8, 256),
        (96, 256, 7168, 8, 256),
        (128, 256, 7168, 8, 256),
        (256, 256, 7168, 8, 256),
    ],
)
@pytest.mark.parametrize("routed_weight", [False, True])
@pytest.mark.parametrize("group_size", [64])
# @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("dtype", [torch.float16])
@pytest.mark.parametrize("has_zp", [True])
@pytest.mark.parametrize("persistent", [False])
@pytest.mark.parametrize("silu_fused", [False])
def test_fused_moe_int4_w4a16(
    M: int,
    N: int,
    K: int,
    top_k: int,
    E: int,
    routed_weight: bool,
    dtype: torch.dtype,
    group_size: int,
    has_zp: bool,
    persistent: bool,
    silu_fused: bool,
):

    if (
        M == 1
        and N == 64
        and K == 128
        and top_k == 1
        and E == 2
        and group_size == 8
        and routed_weight
        and not persistent
        and has_zp
        and not silu_fused
    ):
        pytest.skip("Results in accuracy failure because of Triton compiler change")

    torch.manual_seed(20)
    (
        a,
        b,
        triton_out,
        triton_out_silu,
        b_zp,
        b_scale,
        topk_weights,
        topk_ids,
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        config,
    ) = input_helper_int4_w4a16(
        M,
        N,
        K,
        top_k,
        E,
        routed_weight=routed_weight,
        dtype=dtype,
        group_size=group_size,
        has_zp=has_zp,
    )

    # if persistent:
    #     (
    #         triton_moe_silu_set_use_persistent_kernel(True)
    #         if silu_fused
    #         else triton_moe_set_use_persistent_kernel(True)
    #     )
    # else:
    #     (
    #         triton_moe_silu_set_use_persistent_kernel(False)
    #         if silu_fused
    #         else triton_moe_set_use_persistent_kernel(False)
    #     )

    _triton_moe = triton_moe_silu if silu_fused else triton_moe
    _triton_moe(
        a,
        b,
        triton_out_silu if silu_fused else triton_out,
        None,
        b_scale,
        b_zp,
        topk_weights,
        topk_ids,
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        routed_weight,
        top_k,
        torch_to_triton_dtype[dtype],
        use_fp8_w8a8=False,
        use_int8_w8a16=False,
        use_int4_w4a16=True,
        block_shape=[0, group_size],
        config=config,
    )

    torch_out = torch.empty_like(triton_out)
    torch_out = torch_moe_ref(
        a,
        b,
        torch_out,
        None,
        b_scale,
        b_zp,
        group_size,
        topk_ids,
        topk_weights,
        routed_weight,
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        dtype,
        False,
        False,
        True,
    )
    # torch_out.view(torch.uint16).cpu().numpy().tofile(f"{dtype}_torch_out_failed.bin")
    # triton_out.view(torch.uint16).cpu().numpy().tofile(f"{dtype}_triton_out_failed.bin")
    # import numpy as np
    # np.savetxt(f"{dtype}_torch_out_failed.txt", torch_out.view(-1, 32).to(torch.float32).cpu().numpy())
    # np.savetxt(f"{dtype}_triton_out_failed.txt", triton_out.view(-1, 32).to(torch.float32).cpu().numpy())
    if silu_fused:
        torch_out_silu = torch_silu_and_mul_ref(torch_out.view(-1, N))

    if silu_fused:
        torch.testing.assert_close(
            triton_out_silu, torch_out_silu, atol=2e-1, rtol=2e-1
        )
    else:
        torch.testing.assert_close(triton_out, torch_out, atol=2e-2, rtol=2e-2)

@pytest.mark.parametrize(
    "M, N, K, top_k, E",
    [
        (1, 256, 7168, 8, 256),
        (2, 256, 7168, 8, 256),
        (4, 256, 7168, 8, 256),
        (8, 256, 7168, 8, 256),
        (16, 256, 7168, 8, 256),
        (64, 256, 7168, 8, 256),
        (96, 256, 7168, 8, 256),
        (128, 256, 7168, 8, 256),
        (256, 256, 7168, 8, 256),
    ],
)
@pytest.mark.parametrize("routed_weight", [False, True])
@pytest.mark.parametrize("group_size", [64])
# @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("dtype", [torch.float16])
@pytest.mark.parametrize("has_zp", [True])
def test_fused_moe_int4_w4a8(
    M: int,
    N: int,
    K: int,
    top_k: int,
    E: int,
    routed_weight: bool,
    dtype: torch.dtype,
    group_size: int,
    has_zp: bool,
):
    torch.manual_seed(20)
    (
        a,
        b,
        triton_out,
        triton_out_silu,
        b_zp,
        b_scale,
        topk_weights,
        topk_ids,
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        config,
    ) = input_helper_int4_w4a16(
        M,
        N,
        K,
        top_k,
        E,
        routed_weight=routed_weight,
        dtype=dtype,
        group_size=group_size,
        has_zp=has_zp,
        use_int4_w4a8=True,
    )

    a_quant, a_scale = per_block_quant_wrapper((1, group_size))(per_token_quant_hip)(a)
    _triton_moe = triton_moe
    _triton_moe(
        a_quant,
        b,
        triton_out,
        a_scale,
        b_scale,
        b_zp,
        topk_weights,
        topk_ids,
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        routed_weight,
        top_k,
        torch_to_triton_dtype[dtype],
        use_int4_w4a8=True,
        block_shape=[0, group_size],
        config=config,
    )

    torch_out = torch.empty_like(triton_out)
    torch_out = torch_moe_ref(
        a,
        b,
        torch_out,
        None,
        b_scale,
        b_zp,
        group_size,
        topk_ids,
        topk_weights,
        routed_weight,
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        dtype,
        fp8_w8a8=False,
        int8_w8a16=False,
        int4_w4a16=True,
    )

    torch.testing.assert_close(triton_out, torch_out, atol=2e-2, rtol=2e-2)


# Note: TODO These 2 result in accuracy issues (64, 14336, 4096, 2, 8), (1, 1024, 16384, 1, 2)
# @pytest.mark.parametrize(
#     "M, N, K, top_k, E",
#     [
#         (64, 14336, 4096, 2, 8),
#         (16, 14336, 1, 2, 4),
#         (4, 4, 8, 1, 2),
#         (1, 14336, 128, 2, 4),
#         (3, 14336, 128, 2, 4),
#         (16, 14336, 128, 1, 4),
#         (16, 14336, 128, 1, 1),
#         (64, 7186, 128, 2, 8),
#         (64, 3584, 128, 2, 8),
#         (64, 1792, 128, 2, 8),
#         (64, 64, 128, 2, 8),
#         (1, 1024, 16384, 1, 2),
#     ],
# )
# @pytest.mark.parametrize("routed_weight", [False, True])
# # @pytest.mark.parametrize('fp8_w8a8, int8_w8a16', [(False, False), (True, False), (False, True)]) #TODO: Accuracy issues with fp8
# @pytest.mark.parametrize("fp8_w8a8, int8_w8a16", [(False, False)])
# @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
# @pytest.mark.parametrize("persistent", [False, True])
# def test_fused_moe_gelu(
#     M: int,
#     N: int,
#     K: int,
#     top_k: int,
#     E: int,
#     routed_weight: bool,
#     fp8_w8a8: bool,
#     int8_w8a16: bool,
#     persistent: bool,
#     dtype,
# ):
#     torch.manual_seed(20)
#     torch.set_printoptions(threshold=100000)
#     if persistent:
#         triton_moe_gelu_set_use_persistent_kernel(True)
#     else:
#         triton_moe_gelu_set_use_persistent_kernel(False)

#     (
#         a,
#         b,
#         triton_out,
#         triton_out_silu,
#         b_zp,
#         a_scale,
#         b_scale,
#         topk_weights,
#         topk_ids,
#         sorted_token_ids,
#         expert_ids,
#         num_tokens_post_padded,
#         config,
#     ) = input_helper(
#         M,
#         N,
#         K,
#         top_k,
#         E,
#         routed_weight=routed_weight,
#         dtype=dtype,
#         fp8_w8a8=fp8_w8a8,
#         int8_w8a16=int8_w8a16,
#     )

#     if DEBUG_MODE:
#         print(f"M={M}, N={N}, K={K}, top_K={top_k}, E={E}")
#         print(f"config={config}")
#         print(f"a.shape={a.shape} a={a}")
#         print(f"b.shape={b.shape} b={b}")
#         print(f"sorted_token_ids.shape={sorted_token_ids.shape}")
#         print(f"sorted_token_ids={sorted_token_ids}")
#         print(f"expert_ids.shape={expert_ids.shape}")
#         print(f"expert_ids={expert_ids}")
#         print(f"num_tokens_post_padded={num_tokens_post_padded}")
#     triton_moe_gelu(
#         a,
#         b,
#         triton_out,
#         a_scale,
#         b_scale,
#         topk_weights,
#         topk_ids,
#         sorted_token_ids,
#         expert_ids,
#         num_tokens_post_padded,
#         routed_weight,
#         top_k,
#         torch_to_triton_dtype[dtype],
#         fp8_w8a8,
#         int8_w8a16,
#         config=config,
#     )

#     torch_out = torch.empty_like(triton_out)
#     torch_out = torch_moe_ref(
#         a,
#         b,
#         torch_out,
#         a_scale,
#         b_scale,
#         None,
#         0,
#         topk_ids,
#         topk_weights,
#         routed_weight,
#         sorted_token_ids,
#         expert_ids,
#         num_tokens_post_padded,
#         dtype,
#         fp8_w8a8,
#         int8_w8a16,
#         False,
#         gelu=True,
#     )

#     if DEBUG_MODE:
#         print(f"triton_out={triton_out}")
#         print(f"torch_out={torch_out}")
#     # Validate correctness
#     torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-1)


# TODO (64, 7186, 128, 2, 8), (64, 3584, 128, 2, 8), (4, 4, 8, 1, 2), (64, 1792, 128, 2, 8), (64, 64, 128, 2, 8) don't work because of the percision issue with atomics
# @pytest.mark.parametrize(
#     "M, N, K, top_k, E",
#     [
#         (16, 14336, 4096, 2, 8),
#         (16, 14336, 1, 2, 4),
#         (1, 14336, 128, 2, 4),
#         (3, 14336, 128, 2, 4),
#         (16, 14336, 128, 1, 4),
#         (16, 14336, 128, 1, 1),
#         (1, 1024, 16384, 1, 2),
#     ],
# )
# @pytest.mark.parametrize("routed_weight", [False, True])
# # @pytest.mark.parametrize('fp8_w8a8, int8_w8a16', [(False, False), (True, False), (False, True)]) #TODO: Accuracy issues with fp8
# @pytest.mark.parametrize("fp8_w8a8, int8_w8a16", [(False, False)])
# @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
# # @pytest.mark.parametrize('dtype', [torch.float16, torch.bfloat16])
# @pytest.mark.parametrize("persistent", [True, False])
# def test_moe_e2e(
#     M: int,
#     N: int,
#     K: int,
#     top_k: int,
#     E: int,
#     routed_weight: bool,
#     fp8_w8a8: bool,
#     int8_w8a16: bool,
#     persistent: bool,
#     dtype,
# ):
#     torch.manual_seed(20)
#     torch.set_printoptions(threshold=100000)
#     if persistent:
#         triton_e2e_moe_set_use_persistent_kernel(True)
#     else:
#         triton_e2e_moe_set_use_persistent_kernel(False)

#     intermediate = None
#     if persistent:
#         intermediate = torch.zeros(
#             (M * top_k, N // 2), dtype=torch.float32, device="cuda"
#         )

#     (
#         a,
#         w1,
#         w2,
#         triton_out,
#         a_scale,
#         w1_scale,
#         w2_scale,
#         topk_weights,
#         topk_ids,
#         sorted_token_ids,
#         expert_ids,
#         num_tokens_post_padded,
#         config,
#     ) = input_helper_e2e(
#         M,
#         N,
#         K,
#         top_k,
#         E,
#         routed_weight=routed_weight,
#         dtype=dtype,
#         fp8_w8a8=fp8_w8a8,
#         int8_w8a16=int8_w8a16,
#         persistent=persistent,
#     )

#     if DEBUG_MODE:
#         print(f"M={M}, N={N}, K={K}, top_K={top_k}, E={E}")
#         print(f"config={config}")
#         print(f"a.shape={a.shape} a={a}")
#         print(f"w1.shape={w1.shape} w1={w1}")
#         print(f"w2.shape={w2.shape} w2={w2}")
#         print(f"sorted_token_ids.shape={sorted_token_ids.shape}")
#         print(f"sorted_token_ids={sorted_token_ids}")
#         print(f"expert_ids.shape={expert_ids.shape}")
#         print(f"expert_ids={expert_ids}")
#         print(f"num_tokens_post_padded={num_tokens_post_padded}")
#     triton_out = triton_e2e_moe(
#         a,
#         w1,
#         w2,
#         intermediate,
#         triton_out,
#         a_scale,
#         w1_scale,
#         w2_scale,
#         topk_weights,
#         sorted_token_ids,
#         topk_ids,
#         expert_ids,
#         num_tokens_post_padded,
#         routed_weight,
#         top_k,
#         fp8_w8a8,
#         int8_w8a16,
#         config,
#     )

#     torch_out = torch.empty_like(triton_out)
#     torch_out = torch_e2e_moe(
#         a,
#         w1,
#         w2,
#         torch_out,
#         a_scale,
#         w1_scale,
#         w2_scale,
#         topk_ids,
#         topk_weights,
#         routed_weight,
#         dtype,
#         fp8_w8a8,
#         int8_w8a16,
#     )

#     if DEBUG_MODE:
#         print(f"triton_out={triton_out}")
#         print(f"torch_out={torch_out}")
#     # Validate correctness
#     torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-1)
if __name__ == "__main__":
    sizes = [
        (1, 256, 7168, 8, 256),
        (2, 256, 7168, 8, 256),
        (4, 256, 7168, 8, 256),
        (8, 256, 7168, 8, 256),
        (16, 256, 7168, 8, 256),
        # (64, 256, 7168, 8, 256),
        # (96, 256, 7168, 8, 256),
        # (128, 256, 7168, 8, 256),
        # (256, 256, 7168, 8, 256),
    ]
    for i in range(len(sizes)):
        print(i)
        (M, N, K, top_k, E) = sizes[i]
        ## int8 w8a8
        # test_fused_moe(
        #     M=M,
        #     N=N,
        #     K=K,
        #     top_k=top_k,
        #     E=E,
        #     routed_weight=False,
        #     fp8_w8a8=False,
        #     int8_w8a16=False,
        #     int8_w8a8=True,
        #     block_shape=[128, 128],
        #     dtype=torch.float16,
        #     persistent=False,
        #     silu_fused=False
        # )
        ## fp8 w8a8
        test_fused_moe(
            M=M,
            N=N,
            K=K,
            top_k=top_k,
            E=E,
            routed_weight=False,
            fp8_w8a8=True,
            int8_w8a16=False,
            int8_w8a8=False,
            block_shape=[128, 128],
            dtype=torch.bfloat16,
            persistent=False,
            silu_fused=False
        )

        print("--------------------------------run complete--------------------------------")
