// SPDX-License-Identifier: MIT
 
#include <hip/hip_runtime.h>
#include <hip/hip_fp16.h>
#include <torch/all.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include "aiter_hip_common.h"
#include "moe_op.h"
//#include "py_itfs_common.h"
#include <iostream>
#include <map>
#include <stdexcept>

class Moe1SolutionLookup {
public:
    Moe1SolutionLookup() {
        SolutionMap_DeepSeek = {
            {1, 10000}, {2, 10000}, {4, 10000}, {8, 10000}, {12, 10000},{16, 10000}, {24, 10000},
            {32, 10000}, {48, 10000}, {64, 10000},{96, 10000}, {128, 10000}, {160, 10000}, {256, 10000}, {512, 10000}
        };
        SolutionMap_QianWen = {
            {1, 30000}, {2, 30000}, {4, 30000}, {8, 30000}, {12, 30000},{16, 30000}, {24, 30000},
            {32, 30000}, {48, 30000}, {64, 30000},{96, 30000}, {128, 30000}, {160, 30000}, {256, 30000}, {512, 30000}
        };
    }

    uint32_t getSolution(int m, int k) {
        if (m < 1) throw std::out_of_range("m must be >= 1");
        if (m > 32768) m = 32768;

        if (k == 256) {
          auto it = SolutionMap_DeepSeek.lower_bound(m); // Find the first key >= m
          // if it > m or end(), return the previous one
          if (it == SolutionMap_DeepSeek.end() || it->first > m) {
            if (it == SolutionMap_DeepSeek.begin()) throw std::logic_error("No valid solution found");
              --it;
          }
          return it->second;
        }
        else {
          auto it = SolutionMap_QianWen.lower_bound(m); // Find the first key >= m
          // if it > m or end(), return the previous one
          if (it == SolutionMap_QianWen.end() || it->first > m) {
            if (it == SolutionMap_QianWen.begin()) throw std::logic_error("No valid solution found");
              --it;
          }
          return it->second;
        }
    }
private:
    std::map<int, uint32_t> SolutionMap_DeepSeek;
    std::map<int, uint32_t> SolutionMap_QianWen;
};

class Moe2SolutionLookup {
public:
    Moe2SolutionLookup() {
        SolutionMap_DeepSeek = {
            {1, 20000}, {2, 20000}, {3, 20001}, {4, 20001}, {8, 20001}, {12, 20001},{16, 20001}, {24, 20001},
            {32, 20001}, {48, 20001}, {64, 20001},{96, 20001}, {128, 20001}, {160, 20001}, {256, 20001}, {384, 20002}, {512, 20002}
        };
        SolutionMap_QianWen = {
            {1, 40000}, {2, 40000}, {3, 40000}, {4, 40000}, {8, 40000}, {12, 40000},{16, 40000}, {24, 40000},
            {32, 40000}, {48, 40000}, {64, 40000},{96, 40000}, {128, 40000}, {160, 40000}, {256, 40000}, {384, 40000}, {512, 40000}
        };
    }

    uint32_t getSolution(int m, int k) {
        if (m < 1) throw std::out_of_range("m must be >= 1");
        if (m > 32768) m = 32768;

        if (k == 256) {
          auto it = SolutionMap_DeepSeek.lower_bound(m); // Find the first key >= m
          // if it > m or end(), return the previous one
          if (it == SolutionMap_DeepSeek.end() || it->first > m) {
            if (it == SolutionMap_DeepSeek.begin()) throw std::logic_error("No valid solution found");
              --it;
          }
          return it->second;
        }
        else {
          auto it = SolutionMap_QianWen.lower_bound(m); // Find the first key >= m
          // if it > m or end(), return the previous one
          if (it == SolutionMap_QianWen.end() || it->first > m) {
            if (it == SolutionMap_QianWen.begin()) throw std::logic_error("No valid solution found");
              --it;
          }
          return it->second;
        }
    }
private:
    std::map<int, uint32_t> SolutionMap_DeepSeek;
    std::map<int, uint32_t> SolutionMap_QianWen;
};


//#define DEBUG_BUFFER

struct __attribute__((packed)) KernelArgs
{
     uint32_t gemm_count;
     uint32_t internalArgs;
     uint32_t internalArgs1;
     uint32_t numWorkGroups;
     void*    debugBuffer;
     void*    argsPtr;
};


