import torch
import torch.nn.functional as F
import ctypes
from typing import Optional, Dict
import os
import threading
import pandas as pd
import functools
import aiter
from bisect import bisect_left
from aiter import logger
from aiter import per_token_quant_hip, per_block_quant_wrapper, get_hip_quant
from aiter import ActivationType, QuantType, dtypes
from aiter import silu_and_mul,gelu_and_mul
from aiter.ops.triton.fused_moe import (
    triton_moe_sum,
    triton_silu_and_mul,
    triton_gelu_and_mul
)

from aiter.jit.core import AITER_ROOT_DIR
# from vllm.model_executor.layers.fused_moe.fused_moe import moe_align_block_size
# from vllm.model_executor.layers.quantization.utils.int8_utils import (
#     per_token_group_quant_int8, per_token_quant_int8)
from aiter.ops.triton.group_quant_int8 import per_token_group_quant_int8
from aiter.jit.utils.chip_info import get_gfx, get_cu_num
from functools import lru_cache
from aiter.jit.utils.torch_guard import torch_compile_guard


def moe_sorting_ck(
    topk_ids,
    topk_weights,
    num_experts,
    model_dim,
    moe_buf,
    block_size=32,
    expert_mask=None,
):
    device = topk_ids.device
    M, topk = topk_ids.shape
    topk = topk_ids.shape[1]
    max_num_tokens_padded = topk_ids.numel() + num_experts * block_size - topk
    max_num_m_blocks = int((max_num_tokens_padded + block_size - 1) // block_size)
    sorted_ids = torch.empty((max_num_tokens_padded,), dtype=dtypes.i32, device=device)
    sorted_weights = torch.empty(
        (max_num_tokens_padded,), dtype=dtypes.fp32, device=device
    )
    sorted_expert_ids = torch.empty(
        (max_num_m_blocks,), dtype=dtypes.i32, device=device
    )
    tokens_positions_per_expert = torch.empty(
        (num_experts*2,), dtype=dtypes.i32, device=device
    )
    num_valid_ids = torch.empty((1), dtype=dtypes.i32, device=device)

# for now, moe_sorting_fwd only support int32 topk_ids
    if topk_ids.dtype != dtypes.i32:
        topk_ids = topk_ids.to(dtypes.i32)

    aiter.moe_sorting_fwd(
        topk_ids,
        topk_weights,
        sorted_ids,
        sorted_weights,
        sorted_expert_ids,
        tokens_positions_per_expert,
        num_valid_ids,
        moe_buf,
        num_experts,
        block_size,
        expert_mask,
    )
    return sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, tokens_positions_per_expert, moe_buf

#@staticmethod
def run_fused_experts_asm_impl(hidden_states: torch.Tensor,
                   w1: torch.Tensor,
                   w2: torch.Tensor,
                   topk_weights: torch.Tensor,
                   topk_ids: torch.Tensor,
                   dtype: torch.dtype,
                   inplace: bool = False,
                   activation: str = "silu",
                   use_fp8_w8a8: bool = False,
                   use_int8_w8a8: bool = False,
                   use_int8_w4a8: bool = False,
                   use_int8_w8a16: bool = False,
                   use_int4_w4a16: bool = False,
                   per_channel_quant: bool = False,
                   global_num_experts: int = -1,
                   expert_map: Optional[torch.Tensor] = None,
                   w1_scale: Optional[torch.Tensor] = None,
                   w2_scale: Optional[torch.Tensor] = None,
                   w1_zp: Optional[torch.Tensor] = None,
                   w2_zp: Optional[torch.Tensor] = None,
                   a1_scale: Optional[torch.Tensor] = None,
                   a2_scale: Optional[torch.Tensor] = None,
                   block_shape: Optional[list[int]] = None,
                   use_persist: bool = False,
                   persist_cu: Optional[int] = 0,
                   use_shuffle: Optional[int] = 0,
                   solution_id: Optional[str] = None)-> torch.Tensor:
    return fused_experts_asm_impl(
                    hidden_states,
                    w1,
                    w2,
                    topk_weights,
                    topk_ids,
                    dtype,
                    inplace,
                    activation,
                    use_fp8_w8a8,
                    use_int8_w8a8,
                    use_int8_w4a8,
                    use_int8_w8a16,
                    use_int4_w4a16,
                    per_channel_quant,
                    global_num_experts,
                    expert_map,
                    w1_scale,
                    w2_scale,
                    w1_zp,
                    w2_zp,
                    a1_scale,
                    a2_scale,
                    block_shape,
                    use_persist,
                    persist_cu,
                    use_shuffle,
                    solution_id
                )

def fused_moe_fake(
    hidden_states: torch.Tensor,
    w1: torch.Tensor,
    w2: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    dtype,
    inplace: bool = False,
    activation: str = "silu",
    use_fp8_w8a8: bool = False,
    use_int8_w8a8: bool = False,
    use_int8_w4a8: bool = False,
    use_int8_w8a16: bool = False,
    use_int4_w4a16: bool = False,
    per_channel_quant: bool = False,
    global_num_experts: int = -1,
    expert_map: Optional[torch.Tensor] = None,
    w1_scale: Optional[torch.Tensor] = None,
    w2_scale: Optional[torch.Tensor] = None,
    w1_zp: Optional[torch.Tensor] = None,
    w2_zp: Optional[torch.Tensor] = None,
    a1_scale: Optional[torch.Tensor] = None,
    a2_scale: Optional[torch.Tensor] = None,
    block_shape: Optional[list[int]] = None,
    use_persist: bool = False,
    persist_cu: Optional[int] = 0,
    use_shuffle: Optional[int] = 0,
    solution_id: Optional[str] = None
) -> torch.Tensor:
    device = topk_ids.device
    M, topk = topk_ids.shape
    dtype = dtype
    # E, model_dim, inter_dim = get_inter_dim(w1.shape, w2.shape)
    # FIXME: W2.size must be same as hidden_dim
    moe_buf = torch.empty((M, w2.size(1)), dtype=dtype, device=device)
    return moe_buf
    



@torch_compile_guard(gen_fake=fused_moe_fake)
def fused_experts_asm_impl(hidden_states: torch.Tensor,
                       w1: torch.Tensor,
                       w2: torch.Tensor,
                       topk_weights: torch.Tensor,
                       topk_ids: torch.Tensor,
                       dtype: torch.dtype,
                       inplace: bool = False,
                       activation: str = "silu",
                       use_fp8_w8a8: bool = False,
                       use_int8_w8a8: bool = False,
                       use_int8_w4a8: bool = False,
                       use_int8_w8a16: bool = False,
                       use_int4_w4a16: bool = False,
                       per_channel_quant: bool = False,
                       global_num_experts: int = -1,
                       expert_map: Optional[torch.Tensor] = None,
                       w1_scale: Optional[torch.Tensor] = None,
                       w2_scale: Optional[torch.Tensor] = None,
                       w1_zp: Optional[torch.Tensor] = None,
                       w2_zp: Optional[torch.Tensor] = None,
                       a1_scale: Optional[torch.Tensor] = None,
                       a2_scale: Optional[torch.Tensor] = None,
                       block_shape: Optional[list[int]] = None,
                       use_persist: bool = False,
                       persist_cu: Optional[int] = 0,
                       use_shuffle: Optional[int] = 0,
                       solution_id: Optional[str] = None,
                       routed_scaling_factor: Optional[float] = 1.0)-> torch.Tensor:
    # Check constraints.
    if use_int8_w4a8:
         assert block_shape[0] == 0 and block_shape[1] == 64, "[ERROR]ASM Fused MoE only support w4a8 block_shape=64 now."

    if use_shuffle:
        assert use_fp8_w8a8 or use_int8_w8a8 or (not use_int4_w4a16 and not use_int4_w4a16), "[ERROR]ASM Fused MoE only support f8 now."

    if use_int4_w4a16 or use_int8_w4a8:
        assert hidden_states.shape[1] // 2 == w1.shape[
            2], "Hidden size mismatch"
    else:
        assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch"

    assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
    assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
    assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
    assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
    assert hidden_states.dtype in [
        torch.float32, torch.float16, torch.bfloat16, torch.int8, torch.float8_e4m3fn
    ]

    num_tokens, _ = hidden_states.shape
    E, N, _ = w1.shape
    _, model_dim, inter_dim = w2.shape
    if global_num_experts == -1:
        global_num_experts = E
    top_k_num = topk_ids.shape[1]
    # We execute the fused_moe kernel in chunks to circumvent this issue:
    # https://github.com/vllm-project/vllm/issues/5938
    # need to change according to token
    CHUNK_SIZE = 65536
    M = min(num_tokens, CHUNK_SIZE)

    out_hidden_states = torch.empty((num_tokens, model_dim), dtype=dtype, device=hidden_states.device)
    for chunk in range((num_tokens // CHUNK_SIZE) + 1):
        begin_chunk_idx, end_chunk_idx = (chunk * CHUNK_SIZE,
                                          min((chunk + 1) * CHUNK_SIZE,
                                              num_tokens))
        curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
        tokens_in_chunk, _ = curr_hidden_states.shape

        if tokens_in_chunk == 0:
            break


        curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
        curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]

        #FIXME: just for EP Accuracy Test
        if expert_map is not None:
            d_w1_out = torch.zeros((curr_hidden_states.size(0) * top_k_num, w1.size(1)), dtype=dtype, device=curr_hidden_states.device)
            if use_int4_w4a16 or use_int8_w4a8:
                d_silu = torch.zeros((curr_hidden_states.size(0) * top_k_num, w2.size(2) * 2), dtype=dtype, device=curr_hidden_states.device)
            else:
                d_silu = torch.zeros((curr_hidden_states.size(0) * top_k_num, w2.size(2)), dtype=dtype, device=curr_hidden_states.device)
            d_w2_out = torch.zeros((curr_hidden_states.size(0), top_k_num, w2.size(1)), dtype=dtype, device=curr_hidden_states.device)
        else:
            d_w1_out = torch.empty((curr_hidden_states.size(0) * top_k_num, w1.size(1)), dtype=dtype, device=curr_hidden_states.device)
            if use_int4_w4a16 or use_int8_w4a8:
                d_silu = torch.empty((curr_hidden_states.size(0) * top_k_num, w2.size(2) * 2), dtype=dtype,device=curr_hidden_states.device)
            else:
                d_silu = torch.empty((curr_hidden_states.size(0) * top_k_num, w2.size(2)), dtype=dtype, device=curr_hidden_states.device)
            d_w2_out = torch.empty((curr_hidden_states.size(0), top_k_num, w2.size(1)), dtype=dtype, device=curr_hidden_states.device)

        arch = get_gfx()
        cu_num = get_cu_num()
        odtype = 0
        if dtype == torch.bfloat16:
            odtype = 1
        if use_persist:
            if persist_cu <= 0 or persist_cu >= cu_num:
               persist_cu = cu_num
        else:
            persist_cu = 0
        # INT4 w4a16
        if use_int4_w4a16:
            if solution_id is None:
                solution_id = get_moe_asm_solution(arch, tokens_in_chunk, N/2, w1.size(2)*2, E, top_k_num, MoeQuantType.INT4_W4A16)
            if block_shape is not None and block_shape[1] == 32:
                config = decode_sol_w4a16_gw32()
            else:
                config = decode_sol_w4a16(solution_id)
            sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, tokens_positions_per_expert, moe_buf = (
                moe_sorting_ck(curr_topk_ids, curr_topk_weights, global_num_experts, model_dim, out_hidden_states[begin_chunk_idx:end_chunk_idx], config["BLOCK_SIZE_M"], expert_map)
            )
            if print_log():
                print(f"Asm Moe Size: chunk:{chunk}, arch:{arch}, quant:{MoeQuantType.INT4_W4A16}, tokens:{tokens_in_chunk}, inter_dim:{int(N/2)}, model_dim:{w1.size(2)*2}, expert:{E}, topk:{top_k_num}")
                print(f"solution:{solution_id}, shuffle:{use_shuffle}, persist:{persist_cu}")
                if solution_id == "default":
                    print(f">>> Warning: No matching config pattern found, using default asm solution.")
            solution_id = None

            if dtype == torch.bfloat16:
                if block_shape is not None and block_shape[1] == 32:
                    aiter.asm_fmoe_stage1(d_w1_out,
                                    curr_hidden_states, 
                                    w1, 
                                    w2, 
                                    sorted_ids, 
                                    sorted_weights, 
                                    sorted_expert_ids,
                                    num_valid_ids, 
                                    top_k_num, 
                                    w1_scale, 
                                    w1_scale, 
                                    w1_zp, 
                                    4,
                                    config["SOL_ID1"],
                                    config["BLOCK_SIZE_M"])
                else:
                    aiter.asm_fmoe_stage1(d_w1_out,
                                    curr_hidden_states, 
                                    w1, 
                                    w2, 
                                    sorted_ids, 
                                    sorted_weights, 
                                    sorted_expert_ids,
                                    num_valid_ids, 
                                    top_k_num, 
                                    w1_scale, 
                                    w1_scale, 
                                    w1_zp, 
                                    3,
                                    config["SOL_ID1"],
                                    config["BLOCK_SIZE_M"])
            else:
                aiter.asm_fmoe_stage1(d_w1_out,
                                curr_hidden_states, 
                                w1, 
                                w2, 
                                sorted_ids, 
                                sorted_weights, 
                                sorted_expert_ids,
                                num_valid_ids, 
                                top_k_num, 
                                w1_scale, 
                                w1_scale, 
                                w1_zp, 
                                2,
                                config["SOL_ID1"],
                                config["BLOCK_SIZE_M"])  
            if activation == "silu":
                triton_silu_and_mul(d_silu,d_w1_out)
                # silu_and_mul(d_silu,d_w1_out)
            elif activation == "gelu":
                triton_gelu_and_mul(d_silu,d_w1_out)
                # gelu_and_mul(d_silu,d_w1_out)
            else:
                raise ValueError(f"Unsupported FusedMoe activation: {activation}")
            if dtype == torch.bfloat16:
                if block_shape is not None and block_shape[1] == 32:
                    aiter.asm_fmoe_stage2(d_w2_out, 
                                            d_silu, 
                                            w1, 
                                            w2, 
                                            sorted_ids, 
                                            sorted_weights, 
                                            sorted_expert_ids, 
                                            num_valid_ids, 
                                            top_k_num, 
                                            w2_scale, 
                                            w2_scale, 
                                            w2_zp, 
                                            4,
                                            config["SOL_ID2"],
                                            config["BLOCK_SIZE_M"])
                else:
                    aiter.asm_fmoe_stage2(d_w2_out, 
                                            d_silu, 
                                            w1, 
                                            w2, 
                                            sorted_ids, 
                                            sorted_weights, 
                                            sorted_expert_ids, 
                                            num_valid_ids, 
                                            top_k_num, 
                                            w2_scale, 
                                            w2_scale, 
                                            w2_zp, 
                                            3,
                                            config["SOL_ID2"],
                                            config["BLOCK_SIZE_M"])
            else:
                aiter.asm_fmoe_stage2(d_w2_out, 
                                        d_silu, 
                                        w1, 
                                        w2, 
                                        sorted_ids, 
                                        sorted_weights, 
                                        sorted_expert_ids, 
                                        num_valid_ids, 
                                        top_k_num, 
                                        w2_scale, 
                                        w2_scale, 
                                        w2_zp, 
                                        2,
                                        config["SOL_ID2"],
                                        config["BLOCK_SIZE_M"])
        #int8 channel wise
        elif use_int8_w8a8 and per_channel_quant:
            if solution_id is None:
                solution_id = get_moe_asm_solution(arch, tokens_in_chunk, N/2, w1.size(2), E, top_k_num, MoeQuantType.INT8_W8A8_C, use_shuffle)
            config = decode_sol_w8a8_c(solution_id)
            if persist_cu == cu_num:
                calculate_persist_groups(persist_cu, config, MoeQuantType.INT8_W8A8_C)
            else:
                config["PERSIST_GROUP1"] = persist_cu
                config["PERSIST_GROUP2"] = persist_cu

            sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, tokens_positions_per_expert, moe_buf = (
                moe_sorting_ck(curr_topk_ids, curr_topk_weights, global_num_experts, model_dim, out_hidden_states[begin_chunk_idx:end_chunk_idx], config["BLOCK_SIZE_M"], expert_map)
            )
            if print_log():
                print(f"Asm Moe Size: chunk:{chunk}, arch:{arch}, quant:{MoeQuantType.INT8_W8A8_C}, tokens:{tokens_in_chunk}, inter_dim:{int(N/2)}, model_dim:{w1.size(2)}, expert:{E}, topk:{top_k_num}")
                print(f"solution:{solution_id}, shuffle:{use_shuffle}, persist:{persist_cu}")
                if solution_id== "default":
                    print(f">>> Warning: No matching config pattern found, using default asm solution.")
            solution_id = None

            if curr_hidden_states.dtype == torch.float16 or curr_hidden_states.dtype == torch.bfloat16:
                input_q,input_scale = per_token_quant_hip(curr_hidden_states)
                #input_q,input_scale = per_token_quant_int8(curr_hidden_states)
            else:
                input_q,input_scale = curr_hidden_states,a1_scale
            aiter.asm_fmoe_a8(d_w1_out,
                    input_q, 
                    w1, 
                    w1, 
                    sorted_ids, 
                    curr_topk_weights, 
                    sorted_expert_ids,
                    num_valid_ids, 
                    top_k_num, 
                    input_scale, 
                    w1_scale, 
                    w1_zp, 
                    0,
                    config["SOL_ID1"],
                    odtype,
                    config["PERSIST_GROUP1"],
                    use_shuffle)
            if activation == "silu":
                triton_silu_and_mul(d_silu,d_w1_out)
                # silu_and_mul(d_silu,d_w1_out)
            elif activation == "gelu":
                triton_gelu_and_mul(d_silu,d_w1_out)
                # gelu_and_mul(d_silu,d_w1_out)
            else:
                raise ValueError(f"Unsupported FusedMoe activation: {activation}")

            bridge_q,bridge_scale = per_token_quant_hip(d_silu)
            #bridge_q,bridge_scale = per_token_quant_int8(d_silu)
            aiter.asm_fmoe_a8(d_w2_out,
                    bridge_q, 
                    w2, 
                    w2, 
                    sorted_ids, 
                    curr_topk_weights, 
                    sorted_expert_ids,
                    num_valid_ids, 
                    top_k_num, 
                    bridge_scale, 
                    w2_scale, 
                    w2_zp, 
                    1,
                    config["SOL_ID2"],
                    odtype,
                    config["PERSIST_GROUP2"],
                    use_shuffle)
        #w4a8 block wise = 64
        elif use_int8_w4a8:
            if solution_id is None:
                solution_id = get_moe_asm_solution(arch, tokens_in_chunk, N/2, w1.size(2)*2, E, top_k_num, MoeQuantType.INT4_W4A8)
            config = decode_sol_0(solution_id)
            if persist_cu == cu_num:
                calculate_persist_groups(persist_cu, config, MoeQuantType.INT4_W4A8)
            else:
                config["PERSIST_GROUP1"] = persist_cu
                config["PERSIST_GROUP2"] = persist_cu

            sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, tokens_positions_per_expert, moe_buf = (
                moe_sorting_ck(curr_topk_ids, curr_topk_weights, global_num_experts, model_dim, out_hidden_states[begin_chunk_idx:end_chunk_idx], config["BLOCK_SIZE_M"], expert_map)
            )
            if print_log():
                print(f"Asm Moe Size: chunk:{chunk}, arch:{arch}, quant:{MoeQuantType.INT4_W4A8}, tokens:{tokens_in_chunk}, inter_dim:{int(N/2)}, model_dim:{w1.size(2)*2}, expert:{E}, topk:{top_k_num}")
                print(f"solution:{solution_id}, shuffle:{use_shuffle}, persist:{persist_cu}")
                if solution_id== "default":
                    print(f">>> Warning: No matching config pattern found, using default asm solution.")
            solution_id = None

            if curr_hidden_states.dtype == torch.float16 or curr_hidden_states.dtype==torch.bfloat16:
                input_q,input_scale = per_token_group_quant_int8(curr_hidden_states, block_shape[1])
            else:
                input_q,input_scale = curr_hidden_states,a1_scale
            #quant_func = get_hip_quant(QuantType.per_1x64)
            #input_q,input_scale = quant_func(curr_hidden_states, quant_dtype=dtypes.i8)
            # input_q,input_scale = per_token_group_quant_int8(curr_hidden_states, block_shape[1])

            aiter.asm_fmoe_a8(d_w1_out,
                    input_q, 
                    w1, 
                    w1, 
                    sorted_ids, 
                    curr_topk_weights, 
                    sorted_expert_ids,
                    num_valid_ids, 
                    top_k_num, 
                    input_scale, 
                    w1_scale, 
                    w1_zp, 
                    10,
                    config["SOL_ID1"],
                    odtype,
                    config["PERSIST_GROUP1"])
            if activation == "silu":
                triton_silu_and_mul(d_silu,d_w1_out)
                # silu_and_mul(d_silu,d_w1_out)
            elif activation == "gelu":
                triton_gelu_and_mul(d_silu,d_w1_out)
                # gelu_and_mul(d_silu,d_w1_out)
            else:
                raise ValueError(f"Unsupported FusedMoe activation: {activation}")

            #quant_func = get_hip_quant(QuantType.per_1x64)
            #bridge_q,bridge_scale = quant_func(d_silu, quant_dtype=dtypes.i8)
            bridge_q,bridge_scale = per_token_group_quant_int8(d_silu, block_shape[1])

            aiter.asm_fmoe_a8(d_w2_out,
                    bridge_q, 
                    w2, 
                    w2, 
                    sorted_ids, 
                    curr_topk_weights, 
                    sorted_expert_ids,
                    num_valid_ids, 
                    top_k_num, 
                    bridge_scale, 
                    w2_scale, 
                    w2_zp,
                    11,
                    config["SOL_ID2"],
                    odtype,
                    config["PERSIST_GROUP2"])
        #int8 block wise = 128
        elif use_int8_w8a8:
            if solution_id is None:
                solution_id = get_moe_asm_solution(arch, tokens_in_chunk, N/2, w1.size(2), E, top_k_num, MoeQuantType.INT8_W8A8, use_shuffle)
            config = decode_sol_0(solution_id, use_shuffle)
            if persist_cu == cu_num:
                calculate_persist_groups(persist_cu, config, MoeQuantType.INT8_W8A8)
            else:
                config["PERSIST_GROUP1"] = persist_cu
                config["PERSIST_GROUP2"] = persist_cu

            sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, tokens_positions_per_expert, moe_buf = (
                moe_sorting_ck(curr_topk_ids, curr_topk_weights, global_num_experts, model_dim, out_hidden_states[begin_chunk_idx:end_chunk_idx], config["BLOCK_SIZE_M"], expert_map)
            )
            if print_log():
                print(f"Asm Moe Size: chunk:{chunk}, arch:{arch}, quant:{MoeQuantType.INT8_W8A8}, tokens:{tokens_in_chunk}, inter_dim:{int(N/2)}, model_dim:{w1.size(2)}, expert:{E}, topk:{top_k_num}")
                print(f"solution:{solution_id}, shuffle:{use_shuffle}, persist:{persist_cu}")
                if solution_id== "default":
                    print(f">>> Warning: No matching config pattern found, using default asm solution.")
            solution_id = None

            if curr_hidden_states.dtype == torch.float16 or curr_hidden_states.dtype==torch.bfloat16:
                input_q,input_scale = per_block_quant_wrapper((1,block_shape[1]))(per_token_quant_hip)(curr_hidden_states, quant_dtype=torch.int8)
            else:
                input_q,input_scale = curr_hidden_states,a1_scale
            aiter.asm_fmoe_a8(d_w1_out,
                    input_q, 
                    w1, 
                    w1, 
                    sorted_ids, 
                    curr_topk_weights, 
                    sorted_expert_ids,
                    num_valid_ids, 
                    top_k_num, 
                    input_scale, 
                    w1_scale, 
                    w1_zp, 
                    2,
                    config["SOL_ID1"],
                    odtype,
                    config["PERSIST_GROUP1"],
                    use_shuffle)
            if activation == "silu":
                triton_silu_and_mul(d_silu,d_w1_out)
                # silu_and_mul(d_silu,d_w1_out)
            elif activation == "gelu":
                triton_gelu_and_mul(d_silu,d_w1_out)
                # gelu_and_mul(d_silu,d_w1_out)
            else:
                raise ValueError(f"Unsupported FusedMoe activation: {activation}")

            #FIXME: aiter quant method performance is little worse than triton. Change it latter!!
            bridge_q, bridge_scale = per_block_quant_wrapper((1,block_shape[1]))(per_token_quant_hip)(d_silu, quant_dtype=torch.int8)
            # bridge_q,bridge_scale = per_token_group_quant_int8(d_silu, block_shape[1])
            aiter.asm_fmoe_a8(d_w2_out,
                    bridge_q, 
                    w2, 
                    w2, 
                    sorted_ids, 
                    curr_topk_weights, 
                    sorted_expert_ids,
                    num_valid_ids, 
                    top_k_num, 
                    bridge_scale, 
                    w2_scale, 
                    w2_zp, 
                    3,
                    config["SOL_ID2"],
                    odtype,
                    config["PERSIST_GROUP2"],
                    use_shuffle)
        #f8 channel wise
        elif use_fp8_w8a8 and per_channel_quant:
            if solution_id is None:
                solution_id = get_moe_asm_solution(arch, tokens_in_chunk, N/2, w1.size(2), E, top_k_num, MoeQuantType.F8_W8A8_C, use_shuffle)
            config = decode_sol_w8a8_c(solution_id)
            if persist_cu == cu_num:
                calculate_persist_groups(persist_cu, config, MoeQuantType.F8_W8A8_C)
            else:
                config["PERSIST_GROUP1"] = persist_cu
                config["PERSIST_GROUP2"] = persist_cu

            sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, tokens_positions_per_expert, moe_buf = (
                moe_sorting_ck(curr_topk_ids, curr_topk_weights, global_num_experts, model_dim, out_hidden_states[begin_chunk_idx:end_chunk_idx], config["BLOCK_SIZE_M"], expert_map)
            )
            if print_log():
                print(f"Asm Moe Size: chunk:{chunk}, arch:{arch}, quant:{MoeQuantType.F8_W8A8_C}, tokens:{tokens_in_chunk}, inter_dim:{int(N/2)}, model_dim:{w1.size(2)}, expert:{E}, topk:{top_k_num}")
                print(f"solution:{solution_id}, shuffle:{use_shuffle}, persist:{persist_cu}")
                if solution_id== "default":
                    print(f">>> Warning: No matching config pattern found, using default asm solution.")
            solution_id = None

            if curr_hidden_states.dtype == torch.float16 or curr_hidden_states.dtype==torch.bfloat16:
                input_q,input_scale = per_token_quant_hip(curr_hidden_states, quant_dtype=torch.float8_e4m3fn)
            else:
                input_q,input_scale = curr_hidden_states,a1_scale
            aiter.asm_fmoe_a8(d_w1_out,
                    input_q, 
                    w1, 
                    w1, 
                    sorted_ids, 
                    curr_topk_weights, 
                    sorted_expert_ids,
                    num_valid_ids, 
                    top_k_num, 
                    input_scale, 
                    w1_scale, 
                    w1_zp, 
                    4,
                    config["SOL_ID1"],
                    odtype,
                    config["PERSIST_GROUP1"],
                    use_shuffle)
            if activation == "silu":
                triton_silu_and_mul(d_silu,d_w1_out)
                # silu_and_mul(d_silu,d_w1_out)
            elif activation == "gelu":
                triton_gelu_and_mul(d_silu,d_w1_out)
                # gelu_and_mul(d_silu,d_w1_out)
            else:
                raise ValueError(f"Unsupported FusedMoe activation: {activation}")

            bridge_q,bridge_scale= per_token_quant_hip(d_silu, quant_dtype=torch.float8_e4m3fn)
            aiter.asm_fmoe_a8(d_w2_out,
                    bridge_q, 
                    w2, 
                    w2, 
                    sorted_ids, 
                    curr_topk_weights, 
                    sorted_expert_ids,
                    num_valid_ids, 
                    top_k_num, 
                    bridge_scale, 
                    w2_scale, 
                    w2_zp, 
                    5,
                    config["SOL_ID2"],
                    odtype,
                    config["PERSIST_GROUP2"],
                    use_shuffle)
        #f8 block wise = 128
        elif use_fp8_w8a8:
            if solution_id is None:
                solution_id = get_moe_asm_solution(arch, tokens_in_chunk, N/2, w1.size(2), E, top_k_num, MoeQuantType.F8_W8A8, use_shuffle)
            config = decode_sol_0(solution_id, use_shuffle)
            if persist_cu == cu_num:
                calculate_persist_groups(persist_cu, config, MoeQuantType.F8_W8A8)
            else:
                config["PERSIST_GROUP1"] = persist_cu
                config["PERSIST_GROUP2"] = persist_cu

            sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, tokens_positions_per_expert, moe_buf = (
                moe_sorting_ck(curr_topk_ids, curr_topk_weights, global_num_experts, model_dim, out_hidden_states[begin_chunk_idx:end_chunk_idx], config["BLOCK_SIZE_M"], expert_map)
            )
            if print_log():
                print(f"Asm Moe Size: chunk:{chunk}, arch:{arch}, quant:{MoeQuantType.F8_W8A8}, tokens:{tokens_in_chunk}, inter_dim:{int(N/2)}, model_dim:{w1.size(2)}, expert:{E}, topk:{top_k_num}")
                print(f"solution:{solution_id}, shuffle:{use_shuffle}, persist:{persist_cu}")
                if solution_id== "default":
                    print(f">>> Warning: No matching config pattern found, using default asm solution.")
            solution_id = None

            if curr_hidden_states.dtype == torch.float16 or curr_hidden_states.dtype==torch.bfloat16:
                input_q,input_scale = per_block_quant_wrapper((1,block_shape[1]))(per_token_quant_hip)(curr_hidden_states, quant_dtype=torch.float8_e4m3fn)
            else:
                input_q,input_scale = curr_hidden_states,a1_scale
            aiter.asm_fmoe_a8(d_w1_out,
                    input_q, 
                    w1, 
                    w1, 
                    sorted_ids, 
                    curr_topk_weights, 
                    sorted_expert_ids,
                    num_valid_ids, 
                    top_k_num, 
                    input_scale, 
                    w1_scale, 
                    w1_zp, 
                    6,
                    config["SOL_ID1"],
                    odtype,
                    config["PERSIST_GROUP1"],
                    use_shuffle)
            if activation == "silu":
                triton_silu_and_mul(d_silu,d_w1_out)
                # silu_and_mul(d_silu,d_w1_out)
            elif activation == "gelu":
                triton_gelu_and_mul(d_silu,d_w1_out)
                # gelu_and_mul(d_silu,d_w1_out)
            else:
                raise ValueError(f"Unsupported FusedMoe activation: {activation}")

            bridge_q,bridge_scale = per_block_quant_wrapper((1,block_shape[1]))(per_token_quant_hip)(d_silu, quant_dtype=torch.float8_e4m3fn)
            aiter.asm_fmoe_a8(d_w2_out,
                    bridge_q, 
                    w2, 
                    w2, 
                    sorted_ids, 
                    curr_topk_weights, 
                    sorted_expert_ids,
                    num_valid_ids, 
                    top_k_num, 
                    bridge_scale, 
                    w2_scale, 
                    w2_zp, 
                    7,
                    config["SOL_ID2"],
                    odtype,
                    config["PERSIST_GROUP2"],
                    use_shuffle)
        #
        else:
            if solution_id is None:
                solution_id = get_moe_asm_solution(arch, tokens_in_chunk, N/2, w1.size(2), E, top_k_num, MoeQuantType.NO_QUANT, use_shuffle)
            config = decode_sol_w8a8_c(solution_id)
            if persist_cu == cu_num:
                calculate_persist_groups(persist_cu, config, MoeQuantType.NO_QUANT)
            else:
                config["PERSIST_GROUP1"] = persist_cu
                config["PERSIST_GROUP2"] = persist_cu

            sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, tokens_positions_per_expert, moe_buf = (
                moe_sorting_ck(curr_topk_ids, curr_topk_weights, global_num_experts, model_dim, out_hidden_states[begin_chunk_idx:end_chunk_idx], config["BLOCK_SIZE_M"], expert_map)
            )
            if print_log():
                print(f"Asm Moe Size: chunk:{chunk}, arch:{arch}, quant:{MoeQuantType.NO_QUANT}, tokens:{tokens_in_chunk}, inter_dim:{int(N/2)}, model_dim:{w1.size(2)}, expert:{E}, topk:{top_k_num}")
                print(f"solution:{solution_id}, shuffle:{use_shuffle}, persist:{persist_cu}")
                if solution_id== "default":
                    print(f">>> Warning: No matching config pattern found, using default asm solution.")
            solution_id = None

            aiter.asm_fmoe_a8(d_w1_out,
                    curr_hidden_states, 
                    w1, 
                    w1, 
                    sorted_ids, 
                    curr_topk_weights, 
                    sorted_expert_ids,
                    num_valid_ids, 
                    top_k_num, 
                    None, 
                    None, 
                    None, 
                    20,
                    config["SOL_ID1"],
                    odtype,
                    config["PERSIST_GROUP1"],
                    use_shuffle)
            #return d_w1_out
            if activation == "silu":
                triton_silu_and_mul(d_silu,d_w1_out)
                # silu_and_mul(d_silu,d_w1_out)
            elif activation == "gelu":
                triton_gelu_and_mul(d_silu,d_w1_out)
                # gelu_and_mul(d_silu,d_w1_out)
            else:
                raise ValueError(f"Unsupported FusedMoe activation: {activation}")

            aiter.asm_fmoe_a8(d_w2_out,
                    d_silu, 
                    w2, 
                    w2, 
                    sorted_ids, 
                    curr_topk_weights, 
                    sorted_expert_ids,
                    num_valid_ids, 
                    top_k_num, 
                    None, 
                    None, 
                    None, 
                    21,
                    config["SOL_ID2"],
                    odtype,
                    config["PERSIST_GROUP2"],
                    use_shuffle)
        triton_moe_sum(d_w2_out, moe_buf if not inplace else hidden_states, routed_scaling_factor)

    return out_hidden_states if not inplace else hidden_states

@lru_cache(maxsize=1)
def print_log():
    value = os.getenv("ASM_MOE_LOG")
    if value is not None:
        return True
    else:
        return False

class MoeQuantType:
    NO_QUANT = "no_quant"
    INT4_W4A16 = "int4_w4a16"
    INT4_W4A8 = "int4_w4a8"
    INT8_W8A8 = "int8_w8a8_block"
    INT8_W8A8_C = "int8_w8a8_channel"
    F8_W8A8 = "f8_w8a8_block"
    F8_W8A8_C = "f8_w8a8_channel"

    ALL_TYPES = [NO_QUANT, INT4_W4A16, INT4_W4A8, INT8_W8A8, INT8_W8A8_C, F8_W8A8, F8_W8A8_C]

    @classmethod
    def is_valid(cls, qtype_str: str) -> bool:
        return qtype_str in cls.ALL_TYPES

    @classmethod
    def get_default(cls) -> str:
        return cls.NO_QUANT

_cached_data_by_quant = {}
_data_lock = threading.RLock()

CSV_FILE_MAPPING = {
    "no_quant": "tuned_fmoe_asm.csv",
    "int4_w4a16": "tuned_fmoe_asm_w4a16.csv",
    "int4_w4a8": "tuned_fmoe_asm_w4a8_group.csv",
    "int8_w8a8_block": "tuned_fmoe_asm_w8a8_group.csv",
    "int8_w8a8_channel": "tuned_fmoe_asm_w8a8_channel.csv",
    "f8_w8a8_block": "tuned_fmoe_asm_w8a8_group.csv",
    "f8_w8a8_channel": "tuned_fmoe_asm_w8a8_channel.csv",
    "no_quant_s": "tuned_fmoe_asm_shuffle.csv",
    "int8_w8a8_block_s": "tuned_fmoe_asm_w8a8_group_shuffle.csv",
    "int8_w8a8_channel_s": "tuned_fmoe_asm_w8a8_channel_shuffle.csv",
    "f8_w8a8_block_s": "tuned_fmoe_asm_w8a8_group_shuffle.csv",
    "f8_w8a8_channel_s": "tuned_fmoe_asm_w8a8_channel_shuffle.csv",
}

def get_csv_path(quant_type):
    if quant_type in CSV_FILE_MAPPING:
        filename = CSV_FILE_MAPPING[quant_type]
    else:
        filename = f"{quant_type}.csv"
    return os.path.join(AITER_ROOT_DIR, "aiter", "configs", filename)

def load_and_cache_csv_for_quant(quant_type, use_shuffle):

    global _cached_data_by_quant

    if use_shuffle == 1:
        quant_type_file = quant_type + "_s"
    else:
        quant_type_file = quant_type
    csv_path = get_csv_path(quant_type_file)
    
    if not os.path.exists(csv_path):
        print(f"Asm moe tuned csv not found: {csv_path}")
        return False
    
    with _data_lock:
        try:
            if (quant_type_file in _cached_data_by_quant):
                return True

            print(f"Load asm moe tuned csv: {csv_path}")
            moe_asm_cfg = pd.read_csv(csv_path)
            
            # Group by key parameters
            group_cols = ['arch', 'inter_dim', 'model_dim', 'expert', 'topk', 'quant_type']
            cached_groups = {}
            
            for group_key, group_df in moe_asm_cfg.groupby(group_cols):
                # Ensure tokens are integers and sorted in ascending order
                group_df = group_df.sort_values('token')
                tokens_array = group_df['token'].values.astype(int)
                sol_ids_array = group_df['sol_id'].values
                
                # Store in cache
                cached_groups[group_key] = (tokens_array, sol_ids_array, group_df)

            # Update cache
            _cached_data_by_quant[quant_type_file] = cached_groups

            return True

        except Exception as e:
            print(f"Load asm moe tuned csv failed {csv_path}: {e}")
            return False

@lru_cache(maxsize=4096)
def get_moe_asm_solution(
    arch,
    token,
    inter_dim,
    model_dim,
    expert,
    topk,
    quant_type,
    use_shuffle=0,
    q_size_n=0,
    q_size_k=0
):
    if not load_and_cache_csv_for_quant(quant_type, use_shuffle):
        return "default"

    with _data_lock:
        if use_shuffle == 1:
            quant_type_file = quant_type + "_s"
        else:
            quant_type_file = quant_type
        cached_groups = _cached_data_by_quant.get(quant_type_file)
        if cached_groups is None:
            return "default"
        
        cache_key = (str(arch), int(inter_dim), int(model_dim), 
                     int(expert), int(topk), str(quant_type))
        
        if cache_key not in cached_groups:
            return "default"
        
        tokens_array, sol_ids_array, _ = cached_groups[cache_key]
        token_int = int(token)
        n = len(tokens_array)
        
        if n == 0:
            return "default"
        
        if token_int <= tokens_array[0]:
            return sol_ids_array[0]
        if token_int >= tokens_array[-1]:
            return sol_ids_array[-1]
        
        # Binary search
        idx = bisect_left(tokens_array, token_int)
        
        # Exact match
        if idx < n and tokens_array[idx] == token_int:
            return sol_ids_array[idx]
        
        # No exact match, find the closest token
        left_idx = idx - 1
        right_idx = idx
        
        left_dist = token_int - tokens_array[left_idx]
        right_dist = tokens_array[right_idx] - token_int

        if left_dist <= right_dist:
            return sol_ids_array[left_idx]
        else:
            return sol_ids_array[right_idx]
def decode_sol_w4a16(solution) -> Dict[str, int]:
    if solution == "default":
        config = {
            "SOL_ID1": 11002,
            "SOL_ID2": 21002,
            "BLOCK_SIZE_M": 32,
        }
        return config

    parts = solution.split("+")
    if len(parts) == 2:
        sol_id1 = int(parts[0])
        sol_id2 = int(parts[1])
    else:
        raise ValueError("Invalid solution_id")

    if 10000 <= sol_id1 <= 10999:
        block_size_m = 16
    elif 11000 <= sol_id1 <= 11999:
        block_size_m = 32
    elif 12000 <= sol_id1 <= 12999:
        block_size_m = 64
    elif 13000 <= sol_id1 <= 13999:
        block_size_m = 128
    else:
        raise ValueError(f"key1 value {sol_id1} is not in the expected ranges (10000-13999)")
    config = {
        "SOL_ID1": sol_id1,
        "SOL_ID2": sol_id2,
        "BLOCK_SIZE_M": block_size_m,
    }
    return config

def decode_sol_w4a16_gw32() -> Dict[str, int]:
    config = {
        "SOL_ID1": 50032,
        "SOL_ID2": 60032,
        "BLOCK_SIZE_M": 32, 
    }
    return config


def decode_sol_0(solution, use_shuffle=0) -> Dict[str, int]:
    if solution == "default":
        config = {
            "SOL_ID1": 10000,
            "SOL_ID2": 20000,
            "BLOCK_SIZE_M": 16,
        }
        if use_shuffle:
            config["SOL_ID1"] = 10001
        return config

    parts = solution.split("+")
    if len(parts) == 2:
        sol_id1 = int(parts[0])
        sol_id2 = int(parts[1])
    else:
        raise ValueError("Invalid solution_id")

    if 10000 <= sol_id1 <= 10999:
        block_size_m = 16
    elif 11000 <= sol_id1 <= 11999:
        block_size_m = 32
    elif 12000 <= sol_id1 <= 12999:
        block_size_m = 64
    elif 13000 <= sol_id1 <= 13999:
        block_size_m = 128
    elif 14000 <= sol_id1 <= 14999:
        block_size_m = 256
    else:
        raise ValueError(f"key1 value {sol_id1} is not in the expected ranges (10000-13999)")
    config = {
        "SOL_ID1": sol_id1,
        "SOL_ID2": sol_id2,
        "BLOCK_SIZE_M": block_size_m,
    }
    return config

def decode_sol_w8a8_c(solution) -> Dict[str, int]:

    if solution == "default":
        config = {
            "SOL_ID1": 11000,
            "SOL_ID2": 21001,
            "BLOCK_SIZE_M": 32,
        }
        return config

    parts = solution.split("+")
    if len(parts) == 2:
        sol_id1 = int(parts[0])
        sol_id2 = int(parts[1])
    else:
        raise ValueError("Invalid solution_id")

    if 10000 <= sol_id1 <= 10999:
        block_size_m = 16
    elif 11000 <= sol_id1 <= 11999:
        block_size_m = 32
    elif 12000 <= sol_id1 <= 12999:
        block_size_m = 64
    elif 13000 <= sol_id1 <= 13999:
        block_size_m = 128
    else:
        raise ValueError(f"sol_id1 value {sol_id1} is not in the expected ranges (10000-12999)")
    config = {
        "SOL_ID1": sol_id1,
        "SOL_ID2": sol_id2,
        "BLOCK_SIZE_M": block_size_m,
    }
    return config

def calculate_persist_groups(persist_cu, config, quant_type):

    # Maximum number of sol_id workgroups executable per CU.
    if quant_type == MoeQuantType.INT4_W4A8:
        sol_id_table = {
            20000: 4,
            21000: 4,
            22000: 2,
            23000: 2,
        }
    elif quant_type == MoeQuantType.INT8_W8A8_C or quant_type == MoeQuantType.F8_W8A8_C:
        sol_id_table = {
            20000: 7,
            20001: 5,
            21005: 7,
            21006: 5,
            22003: 4,
            22004: 3,
            23000: 3,
            23001: 2,
            23002: 2,
        }
    elif quant_type == MoeQuantType.INT8_W8A8 or quant_type == MoeQuantType.F8_W8A8:
        sol_id_table = {
            20000: 4,
            21000: 4,
            22000: 2,
            23000: 2,
            23001: 4,
            24000: 2,
            24001: 4,
        }
    else:
        return

    for i in [1, 2]:
        if config[f"SOL_ID{i}"] in sol_id_table:
            config[f"PERSIST_GROUP{i}"] = persist_cu * sol_id_table[config[f'SOL_ID{i}']]
        else:
            config[f"PERSIST_GROUP{i}"] = persist_cu