import torch
import os
from typing import Optional, List
import functools
from bisect import bisect_left
import aiter
from aiter import ActivationType, QuantType, dtypes
from aiter.jit.core import AITER_ROOT_DIR
from aiter import ck_moe, ck_shuffle_moe
from aiter.jit.utils.torch_guard import torch_compile_guard
from aiter.jit.utils.chip_info import get_gfx
from aiter.fused_moe import moe_sorting
from aiter import per_token_quant_hip, per_block_quant_wrapper


BLOCK_SIZE_M = 32

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"
    
    ALL_TYPES = [NO_QUANT, INT4_W4A16, INT4_W4A8, INT8_W8A8, INT8_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


ck_tuned_file = os.path.join(AITER_ROOT_DIR, "aiter", "configs", "ck_tune", "tuned_fmoe_ck.csv")
ck_tuned_int8_w8a8_group_file = os.path.join(AITER_ROOT_DIR, "aiter", "configs", "ck_tune", "tuned_fmoe_ck_int8_w8a8_group.csv")

moe_ck_cfg = None
moe_ck_noquant_cfg = None
moe_ck_int8_w8a8_group_cfg = None
moe_ck_noquant_index = None
moe_ck_int8_w8a8_group_index = None
current_quant_type = None

def get_moe_ck_solution(
    indtype,
    token,
    inter_dim,
    model_dim,
    expert,
    topk,
    quant_type,
    q_size_n=0,
    q_size_k=0
):

    def get_moe_cfg(ck_tuned_file):
        import pandas as pd
        try:
            moe_cfg = pd.read_csv(ck_tuned_file)
        except Exception as e:
            print(f">>> Warning: Failed to read config file {ck_tuned_file}: {e}")
            return None
        return moe_cfg
    
    global moe_ck_cfg
    if moe_ck_cfg is None:
        moe_ck_cfg = get_moe_cfg(ck_tuned_file)
        if moe_ck_cfg is None:
            print(f">>> Warning: config file {ck_tuned_file} is not found, using default ck solution.")
            return functools.partial(ck_moe, solution_id = 0)
        
    mask = (
        (moe_ck_cfg["indtype"] == str(indtype)) &
        (moe_ck_cfg["inter_dim"] == inter_dim) &
        (moe_ck_cfg["model_dim"] == model_dim) &
        (moe_ck_cfg["expert"] == expert) &
        (moe_ck_cfg["topk"] == topk) &
        (moe_ck_cfg["quant_type"] == str(quant_type)) &
        (moe_ck_cfg["q_size_n"] == q_size_n) &
        (moe_ck_cfg["q_size_k"] == q_size_k)
    )
    matching_configs = moe_ck_cfg[mask]
    if matching_configs.empty:
        sol_id = 0
        print(f">>> Warning: No matching config pattern found, using default ck solution.")
        return functools.partial(ck_moe, solution_id=sol_id)
    
    # 1. 精确匹配 token
    exact_match = matching_configs[matching_configs["token"] == token]
    if not exact_match.empty:
        sol_id = int(exact_match.iloc[0]["sol_id"])
        print(f">>> Info: Exact token match found for token={token}, using sol_id={sol_id}.")
        return functools.partial(ck_moe, solution_id=sol_id)
    
    # 2. 找最接近的 token
    matching_configs["token_distance"] = abs(matching_configs["token"] - token)
    closest_match = matching_configs.loc[matching_configs["token_distance"].idxmin()]
    
    closest_token = closest_match["token"]
    distance = closest_match["token_distance"]
    sol_id = int(closest_match["sol_id"])
    
    print(f">>> Info: Closest token match found: token={closest_token} (distance={distance}) for target token={token}, using sol_id={sol_id}.")
    return functools.partial(ck_moe, solution_id=sol_id)


def build_moe_index(df):
    """Convert the tuning table into a pure-Python lookup structure."""
    moe_index = {}
    for row in df.itertuples(index=False):
        key = (
            row.arch,
            int(row.inter_dim),
            int(row.model_dim),
            int(row.expert),
            int(row.topk),
            str(row.quant_type),
            int(row.q_size_n),
            int(row.q_size_k),
        )
        entry = moe_index.get(key)
        if entry is None:
            entry = {"token_to_sol": {}, "tokens": []}
            moe_index[key] = entry
        token_val = int(row.token)
        entry["token_to_sol"][token_val] = int(row.sol_id)
        entry["tokens"].append(token_val)

    for entry in moe_index.values():
        entry["tokens"].sort()

    return moe_index


def _find_closest_token(sorted_tokens, target_token):
    idx = bisect_left(sorted_tokens, target_token)
    if idx == 0:
        return sorted_tokens[0]
    if idx == len(sorted_tokens):
        return sorted_tokens[-1]

    before = sorted_tokens[idx - 1]
    after = sorted_tokens[idx]
    if (target_token - before) <= (after - target_token):
        return before
    return after

def get_moe_ck_solution_id(
    arch,
    quant_type,
    token,
    inter_dim,
    model_dim,
    expert,
    topk,
    q_size_n=0,
    q_size_k=0
):
    def get_moe_cfg(ck_tuned_file):
        import pandas as pd
        try:
            moe_cfg = pd.read_csv(ck_tuned_file)
        except Exception as e:
            print(f">>> Warning: Failed to read config file {ck_tuned_file}: {e}")
            return None
        return moe_cfg

    global moe_ck_cfg, current_quant_type
    global moe_ck_noquant_cfg, moe_ck_int8_w8a8_group_cfg
    global moe_ck_noquant_index, moe_ck_int8_w8a8_group_index

    current_index = None

    if moe_ck_cfg is None or quant_type != current_quant_type:
        if quant_type == MoeQuantType.INT8_W8A8:
            if moe_ck_int8_w8a8_group_cfg is None:
                moe_ck_int8_w8a8_group_cfg = get_moe_cfg(ck_tuned_int8_w8a8_group_file)
                if moe_ck_int8_w8a8_group_cfg is not None:
                    moe_ck_int8_w8a8_group_index = build_moe_index(moe_ck_int8_w8a8_group_cfg)
            moe_ck_cfg = moe_ck_int8_w8a8_group_cfg
        elif quant_type == MoeQuantType.NO_QUANT:
            if moe_ck_noquant_cfg is None:
                moe_ck_noquant_cfg = get_moe_cfg(ck_tuned_file)
                if moe_ck_noquant_cfg is not None:
                    moe_ck_noquant_index = build_moe_index(moe_ck_noquant_cfg)
            moe_ck_cfg = moe_ck_noquant_cfg
        else:
            print(f">>> Warning: quant_type {quant_type} not supported for CK lookup, fallback to no-quant table.")
            if moe_ck_noquant_cfg is None:
                moe_ck_noquant_cfg = get_moe_cfg(ck_tuned_file)
                if moe_ck_noquant_cfg is not None:
                    moe_ck_noquant_index = build_moe_index(moe_ck_noquant_cfg)
            moe_ck_cfg = moe_ck_noquant_cfg
            quant_type = MoeQuantType.NO_QUANT

        current_quant_type = quant_type

    if quant_type == MoeQuantType.INT8_W8A8:
        current_index = moe_ck_int8_w8a8_group_index
    else:
        current_index = moe_ck_noquant_index

    if moe_ck_cfg is None:
        print(f">>> Warning: config file is not found, using default ck solution.")
        return 0
    if current_index is None:
        print(f">>> Warning: ck index is not built, using default ck solution.")
        return 0
    
    key = (arch, inter_dim, model_dim, expert, topk, str(quant_type), q_size_n, q_size_k)
    candidates = current_index.get(key)
    

    if not candidates:
        print(f">>> Warning: No matching config pattern found for key={key}, using default ck solution.")
        return 0
    

    # 1. 精确匹配 token
    token = int(token)
    token_to_sol = candidates["token_to_sol"]
    sol_id = token_to_sol.get(token)
    if sol_id is not None:
        return int(sol_id)

    # 2. 找最接近的 token
    closest_token = _find_closest_token(candidates["tokens"], token)
    sol_id = token_to_sol[closest_token]
    return int(sol_id)


def ck_moe_stage_1(
        hidden_states,
        w1,  # [E, inter_dim*2, model_dim]
        w2,  # [E, model_dim, inter_dim]
        sorted_token_ids,  # [max_num_tokens_padded]
        sorted_expert_ids,  # [max_num_m_blocks]
        tokens_positions_per_expert,  # [num_experts*2]
        num_valid_ids,  # [1]
        use_int8_w8a8_block: bool,
        use_fp8_w8a8_block: bool,
        w1_scale,
        a1_scale,
        dtype,
        topk,
        block_shape_n=0,
        block_shape_k=0,
        block_size=16,
        Activation=ActivationType.Silu,
        sorted_weights=None,  # [max_num_tokens_padded]
):
    token_num = hidden_states.shape[0]
    D = w1.shape[1]
    # max_num_tokens_padded = sorted_expert_ids.shape[0]*block_size
    if Activation == ActivationType.Silu:
        act_op = 1
    else:
        act_op = 0

    if w1.dtype is torch.uint32:
        D = D * 8

    gemm_out_type = torch.float16
    # for now, ck_moe_stage_1 has not do the activation inside, so 'D = 2 * inter_dim'
    # out = torch.empty((token_num * topk, D), dtype=gemm_out_type, device=hidden_states.device)
    out = torch.empty((token_num * topk, D//2), dtype=gemm_out_type, device=hidden_states.device)

    aiter.ck_moe_stage_1(
        hidden_states,
        w1,
        w2,
        sorted_token_ids,
        sorted_expert_ids,
        tokens_positions_per_expert,
        num_valid_ids,
        out,
        topk,
        use_int8_w8a8_block,
        use_fp8_w8a8_block,
        w1_scale,
        a1_scale,
        block_shape_n,
        block_shape_k,
        block_size,
        sorted_weights,
        act_op,
    )

    # silu and multiply
    # silu_out = torch.empty((token_num * topk, D // 2), dtype=dtype, device=hidden_states.device)
    # aiter.silu_and_mul(silu_out, out.to(dtype))
    # return silu_out

    return out.to(dtype)

def ck_moe_stage_2(
        hidden_states,
        w1,  # [E, inter_dim*2, model_dim]
        w2,  # [E, model_dim, inter_dim]
        sorted_token_ids,  # [max_num_tokens_padded]
        sorted_expert_ids,  # [max_num_m_blocks]
        tokens_positions_per_expert,  # [num_experts*2]
        num_valid_ids,  # [1]
        use_int8_w8a8_block: bool,
        use_fp8_w8a8_block: bool,
        w2_scale,
        a2_scale,
        dtype,
        topk,
        block_shape_n=0,
        block_shape_k=0,
        block_size=16,
        sorted_weights=None,  # [max_num_tokens_padded]
        moe_buf=None,         # [token_num, model_dim]
):
    hidden_states.reshape(-1, hidden_states.shape[-1])

    if moe_buf is None:
        out = torch.zeros(          # must be zeros, because use atomic add inside
            (hidden_states.shape[0]//topk, w2.shape[1]),   # [token_num, model_dim]
            dtype=dtypes.fp32,      # gpu not support fp16 atomic add
            device=hidden_states.device,
        )
    else:
        out = moe_buf
    
    aiter.ck_moe_stage_2(
        hidden_states,
        w1,
        w2,
        sorted_token_ids,
        sorted_expert_ids,
        tokens_positions_per_expert,
        num_valid_ids,
        out,
        topk,
        use_int8_w8a8_block,
        use_fp8_w8a8_block,
        w2_scale,
        a2_scale,
        block_shape_n,
        block_shape_k,
        block_size,
        sorted_weights,
    )
    return out.to(dtype)


def fused_moe_fake(
        hidden_states: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        use_int8_w8a16: Optional[bool] = False,
        use_int4_w4a16: Optional[bool] = False,
        use_int8_w8a8_block: Optional[bool] = False,
        use_int4_w4a8_block: Optional[bool] = False,
        w1_zp: Optional[torch.Tensor] = None,
        w2_zp: Optional[torch.Tensor] = None,
        w1_scale: Optional[torch.Tensor] = None,
        w2_scale: Optional[torch.Tensor] = None,
        a1_scale: Optional[torch.Tensor] = None,
        a2_scale: Optional[torch.Tensor] = None,
        block_shape_n: Optional[int] = 0,
        block_shape_k: Optional[int] = 0,
        block_m: Optional[int] = 32,
        solution_id: Optional[int] = 0,
        expert_mask: Optional[torch.Tensor] = 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=torch.float32, device=device)
    return moe_buf


@torch_compile_guard(gen_fake=fused_moe_fake)
def ck_fused_experts_2stage_impl(hidden_states: torch.Tensor,
                                w1: torch.Tensor,
                                w2: torch.Tensor,
                                topk_weights: torch.Tensor,
                                topk_ids: torch.Tensor,
                                odtype:torch.dtype, #compute or output type for i8& f8
                                inplace: bool = False,
                                activation: str = "silu",
                                use_fp8_w8a8: bool = False,
                                use_int8_w8a8: bool = False,
                                use_int8_w8a16: bool = False,
                                use_int4_w4a16: bool = False,
                                use_int4_w4a8: 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,
                                solution_id: Optional[int] = None)-> torch.Tensor:

    num_tokens, _ = hidden_states.shape
    E, N, _ = w1.shape
    _, model_dim, inter_dim = w2.shape
    top_k_num = topk_ids.shape[1]
    quant_block_n, quant_block_k = block_shape[0],block_shape[1] if block_shape is not None else (0,0)

    sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, tokens_positions_per_expert, moe_buf = moe_sorting(
        topk_ids, topk_weights, E, model_dim, torch.float32, BLOCK_SIZE_M
    )

    # print(f"########### token_per_expert: {tokens_positions_per_expert}")

    if use_int8_w8a8:
        if per_channel_quant:
            print(">>> ck fused moe int8 w8a8 per channel not supported yet.")
            return None
        else: # block scale

            # quantization input if needed
            if hidden_states.dtype == torch.float16 or hidden_states.dtype==torch.bfloat16:
                input_q, input_scale = per_block_quant_wrapper((1, quant_block_k))(per_token_quant_hip)(hidden_states, quant_dtype=torch.int8)
            else:
                input_q, input_scale = hidden_states, a1_scale

            out_st1 = ck_moe_stage_1(
                input_q,    # 暂时由外部quant input
                w1,
                w2,
                sorted_ids,
                sorted_expert_ids,
                tokens_positions_per_expert,
                num_valid_ids,
                True,
                False,
                w1_scale,
                input_scale,
                odtype,     # fp16/bf16 compute
                top_k_num,
                block_shape_n=quant_block_n,
                block_shape_k=quant_block_k,
                block_size=BLOCK_SIZE_M,
                Activation=ActivationType.Silu if activation=="silu" else ActivationType.Gelu,
                sorted_weights=None)    # stage1不处理topk weights


            # quantization stage1 output
            out_st1 = out_st1.reshape(-1, out_st1.shape[-1])

            bridge_q, bridge_scale = per_block_quant_wrapper((1, quant_block_k))(per_token_quant_hip)(out_st1, quant_dtype=torch.int8)

            out = ck_moe_stage_2(
                bridge_q,
                w1,
                w2,
                sorted_ids,
                sorted_expert_ids,
                tokens_positions_per_expert,
                num_valid_ids,
                True,
                False,
                w2_scale,
                bridge_scale,
                odtype,     # fp16/bf16 compute
                top_k_num,
                block_shape_n=quant_block_n,
                block_shape_k=quant_block_k,
                block_size=BLOCK_SIZE_M,
                sorted_weights=sorted_weights,  # stage2处理topk weights
                moe_buf=moe_buf
            )

            # return (out, out_st1)
            return out

    elif use_fp8_w8a8:
        if per_channel_quant:
            print(">>> ck fused moe fp8 w8a8 per channel not supported yet.")
            return None
        else:
            # quantization input if needed
            if hidden_states.dtype == torch.float16 or hidden_states.dtype==torch.bfloat16:
                input_q, input_scale = per_block_quant_wrapper((1, quant_block_k))(per_token_quant_hip)(hidden_states, quant_dtype=torch.float8_e4m3fn)
            else:
                input_q, input_scale = hidden_states, a1_scale

            out_st1 = ck_moe_stage_1(
                input_q,    # 暂时由外部quant input
                w1,
                w2,
                sorted_ids,
                sorted_expert_ids,
                tokens_positions_per_expert,
                num_valid_ids,
                False,
                True,
                w1_scale,
                input_scale,
                odtype,     # fp16/bf16 compute
                top_k_num,
                block_shape_n=quant_block_n,
                block_shape_k=quant_block_k,
                block_size=BLOCK_SIZE_M,
                Activation=ActivationType.Silu if activation=="silu" else ActivationType.Gelu,
                sorted_weights=None)    # stage1不处理topk weights


            # quantization stage1 output
            out_st1 = out_st1.reshape(-1, out_st1.shape[-1])

            bridge_q, bridge_scale = per_block_quant_wrapper((1, quant_block_k))(per_token_quant_hip)(out_st1, quant_dtype=torch.float8_e4m3fn)

            out = ck_moe_stage_2(
                bridge_q,
                w1,
                w2,
                sorted_ids,
                sorted_expert_ids,
                tokens_positions_per_expert,
                num_valid_ids,
                False,
                True,
                w2_scale,
                bridge_scale,
                odtype,     # fp16/bf16 compute
                top_k_num,
                block_shape_n=quant_block_n,
                block_shape_k=quant_block_k,
                block_size=BLOCK_SIZE_M,
                sorted_weights=sorted_weights,  # stage2处理topk weights
                moe_buf=moe_buf
            )

            # return (out, out_st1)
            return out

    else:
        return None

@torch_compile_guard(gen_fake=fused_moe_fake)
def ck_fused_experts_1stage_impl(
        hidden_states: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        odtype:torch.dtype, #compute or output type for i8& f8
        use_int8_w8a16: Optional[bool] = False,
        use_int4_w4a16: Optional[bool] = False,
        use_int8_w8a8_block: Optional[bool] = False,
        use_int4_w4a8_block: Optional[bool] = False,
        w1_zp: Optional[torch.Tensor] = None,
        w2_zp: Optional[torch.Tensor] = None,
        w1_scale: Optional[torch.Tensor] = None,
        w2_scale: Optional[torch.Tensor] = None,
        a1_scale: Optional[torch.Tensor] = None,
        a2_scale: Optional[torch.Tensor] = None,
        block_shape_n: Optional[int] = 0,
        block_shape_k: Optional[int] = 0,
        block_m: Optional[int] = 32,
        use_shuffle: Optional[bool] = False,
        solution_id: Optional[int] = 0,
        expert_mask: Optional[torch.Tensor] = None)-> torch.Tensor:

    if use_shuffle and use_shuffle==True:
        out = ck_shuffle_moe(
                hidden_states,
                w1,
                w2,
                topk_weights,
                topk_ids,
                use_int8_w8a16,
                use_int4_w4a16,
                use_int8_w8a8_block,
                use_int4_w4a8_block,
                w1_zp,
                w2_zp,
                w1_scale,
                w2_scale,
                a1_scale,
                a2_scale,
                block_shape_n,
                block_shape_k,
                block_m,
                solution_id,
                expert_mask)
        return out.to(odtype)
    else:
        out = ck_moe(
                hidden_states,
                w1,
                w2,
                topk_weights,
                topk_ids,
                use_int8_w8a16,
                use_int4_w4a16,
                use_int8_w8a8_block,
                use_int4_w4a8_block,
                w1_zp,
                w2_zp,
                w1_scale,
                w2_scale,
                a1_scale,
                a2_scale,
                block_shape_n,
                block_shape_k,
                block_m,
                solution_id,
                expert_mask)
        return out.to(odtype)

        # sum_out = torch.empty_like(hidden_states, dtype=out.dtype, device=out.device)
        # moe_sum(out, sum_out)
        # return sum_out.to(odtype)

def bits30_31(solution_id: int) -> int:
        unsigned32 = solution_id & 0xFFFFFFFF  # treat as 32-bit two’s complement
        return (unsigned32 & 0xC0000000) >> 30  # 0xC0000000 = bits 31–30 set

# The outside interface
def run_fused_experts_ck_impl(
        hidden_states: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        odtype:torch.dtype, #compute or output type for i8& f8
        inplace: bool = False,
        activation: str = "silu",
        use_fp8_w8a8: bool = False,
        use_int8_w8a8: bool = False,
        use_int8_w8a16: bool = False,
        use_int4_w4a16: bool = False,
        use_int4_w4a8: bool = False,
        per_channel_quant: bool = False,
        global_num_experts: int = -1,
        block_m: int = BLOCK_SIZE_M,
        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_shuffle: Optional[bool] = False,
        routed_scaling_factor: Optional[float] = 1.0,
        solution_id: Optional[int] = None)-> torch.Tensor:

    if solution_id == None:
        if use_shuffle and use_shuffle==True:   #only one stage supports shuffle for now.
            solution_id = 0
        else:
            # solution_id = 0
            arch = get_gfx()
            quantType = MoeQuantType.NO_QUANT
            if use_int8_w8a8:
                quantType = MoeQuantType.INT8_W8A8

            E, model_dim, inter_dim = w2.shape
            topk = topk_ids.shape[1]

            if quantType == MoeQuantType.INT8_W8A8 and block_shape[1] == 64:
                solution_id = 1 << 30   # only two stage supports block_shape_k = 64

            else:
                solution_id = get_moe_ck_solution_id(
                    arch,
                    quantType,
                    hidden_states.shape[0],
                    inter_dim,  # inter_dim
                    model_dim,
                    E,
                    topk,
                    block_shape[0] if block_shape is not None else 0,
                    block_shape[1] if block_shape is not None else 0
                )


    solutionType = bits30_31(solution_id)

    # two stage
    if solutionType == 1:
        return ck_fused_experts_2stage_impl(
                hidden_states,
                w1,
                w2,
                topk_weights,
                topk_ids,
                odtype,
                inplace,
                activation,
                use_fp8_w8a8,
                use_int8_w8a8,
                use_int8_w8a16,
                use_int4_w4a16,
                use_int4_w4a8,
                per_channel_quant,
                global_num_experts,
                expert_map,
                w1_scale,
                w2_scale,
                w1_zp,
                w2_zp,
                a1_scale,
                a2_scale,
                block_shape,
                solution_id)

    # one stage
    else:
        return ck_fused_experts_1stage_impl(
                hidden_states,
                w1,
                w2,
                topk_weights,
                topk_ids,
                odtype,
                use_int8_w8a16,
                use_int4_w4a16,
                use_int8_w8a8,
                use_int4_w4a8,
                w1_zp,
                w2_zp,
                w1_scale,
                w2_scale,
                a1_scale,
                a2_scale,
                block_shape[0] if block_shape is not None else 0,
                block_shape[1] if block_shape is not None else 0,
                block_m,
                use_shuffle,
                solution_id,
                expert_map)