struct __attribute__((packed)) GroupedGemmArgs
{
     uint32_t m; //!< size m
     uint32_t n; //!< size n
     uint32_t batch; //!< size batch
     uint32_t k; //!< size k
     void*    d; //!< The d matrix input pointer.
     void*    c; //!< The c matrix input pointer.
     void*    a; //!< The a matrix input pointer.
     void*    b; //!< The b matrix input pointer.
     uint32_t strideD1; //!< The d leading dimension.
     uint32_t strideD2; //!< The d batch stride
     uint32_t strideC1; //!< The c leading dimension.
     uint32_t strideC2; //!< The c batch stride
     uint32_t strideA1; //!< The a leading dimension.
     uint32_t strideA2; //!< The a batch stride
     uint32_t strideB1; //!< The b leading dimension.
     uint32_t strideB2; //!< The b batch stride
     float    alpha; //!< The alpha value.
     float    beta; //!< The beta value.
     void*    sorted_token_ids;
     void*    sorted_weights;
     void*    sorted_expert_ids;
     uint32_t num_valid_ids;
     uint32_t top_k;
     void*    scale_a;
     void*    scale_b;
     void*    zero_points;
};


struct __attribute__((packed)) HipFunctionArgs {
  uint32_t gemm_count;
  uint32_t internalArgs;
  uint32_t internalArgs1;
  uint32_t numWorkGroups;
  uint32_t m; //!< size m
  uint32_t n; //!< size n
  uint32_t batch; //!< size n
  uint32_t k; //!< size k
  void*    d; //!< The d matrix input pointer.
  void*    c; //!< The d matrix input pointer.
  void*    a; //!< The a matrix input pointer.
  void*    b; //!< The b matrix input pointer.
  uint32_t strideD1; //!< The d leading dimension.
  uint32_t strideD2; //!< The d batch stride
  uint32_t strideC1; //!< The c leading dimension.
  uint32_t strideC2; //!< The c batch stride
  uint32_t strideA1; //!< The a leading dimension.
  uint32_t strideA2; //!< The a batch stride
  uint32_t strideB1; //!< The b leading dimension.
  uint32_t strideB2; //!< The b batch stride
  float    alpha; //!< The alpha value.
  float    beta; //!< The beta value.
  void*    sorted_token_ids;
  void*    sorted_weights;
  void*    sorted_expert_ids;
  uint32_t num_valid_ids;
  uint32_t top_k;
  void*    scale_a;
  void*    scale_b;
  void*    zero_points;
  void*    experts_num;              // num of experts
  uint32_t persist_groups;
  void*    debugBuffer;
  float    topk_rcip;
};

void printFunctionArgs(const HipFunctionArgs& args) {
    std::cout << "gemm_count: " << args.gemm_count << "\n";
    std::cout << "internalArgs: " << args.internalArgs << "\n";
    std::cout << "internalArgs1: " << args.internalArgs1 << "\n";
    std::cout << "numWorkGroups: " << args.numWorkGroups << "\n";

    std::cout << "m: " << args.m << "\n";
    std::cout << "n: " << args.n << "\n";
    std::cout << "batch: " << args.batch << "\n";
    std::cout << "k: " << args.k << "\n";

    std::cout << "d: " << static_cast<void*>(args.d) << "\n";
    std::cout << "c: " << static_cast<void*>(args.d) << "\n";
    std::cout << "a: " << static_cast<const void*>(args.a) << "\n";
    std::cout << "b: " << static_cast<const void*>(args.b) << "\n";

    std::cout << "strideD1: " << args.strideD1 << "\n";
    std::cout << "strideD2: " << args.strideD2 << "\n";
    std::cout << "strideC1: " << args.strideC1 << "\n";
    std::cout << "strideC2: " << args.strideC2 << "\n";

    std::cout << "strideA1: " << args.strideA1 << "\n";
    std::cout << "strideA2: " << args.strideA2 << "\n";
    std::cout << "strideB1: " << args.strideB1 << "\n";
    std::cout << "strideB2: " << args.strideB2 << "\n";

    std::cout << "alpha: " << args.alpha << "\n";
    std::cout << "beta: " << args.beta << "\n";

    std::cout << "sorted_token_ids: " << static_cast<const void*>(args.sorted_token_ids) << "\n";
    std::cout << "sorted_weights: " << static_cast<const void*>(args.sorted_weights) << "\n";
    std::cout << "sorted_expert_ids: " << static_cast<const void*>(args.sorted_expert_ids) << "\n";
    std::cout << "num_valid_ids: " << args.num_valid_ids << "\n";
    std::cout << "top_k: " << args.top_k << "\n";
    std::cout << "scale_a: " << static_cast<void*>(args.scale_a) << "\n";
    std::cout << "scale_b: " << static_cast<void*>(args.scale_b) << "\n";
    std::cout << "zero_points: " << static_cast<void*>(args.zero_points) << "\n";

    std::cout << "experts_num: " << args.experts_num << "\n";
    std::cout << "persist_groups: " << args.persist_groups << "\n";
    std::cout << "topk_rcip: " << args.topk_rcip << "\n";
}


class FMoeKernel
{
private:
    hipModule_t module;
    hipFunction_t kernel_func;

public:
    FMoeKernel(const char *name, const char *hsaco)
    {
        const char *AITER_ASM_DIR = std::getenv("AITER_ASM_DIR");
        std::cout << "[aiter] hipModuleLoad: " << (std::string(AITER_ASM_DIR) + hsaco).c_str() << " GetFunction: " << name;
        HIP_CALL(hipModuleLoad(&module, (std::string(AITER_ASM_DIR) + hsaco).c_str()));
        HIP_CALL(hipModuleGetFunction(&kernel_func, module, name));
        std::cout << " Success" << std::endl;
    };

    size_t debugBufferElementsPerThread = 16;
    size_t debugBufferSize = 0;
    std::shared_ptr<unsigned int> debugBufferHostPtr;
    unsigned int* debugBufferDevicePtr   = nullptr;

    void CreateDebugBuffer(size_t numWorkGroups, size_t numThreads)
    {
        std::cout << "Smart Lt debugKernel is enabled !!!! " << std::endl;

        size_t debugBufferNumElem = debugBufferElementsPerThread;
        debugBufferNumElem *= numWorkGroups;
        debugBufferNumElem *= numThreads;
        debugBufferSize = debugBufferNumElem * 4;

        hipMalloc(&debugBufferDevicePtr, debugBufferSize);

        debugBufferHostPtr = std::shared_ptr<unsigned int >(
                        (unsigned int *)std::malloc(debugBufferSize),
                        std::free);
        memset(debugBufferHostPtr.get(), 0, debugBufferSize);
        hipMemcpy(debugBufferDevicePtr, debugBufferHostPtr.get(), debugBufferSize, hipMemcpyHostToDevice);
    };

    void debug_buffer_print()
    {
        hipMemcpy(debugBufferHostPtr.get(), debugBufferDevicePtr, debugBufferSize, hipMemcpyDeviceToHost);

        unsigned int * dbg_ptr = debugBufferHostPtr.get();
        const char *field_names[16] = {
        "tid","wg0","wg1","groA","groB",
        "lraA","lraB","lwaA","lwaB"};


        for (unsigned int i = 0; i < debugBufferSize / 4 / debugBufferElementsPerThread; i++) {
            if (i % 64 == 0) {
                printf("\n");
                for (unsigned int j = 0; j < debugBufferElementsPerThread; j++) {
                    printf("%12s,", field_names[j]);
                }
                printf("\n");
            }

            char flags[16] = {'u', 'u', 'u',
                                                    'x', 'x', 'x', 'x', 'x', 'x',
                                                    'x', 'x', 'x', 'x', 'x', 'x', 'x'};
            //if((i%64) < 4 || (i%64) >= 60)
            //if(i<512)
            {
            for (unsigned int j = 0; j < debugBufferElementsPerThread; j++) {
                if (flags[j] == 'u')
                    printf("    %8u,", dbg_ptr[i * debugBufferElementsPerThread + j]);
                else if (flags[j] == 'x')
                    printf("  0x%08x,", dbg_ptr[i * debugBufferElementsPerThread + j]);
                else if (flags[j] == 'f')
                    printf("    %8.4f,", ((float *)dbg_ptr)[i * debugBufferElementsPerThread + j]);
                else if (flags[j] == 'd')
                    printf("    %8d,", dbg_ptr[i * debugBufferElementsPerThread + j]);
            }

            printf("\n");
        }
        }
        printf("\n");

    };

    template <typename T, typename T_O, bool firstStage = false>
    void launch_kernel(const std::vector<uint32_t>& Config,
                       torch::Tensor &out,               // [token_cnt, dim]
                       torch::Tensor &input,             // [token_cnt, dim] M,K
                       torch::Tensor &w1,                // [expert, inter_dim, dim] N,K
                       torch::Tensor &w2,                // [expert, dim, inter_dim]
                       torch::Tensor &sorted_token_ids,  // [max_num_tokens_padded]
                       torch::Tensor &sorted_weights,    // [max_num_tokens_padded]
                       torch::Tensor &sorted_expert_ids, // [max_num_m_blocks]
                       torch::Tensor &num_valid_ids,                    //
                       uint32_t top_k,
                       std::optional<torch::Tensor> scale_a = std::nullopt,
                       std::optional<torch::Tensor> scale_b = std::nullopt,
                       std::optional<torch::Tensor> zero_points = std::nullopt,
                       std::optional<int> mode = 0,
                       std::optional<int> block_size = 16,
                       std::optional<int> persist_groups = 0)
    {

        int size_m, size_n, size_k;
        uint32_t chosen_experts;
        int64_t block_m;
        uint32_t PersistGroups=0;

        if(block_size.has_value())
        {
          block_m = block_size.value();
        }
        if (persist_groups.has_value())
        {
          PersistGroups = persist_groups.value();
        }

        if constexpr (firstStage) {
          chosen_experts = std::min(input.size(0)*8, sorted_token_ids.size(0)/block_m);
        }
        else {
          chosen_experts = std::min(out.size(0)*8, sorted_token_ids.size(0)/block_m);
        }
        size_m = Config[0];
        if constexpr (firstStage) {
          size_n = w1.size(1);
          size_k = w1.size(2);
        }
        else {
          size_n = w2.size(1);
          size_k = w2.size(2);
        }

        HipFunctionArgs hipFunctionArgs;
        hipFunctionArgs.gemm_count = (1 & 0x3FFFFFFF) | (1 << 30);
        hipFunctionArgs.internalArgs = 0x00200001;
        hipFunctionArgs.internalArgs1 = 1;

        if (PersistGroups != 0)
        {
          hipFunctionArgs.numWorkGroups = PersistGroups;
        }
        else
        {
          hipFunctionArgs.numWorkGroups = chosen_experts * ((size_n + Config[1] - 1) / Config[1]);
        }

        hipFunctionArgs.m = size_m;
        hipFunctionArgs.n = size_n;
        hipFunctionArgs.batch = 1;
        if (mode == 2 || mode == 3 || mode == 4) {
          hipFunctionArgs.k = size_k * 2;
        }
        else {
          hipFunctionArgs.k = size_k;
        }
        hipFunctionArgs.d = out.data_ptr();
        hipFunctionArgs.c = out.data_ptr();
        hipFunctionArgs.a = input.data_ptr();
        if constexpr (firstStage) {
          hipFunctionArgs.b = w1.data_ptr();
        }
        else {
          hipFunctionArgs.b = w2.data_ptr();
        }

        hipFunctionArgs.strideD1 = size_n;
        hipFunctionArgs.strideD2 = size_m * size_n;
        hipFunctionArgs.strideC1 = size_n;
        hipFunctionArgs.strideC2 = size_m * size_n;
        if (mode == 2 || mode == 3 || mode == 4) {
            hipFunctionArgs.strideA1 = size_k * 2;
            hipFunctionArgs.strideA2 = size_m * size_k * 2;
            hipFunctionArgs.strideB1 = size_k * 2;
            hipFunctionArgs.strideB2 = size_n * size_k * 2;
        }
        else {
            hipFunctionArgs.strideA1 = size_k;
            hipFunctionArgs.strideA2 = size_m * size_k;
            hipFunctionArgs.strideB1 = size_k;
            hipFunctionArgs.strideB2 = size_n * size_k;
        }

        hipFunctionArgs.alpha = 1;
        hipFunctionArgs.beta = 0;

        hipFunctionArgs.sorted_token_ids = sorted_token_ids.data_ptr();
        hipFunctionArgs.sorted_weights = sorted_weights.data_ptr();
        hipFunctionArgs.sorted_expert_ids = sorted_expert_ids.data_ptr();

        if constexpr (firstStage) {
          hipFunctionArgs.num_valid_ids = input.size(0);
        }
        else {
          hipFunctionArgs.num_valid_ids = out.size(0);
        }

        hipFunctionArgs.top_k = top_k; 

        if(scale_a.has_value() && scale_a.value().has_storage()){
          hipFunctionArgs.scale_a = scale_a.value().data_ptr();
        }
        else{
          hipFunctionArgs.scale_a = input.data_ptr();
        }

        if(scale_b.has_value()&& scale_b.value().has_storage()){
          hipFunctionArgs.scale_b = scale_b.value().data_ptr();
        }
        else {
          hipFunctionArgs.scale_b = input.data_ptr();
        }

        if(zero_points.has_value()&& zero_points.value().has_storage()){
          hipFunctionArgs.zero_points = zero_points.value().data_ptr();
        }
        else{
          hipFunctionArgs.zero_points = input.data_ptr();
        }

        hipFunctionArgs.experts_num = num_valid_ids.data_ptr();
        hipFunctionArgs.persist_groups = PersistGroups;

        hipFunctionArgs.topk_rcip = 1 / float(top_k);
        hipFunctionArgs.debugBuffer = nullptr;

#if 0
        printFunctionArgs(hipFunctionArgs);

#endif

        size_t arg_size = sizeof(hipFunctionArgs);
        void *config[] = {HIP_LAUNCH_PARAM_BUFFER_POINTER,
                          &hipFunctionArgs, HIP_LAUNCH_PARAM_BUFFER_SIZE,
                          &arg_size, HIP_LAUNCH_PARAM_END};

        int bdx = Config[2];

        int gdx = hipFunctionArgs.numWorkGroups;
        int gdy = 1;
        int gdz = 1;

#ifdef DEBUG_BUFFER
        CreateDebugBuffer(gdx, bdx);
        hipFunctionArgs.debugBuffer = debugBufferDevicePtr;
#endif

//        const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
        const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
        HIP_CALL(hipModuleLaunchKernel(kernel_func,
                                       gdx, gdy, gdz,
                                       bdx, 1, 1,
                                       0, stream, nullptr, (void **)&config));

#ifdef DEBUG_BUFFER
        debug_buffer_print();
#endif

    };
};


void asm_fmoe_stage1(torch::Tensor &out,               // [token_cnt, dim]
              torch::Tensor &input,             // [token_cnt, dim] M,K
              torch::Tensor &gate,              // [expert, inter_dim, dim] N,K
              torch::Tensor &down,              // [expert, dim, inter_dim]
              torch::Tensor &sorted_token_ids,  // [max_num_tokens_padded]
              torch::Tensor &sorted_weights,    // [max_num_tokens_padded]
              torch::Tensor &sorted_expert_ids, // [max_num_m_blocks]
              torch::Tensor &num_valid_ids,                     //
              uint32_t top_k,
              std::optional<torch::Tensor> scale_a = std::nullopt,
              std::optional<torch::Tensor> scale_b = std::nullopt,
              std::optional<torch::Tensor> zero_points = std::nullopt,
              std::optional<int> mode = 0,
              std::optional<int> solidx = 0,
              std::optional<int> block_size = 16,
              std::optional<int> persist_groups = 0
)
{

    struct FMoeKernelConfig
    {
        std::string name;
        std::string co_name;
        uint32_t MT0;
        uint32_t MT1;
        uint32_t bdx;
    };

    static std::unordered_map<int, FMoeKernelConfig> moe1_kernel_w4a16_fp16_configs = {
        {10000, {"Cijk_Alik_Bljk_HHS_BH_UserArgs_MT16x128x128_SN_K1_PGR4_TT1_8_WG16_16_3", "w4a16/fp16/Cijk_Alik_Bljk_HHS_BH_UserArgs_MT16x128x128_SN_K1_PGR4_TT1_8_batch_w4a16_gate1_bs128.co", 16, 128, 768}},
        {11000, {"Cijk_Alik_Bljk_HHS_BH_UserArgs_MT32x128x128_SN_K1_PGR4_TT2_8_WG16_16_3", "w4a16/fp16/Cijk_Alik_Bljk_HHS_BH_UserArgs_MT32x128x128_SN_K1_PGR4_TT2_8_batch_w4a16_gate1_bs128.co", 32, 128, 768}},
        {11001, {"Cijk_Alik_Bljk_HHS_BH_UserArgs_MT32x128x128_SN_K1_PGR4_TT2_8_WG16_16_3", "w4a16/fp16/Cijk_Alik_Bljk_HHS_BH_UserArgs_MT32x128x128_SN_K1_PGR4_TT2_8_batch_w4a16_gate1_bs256.co", 32, 128, 768}},
        {11002, {"Cijk_Alik_Bljk_HHS_BH_UserArgs_MT32x128x128_SN_K1_PGR4_TT2_8_WG16_16_3", "w4a16/fp16/Cijk_Alik_Bljk_HHS_BH_UserArgs_MT32x128x128_SN_K1_PGR4_TT2_8_batch_w4a16.co", 32, 128, 768}},
        {30000, {"Cijk_Alik_Bljk_HHS_BH_UserArgs_MT16x128x128_SN_K1_PGR4_TT1_8_WG16_16_3", "w4a16/fp16/Cijk_Alik_Bljk_HHS_BH_UserArgs_MT16x128x128_SN_K1_PGR4_TT1_8_batch_w4a16_gate1_qianwen.co", 16, 128, 768}}};

    static std::unordered_map<int, FMoeKernelConfig> moe1_kernel_w4a16_bf16_configs = {
        {10000, {"Cijk_Alik_Bljk_BBS_BH_UserArgs_MT16x128x128_SN_K1_PGR4_TT1_8_WG16_16_3", "w4a16/bf16/Cijk_Alik_Bljk_BBS_BH_UserArgs_MT16x128x128_SN_K1_PGR4_TT1_8_batch_w4a16_gate1_bs128.co", 16, 128, 768}},
        {11000, {"Cijk_Alik_Bljk_BBS_BH_UserArgs_MT32x128x128_SN_K1_PGR4_TT2_8_WG16_16_3", "w4a16/bf16/Cijk_Alik_Bljk_BBS_BH_UserArgs_MT32x128x128_SN_K1_PGR4_TT2_8_batch_w4a16_gate1_bs128.co", 32, 128, 768}},
        {11001, {"Cijk_Alik_Bljk_BBS_BH_UserArgs_MT32x128x128_SN_K1_PGR4_TT2_8_WG16_16_3", "w4a16/bf16/Cijk_Alik_Bljk_BBS_BH_UserArgs_MT32x128x128_SN_K1_PGR4_TT2_8_batch_w4a16_gate1_bs256.co", 32, 128, 768}},
        {11002, {"Cijk_Alik_Bljk_BBS_BH_UserArgs_MT32x128x128_SN_K1_PGR4_TT2_8_WG16_16_3", "w4a16/bf16/Cijk_Alik_Bljk_BBS_BH_UserArgs_MT32x128x128_SN_K1_PGR4_TT2_8_batch_w4a16.co", 32, 128, 768}},
        {30000, {"Cijk_Alik_Bljk_BBS_BH_UserArgs_MT16x128x128_SN_K1_PGR4_TT1_8_WG16_16_3", "w4a16/bf16/Cijk_Alik_Bljk_BBS_BH_UserArgs_MT16x128x128_SN_K1_PGR4_TT1_8_batch_w4a16_gate1_qianwen.co", 16, 128, 768}}};

    static std::unordered_map<int, FMoeKernelConfig> moe1_kernel_w16a16_configs = {
        {10000, {"Cijk_Alik_Bljk_HHS_BH_UserArgs_MT32x32x128_SN_K1_PGR4_TT2_2_WG16_16_2", "w16a16/Cijk_Alik_Bljk_HHS_BH_UserArgs_MT32x32x128_SN_K1_PGR4_TT2_2_vllm.co", 32, 32, 768}}};
    
    static std::unordered_map<int, FMoeKernelConfig> moe1_kernel_w4a16_bf16_group32_configs = {
        {50032, {"Cijk_Alik_Bljk_BBS_BH_UserArgs_MT32x128x128_SN_K1_PGR4_TT2_8_WG16_16_3", "w4a16/bf16/Cijk_Alik_Bljk_BBS_BH_UserArgs_MT32x128x128_SN_K1_PGR4_TT2_8_batch_w4a16_gate1_bs256_gw32.co", 32, 128, 768}}};
        
    torch::Tensor ScaleA;
    torch::Tensor ScaleB;
    torch::Tensor ZeroPoints;
    int Mode = 0;

    if(scale_a.has_value())
    {
      ScaleA = scale_a.value();
    }
    if (scale_b.has_value())
    {
      ScaleB = scale_b.value();
    }
    if (zero_points.has_value())
    {
      ZeroPoints = zero_points.value();
    }
    if (mode.has_value())
    {
      Mode = mode.value();
    }

    // g1u0
    FMoeKernel *impl_ptr = nullptr;

    static std::unordered_map<std::string, std::unique_ptr<FMoeKernel>> impl_ptr_map;
    std::vector<uint32_t> Config = {16, 128, 768};

    std::unordered_map<int, FMoeKernelConfig> *config_map = nullptr;

    if (Mode == 2) {
      config_map = &moe1_kernel_w4a16_fp16_configs;
    }
    else if (Mode == 3) {
      config_map = &moe1_kernel_w4a16_bf16_configs;
    } 
    else if (Mode == 4) {
      config_map = &moe1_kernel_w4a16_bf16_group32_configs;
    }
    else {
      config_map = &moe1_kernel_w16a16_configs;
    }

    int Solution = 10000;
    if (solidx.has_value())
    {
      Solution = solidx.value();
    }

    if (solidx.value() == 0)
    {
      Moe1SolutionLookup lookup;
      Solution = lookup.getSolution(input.size(0), gate.size(1)/2);
    }

    if (!config_map)
    {
        TORCH_CHECK(false, __func__, " Input only supput Int8!");
    }

    auto it = config_map->find(Solution);
    if (it != config_map->end())
    {
        const auto &config = it->second;
        const char *name = config.name.c_str();
        const char *co_name = config.co_name.c_str();
        Config = {config.MT0, config.MT1, config.bdx};

        auto result = impl_ptr_map.emplace(name, nullptr);
        if (result.second)
        {
            result.first->second = std::make_unique<FMoeKernel>(name, co_name);
        }
        impl_ptr = result.first->second.get();
    }


    TORCH_CHECK(impl_ptr != nullptr,
                __func__, ": unsupport current input type:", input.scalar_type());
    impl_ptr->launch_kernel<uint16_t, uint16_t, true>(Config,
                                                    out,
                                                    input,
                                                    gate,
                                                    down,
                                                    sorted_token_ids,
                                                    sorted_weights,
                                                    sorted_expert_ids,
                                                    num_valid_ids,
                                                    top_k,
                                                    ScaleA,
                                                    ScaleB,
                                                    ZeroPoints,
                                                    Mode,
                                                    block_size,
                                                    persist_groups);
}


void asm_fmoe_stage2(torch::Tensor &out,               // [token_cnt, dim]
              torch::Tensor &input,             // [token_cnt, dim] M,K
              torch::Tensor &gate,              // [expert, inter_dim, dim] N,K
              torch::Tensor &down,              // [expert, dim, inter_dim]
              torch::Tensor &sorted_token_ids,  // [max_num_tokens_padded]
              torch::Tensor &sorted_weights,    // [max_num_tokens_padded]
              torch::Tensor &sorted_expert_ids, // [max_num_m_blocks]
              torch::Tensor &num_valid_ids,                     //
              uint32_t top_k,
              std::optional<torch::Tensor> scale_a = std::nullopt,
              std::optional<torch::Tensor> scale_b = std::nullopt,
              std::optional<torch::Tensor> zero_points = std::nullopt,
              std::optional<int> mode = 0,
              std::optional<int> solidx = 0,
              std::optional<int> block_size = 16,
              std::optional<int> persist_groups = 0
)
{

    struct FMoeKernelConfig
    {
        std::string name;
        std::string co_name;
        uint32_t MT0;
        uint32_t MT1;
        uint32_t bdx;
    };

    static std::unordered_map<int, FMoeKernelConfig> moe2_kernel_w4a16_fp16_configs = {
        {20000, {"Cijk_Alik_Bljk_HHS_BH_UserArgs_MT16x1024x16_SN_K1_TT1_64_WG16_16_3_WGM1", "w4a16/fp16/Cijk_Alik_Bljk_HHS_BH_UserArgs_MT16x1024x16_SN_K1_TT1_64_batch_w4a16_gate2.co", 16, 1024, 768}},
        {20001, {"Cijk_Alik_Bljk_HHS_BH_UserArgs_MT16x3584x16_SN_K1_TT1_224_WG16_16_3_WGM1", "w4a16/fp16/Cijk_Alik_Bljk_HHS_BH_UserArgs_MT16x3584x16_SN_K1_TT1_224_batch_w4a16_gate2.co", 16, 3584, 768}},
        {20002, {"Cijk_Alik_Bljk_HHS_BH_UserArgs_MT16x7168x16_SN_K1_TT1_448_WG16_16_3_WGM1", "w4a16/fp16/Cijk_Alik_Bljk_HHS_BH_UserArgs_MT16x7168x16_SN_K1_TT1_448_batch_w4a16_gate2.co", 16, 7168, 768}},
        {21000, {"Cijk_Alik_Bljk_HHS_BH_UserArgs_MT32x7168x16_SN_K1_TT2_448_WG16_16_3_WGM1", "w4a16/fp16/Cijk_Alik_Bljk_HHS_BH_UserArgs_MT32x7168x16_SN_K1_TT2_448_batch_w4a16_gate2.co", 32, 7168, 768}},
        {21001, {"Cijk_Alik_Bljk_HHS_BH_UserArgs_MT32x1024x16_SN_K1_TT2_64_WG16_16_3_WGM1", "w4a16/fp16/Cijk_Alik_Bljk_HHS_BH_UserArgs_MT32x1024x16_SN_K1_TT2_16_WG16_16_3_WGM1_vllm_batch_gate2.co", 32, 1024, 768}},
        {21002, {"Cijk_Alik_Bljk_HHS_BH_UserArgs_MT32x1024x16_SN_K1_TT2_64_WG16_16_3_WGM1", "w4a16/fp16/Cijk_Alik_Bljk_HHS_BH_UserArgs_MT32x1024x16_SN_K1_TT2_16_WG16_16_3_WGM1_batch_w4a16.co", 32, 1024, 768}},
        {40000, {"Cijk_Alik_Bljk_HHS_BH_UserArgs_MT16x1024x16_SN_K1_TT1_64_WG16_16_3_WGM1", "w4a16/fp16/Cijk_Alik_Bljk_HHS_BH_UserArgs_MT16x1024x16_SN_K1_TT1_64_batch_w4a16_gate2_qianwen.co", 16, 1024, 768}}};

    static std::unordered_map<int, FMoeKernelConfig> moe2_kernel_w4a16_bf16_configs = {
        {20000, {"Cijk_Alik_Bljk_BBS_BH_UserArgs_MT16x1024x16_SN_K1_TT1_64_WG16_16_3_WGM1", "w4a16/bf16/Cijk_Alik_Bljk_BBS_BH_UserArgs_MT16x1024x16_SN_K1_TT1_64_batch_w4a16_gate2.co", 16, 1024, 768}},
        {20001, {"Cijk_Alik_Bljk_BBS_BH_UserArgs_MT16x3584x16_SN_K1_TT1_224_WG16_16_3_WGM1", "w4a16/bf16/Cijk_Alik_Bljk_BBS_BH_UserArgs_MT16x3584x16_SN_K1_TT1_224_batch_w4a16_gate2.co", 16, 3584, 768}},
        {20002, {"Cijk_Alik_Bljk_BBS_BH_UserArgs_MT16x7168x16_SN_K1_TT1_448_WG16_16_3_WGM1", "w4a16/bf16/Cijk_Alik_Bljk_BBS_BH_UserArgs_MT16x7168x16_SN_K1_TT1_448_batch_w4a16_gate2.co", 16, 7168, 768}},
        {21000, {"Cijk_Alik_Bljk_BBS_BH_UserArgs_MT32x7168x16_SN_K1_TT2_448_WG16_16_3_WGM1", "w4a16/bf16/Cijk_Alik_Bljk_BBS_BH_UserArgs_MT32x7168x16_SN_K1_TT2_448_batch_w4a16_gate2.co", 32, 7168, 768}},
        {21001, {"Cijk_Alik_Bljk_BBS_BH_UserArgs_MT32x1024x16_SN_K1_TT2_64_WG16_16_3_WGM1", "w4a16/bf16/Cijk_Alik_Bljk_BBS_BH_UserArgs_MT32x1024x16_SN_K1_TT2_16_WG16_16_3_WGM1_vllm_batch_gate2.co", 32, 1024, 768}},
        {21002, {"Cijk_Alik_Bljk_BBS_BH_UserArgs_MT32x1024x16_SN_K1_TT2_64_WG16_16_3_WGM1", "w4a16/bf16/Cijk_Alik_Bljk_BBS_BH_UserArgs_MT32x1024x16_SN_K1_TT2_16_WG16_16_3_WGM1_batch_w4a16.co", 32, 1024, 768}},
        {40000, {"Cijk_Alik_Bljk_BBS_BH_UserArgs_MT16x1024x16_SN_K1_TT1_64_WG16_16_3_WGM1", "w4a16/bf16/Cijk_Alik_Bljk_BBS_BH_UserArgs_MT16x1024x16_SN_K1_TT1_64_batch_w4a16_gate2_qianwen.co", 16, 1024, 768}}};

    static std::unordered_map<int, FMoeKernelConfig> moe2_kernel_w16a16_configs = {
        {20000, {"Cijk_Alik_Bljk_HHS_BH_UserArgs_MT32x1024x16_SN_K1_TT2_64_WG16_16_3_WGM1", "w16a16/Cijk_Alik_Bljk_HHS_BH_UserArgs_MT32x1024x16_SN_K1_TT2_16_WG16_16_3_WGM1_vllm.co", 32, 1024, 768}}};
    
    static std::unordered_map<int, FMoeKernelConfig> moe2_kernel_w4a16_bf16_group32_configs = {
        {60032, {"Cijk_Alik_Bljk_BBS_BH_UserArgs_MT32x1024x16_SN_K1_TT2_64_WG16_16_3_WGM1", "w4a16/bf16/Cijk_Alik_Bljk_BBS_BH_UserArgs_MT32x1024x16_SN_K1_TT2_16_WG16_16_3_WGM1_vllm_batch_gate2_gw32.co", 32, 1024, 768}}};

    torch::Tensor ScaleA;
    torch::Tensor ScaleB;
    torch::Tensor ZeroPoints;
    int Mode = 0;

    if (scale_a.has_value())
    {
      ScaleA = scale_a.value();
    }
    if (scale_b.has_value())
    {
      ScaleB = scale_b.value();
    }
    if (zero_points.has_value())
    {
      ZeroPoints = zero_points.value();
    }
    if (mode.has_value())
    {
      Mode = mode.value();
    }

    // g1u0
    FMoeKernel *impl_ptr = nullptr;

    static std::unordered_map<std::string, std::unique_ptr<FMoeKernel>> impl_ptr_map;
    std::vector<uint32_t> Config = {16, 7168, 768};

    std::unordered_map<int, FMoeKernelConfig> *config_map = nullptr;

    if (Mode == 2) {
      config_map = &moe2_kernel_w4a16_fp16_configs;
    }
    else if (Mode == 3) {
      config_map = &moe2_kernel_w4a16_bf16_configs;
    }
    else if (Mode == 4) {
      config_map = &moe2_kernel_w4a16_bf16_group32_configs;
    }
    else {
      config_map = &moe2_kernel_w16a16_configs;
    }

    int Solution = 20000;
    if (solidx.has_value())
    {
      Solution = solidx.value();
    }

    if (solidx.value() == 0)
    {
      Moe2SolutionLookup lookup;
      Solution = lookup.getSolution(out.size(0), down.size(2)*2);
    }

    if (!config_map)
    {
        TORCH_CHECK(false, __func__, " Input only supput Int8!");
    }

    auto it = config_map->find(Solution);
    if (it != config_map->end())
    {
        const auto &config = it->second;
        const char *name = config.name.c_str();
        const char *co_name = config.co_name.c_str();
        Config = {config.MT0, config.MT1, config.bdx};

        auto result = impl_ptr_map.emplace(name, nullptr);
        if (result.second)
        {
            result.first->second = std::make_unique<FMoeKernel>(name, co_name);
        }
        impl_ptr = result.first->second.get();
    }

    TORCH_CHECK(impl_ptr != nullptr,
                __func__, ": unsupport current input type:", input.scalar_type());
    impl_ptr->launch_kernel<uint16_t, uint16_t, false>(Config,
                                                    out,
                                                    input,
                                                    gate,
                                                    down,
                                                    sorted_token_ids,
                                                    sorted_weights,
                                                    sorted_expert_ids,
                                                    num_valid_ids,
                                                    top_k,
                                                    ScaleA,
                                                    ScaleB,
                                                    ZeroPoints,
                                                    Mode,
                                                    block_size,
                                                    persist_groups);
}
