// #include <torch/all.h>
// #include <c10/cuda/CUDAGuard.h>
// #include <ATen/cuda/CUDAContext.h>
// #include <cuda_runtime.h>

// #include <cuda_fp16.h>
// #include <cuda_bf16.h>
// #include "moe_wna16_utils.h"


// #include "moe_ops.h"

#pragma once
#include "moe_w8a16_block_wise.h"
#include "moe_w8a16_awq.h"
#include "moe_w4a16.h"
#include "moe_w4a16_2.h"
#include "moe_w4a16_base.h"
#include "moe_w8a8_block_wise.h"
#include "moe_w8a8_block_wise_kernel2.h"
#include "moe_w8a8_block_wise_fp8.h"
#include "moe_w8a8_block_wise_kernel2_fp8.h"

#include "topk_softmax_kernel.h"
#include "moe_align_sum_kernels.h"
#include "moe_sum_kernels_opt_v2.h"
#include "silu_mul_kernels.h"
#include <torch/all.h>
#include <optional>
#include <vector>
#include <cstdint>
#include <fstream>
// #include "moe_w8a8_utils.h"
// #include "moe_w8a8_config.h"
#include "moe_w8a8_opt.h"
#include "moe_w4a16_opt.h"
#include "moe_w4a8_opt.h"
#include "moe_w8a16_chan_opt.h"

#define BIT_SWITCH(bit, BIT, ...)                           \
[&] {                                                       \
  if (bit == 8) {                                           \
    constexpr static int BIT = 8;                           \
    return __VA_ARGS__();                                   \
      }else if (bit == 4) {                    \
    constexpr static int BIT = 4;          \
    return __VA_ARGS__();                   \
  } \
   else {                                                  \
    std::cout<<"unsupported BIT"<<std::endl;                \
  }                                                         \
}()

#define BLOCK_M_SWITCH(BLOCK_SIZE_M, BLOCK_SIZE_M_, ...)    \
[&] {                                                       \
  if (BLOCK_SIZE_M == 16) {                                  \
    constexpr static int BLOCK_SIZE_M_ = 16;                 \
    return __VA_ARGS__();                                   \
  }else if(BLOCK_SIZE_M == 32) {\
    constexpr static int BLOCK_SIZE_M_ = 32;                 \
    return __VA_ARGS__(); \
  }else if(BLOCK_SIZE_M == 48) {\
    constexpr static int BLOCK_SIZE_M_ = 48;                 \
    return __VA_ARGS__(); \
  }else if(BLOCK_SIZE_M == 64) {\
    constexpr static int BLOCK_SIZE_M_ = 64;                 \
    return __VA_ARGS__(); \
  }else {                                                  \
    std::cout<<"unsupported BLOCK_SIZE_M"<<std::endl;       \
  }                                                         \
}()

#define BLOCK_N_SWITCH(BLOCK_SIZE_N, BLOCK_SIZE_N_, ...)    \
[&] {                                                       \
  if (BLOCK_SIZE_N == 64) {                                \
    constexpr static int BLOCK_SIZE_N_ = 64;               \
    return __VA_ARGS__();                                   \
  }else {                                                  \
    std::cout<<"unsupported BLOCK_SIZE_N"<<std::endl;       \
  }                                                         \
}()

#define BLOCK_K_SWITCH(BLOCK_SIZE_K, BLOCK_SIZE_K_, ...)    \
[&] {                                                       \
  if (BLOCK_SIZE_K == 128) {                                \
    constexpr static int BLOCK_SIZE_K_ = 128;               \
    return __VA_ARGS__();                                   \
  }else {                                                  \
    std::cout<<"unsupported BLOCK_SIZE_K"<<std::endl;       \
  }                                                         \
}()

#define BOOL_SWITCH(COND, CONST_NAME, ...)      \
  [&] {                                         \
    if (COND) {                                 \
      constexpr static bool CONST_NAME = true;  \
      return __VA_ARGS__();                     \
    } else {                                    \
      constexpr static bool CONST_NAME = false; \
      return __VA_ARGS__();                     \
    }                                           \
  }()

#define TOPK_SWITCH(topk, TOPK, ...)   \
[&] {                                       \
  if (topk == 8) {                          \
    constexpr static int TOPK = 8;          \
    return __VA_ARGS__();                   \
  }else if (topk == 1) {                    \
    constexpr static int TOPK = 1;          \
    return __VA_ARGS__();                   \
  }else {                                  \
    std::cout<<"unsupported TOPK"<<std::endl;\
  }                                         \
}()

#define GROUP_SIZE_N_SWITCH(group_size_n, GROUP_SIZE_N, ...)      \
[&] {                                                       \
  if (group_size_n == 128) {                                   \
    constexpr static int GROUP_SIZE_N = 128;                   \
    return __VA_ARGS__();                                   \
  }else {                                                  \
    std::cout<<"unsupported GROUP_SIZE_N"<<std::endl;         \
  }                                                         \
}()

#define GROUP_SIZE_K_SWITCH(group_size_k, GROUP_SIZE_K, ...)      \
[&] {                                                       \
  if (group_size_k == 128) {                                   \
    constexpr static int GROUP_SIZE_K = 128;                   \
    return __VA_ARGS__();                                     \
      }else if(group_size_k == 64){                         \
    constexpr static int GROUP_SIZE_K = 64;            \
    return __VA_ARGS__();                                   \
  }else {                                                  \
    std::cout<<"unsupported GROUP_SIZE_K"<<std::endl;         \
  }                                                         \
}()

#define BLOCK_SIZE_M_LOOPS_SWITCH(block_size_m_loops, BLOCK_SIZE_M_LOOPS, ...) \
[&] {                                                       \
  if (block_size_m_loops == 1) {                            \
    constexpr static int BLOCK_SIZE_M_LOOPS = 1;            \
    return __VA_ARGS__();                                   \
  } else {                                                  \
    std::cout<<"unsupported BLOCK_SIZE_M_LOOPS"<<std::endl;         \
  }                                                         \
}()

#define BLOCK_SIZE_N_LOOPS_SWITCH(block_size_n_loops, BLOCK_SIZE_N_LOOPS, ...) \
[&] {                                                       \
  if (block_size_n_loops ==4) {                            \
    constexpr static int BLOCK_SIZE_N_LOOPS = 4;            \
    return __VA_ARGS__();                                   \
  }else if(block_size_n_loops ==1) {\
     constexpr static int BLOCK_SIZE_N_LOOPS = 1;            \
    return __VA_ARGS__();                                   \
  }else if(block_size_n_loops ==2) {\
     constexpr static int BLOCK_SIZE_N_LOOPS = 2;            \
    return __VA_ARGS__();                                   \
  }else {                                                  \
    std::cout<<"unsupported BLOCK_SIZE_N_LOOPS"<<std::endl;         \
  }                                                         \
}()

#define BLOCK_SIZE_K_LOOPS_SWITCH(block_size_k_loops, BLOCK_SIZE_K_LOOPS, ...) \
[&] {                                                       \
  if (block_size_k_loops == 1) {                            \
    constexpr static int BLOCK_SIZE_K_LOOPS = 1;            \
    return __VA_ARGS__();                                   \
  }else if(block_size_k_loops == 2) {\
    constexpr static int BLOCK_SIZE_K_LOOPS = 2;            \
    return __VA_ARGS__();                                   \
  }else if(block_size_k_loops == 4) {\
    constexpr static int BLOCK_SIZE_K_LOOPS = 4;            \
    return __VA_ARGS__();                                   \
  }else if(block_size_k_loops == 7) {\
    constexpr static int BLOCK_SIZE_K_LOOPS = 7;            \
    return __VA_ARGS__();                                   \
  }else if(block_size_k_loops == 8) {\
    constexpr static int BLOCK_SIZE_K_LOOPS = 8;            \
    return __VA_ARGS__();                                   \
  }else if(block_size_k_loops == 14) {\
    constexpr static int BLOCK_SIZE_K_LOOPS = 14;            \
    return __VA_ARGS__();                                   \
  }else if(block_size_k_loops == 28) {\
    constexpr static int BLOCK_SIZE_K_LOOPS = 28;            \
    return __VA_ARGS__();                                   \
  }else if(block_size_k_loops == 56) {\
    constexpr static int BLOCK_SIZE_K_LOOPS = 56;            \
    return __VA_ARGS__();                                   \
  }else {                                                  \
    std::cout<<"unsupported BLOCK_SIZE_K_LOOPS"<<std::endl;         \
  }                                                         \
}()

#define GROUP_SIZE_SWITCH(group_size, GROUP_SIZE, ...)      \
[&] {                                                       \
  if (group_size == 64) {                                   \
    constexpr static int GROUP_SIZE = 64;                   \
    return __VA_ARGS__();                                   \
  }else {                                                  \
    std::cout<<"unsupported GROUP_SIZE"<<std::endl;         \
  }                                                         \
}()






template <typename OutT, typename ElemT>
static void moe_marlin_w8a8_dispatch_gemm_stages(
    bool first_stage,
    const torch::Tensor& input,
    const torch::Tensor& b_qweight,
    torch::Tensor& output_alias,
    const torch::Tensor& a_scale,
    const torch::Tensor& b_scale,
    const float* topk_weights_ptr,
    const torch::Tensor& sorted_token_ids,
    const torch::Tensor& expert_ids,
    int num_pad,
    const torch::Tensor& num_tokens_post_pad,
    int size_m,
    int size_n,
    int size_k,
    int stride_asm,
    int stride_ask,
    int stride_bse,
    int stride_bsn,
    int stride_bsk,
    int64_t top_k,
    uint32_t real_topk,
    bool is_marlin,
    int64_t mode,
    int config_m,
    bool tensorwise_scale = false) {
  static_assert(
      std::is_same_v<ElemT, at::Float8_e4m3fn> || std::is_same_v<ElemT, int8_t>,
      "ElemT must be at::Float8_e4m3fn or int8_t");
  if (first_stage) {
    const int64_t EM = sorted_token_ids.size(0);
    GemmParams<char, OutT> params_in(
        (const char*)input.data_ptr<ElemT>(),
        (const char*)b_qweight.data_ptr<ElemT>(),
        (OutT*)output_alias.data_ptr(),
        (float*)a_scale.data_ptr(),
        (float*)b_scale.data_ptr(),
        topk_weights_ptr,
        sorted_token_ids.data_ptr<int32_t>(),
        expert_ids.data_ptr<int32_t>(),
        num_pad,
        num_tokens_post_pad.data_ptr<int32_t>(),
        size_m,
        size_n,
        size_k,
        stride_asm,
        stride_ask,
        stride_bse,
        stride_bsn,
        stride_bsk,
        EM,
        top_k,
        real_topk,
        is_marlin,
        tensorwise_scale);

    if (config_m <= 512) {
      if constexpr (std::is_same_v<ElemT, at::Float8_e4m3fn>) {
        auto it = kernel_maps_gemm1_decode_fp8<OutT>.find(mode);
        if (it != kernel_maps_gemm1_decode_fp8<OutT>.end()) {
          it->second(params_in);
        } else {
          printf("bfloat version gemm1 No matching kernel configuration found, using default settings \n");
        }
      } else {
        auto it = kernel_maps_gemm1_decode<OutT>.find(mode);
        if (it != kernel_maps_gemm1_decode<OutT>.end()) {
          it->second(params_in);
        } else {
          if constexpr (std::is_same_v<OutT, half>) {
            printf("half version gemm1 No matching kernel configuration found, using default settings \n");
          } else {
            printf("bfloat version gemm1 No matching kernel configuration found, using default settings \n");
          }
        }
      }
    } else {
      if constexpr (std::is_same_v<ElemT, at::Float8_e4m3fn>) {
        auto it = kernel_maps_gemm1_prefill_fp8<OutT>.find(mode);
        if (it != kernel_maps_gemm1_prefill_fp8<OutT>.end()) {
          it->second(params_in);
        } else {
          printf("bfloat version gemm1 No matching kernel configuration found, using default settings \n");
        }
      } else {
        auto it = kernel_maps_gemm1_prefill<OutT>.find(mode);
        if (it != kernel_maps_gemm1_prefill<OutT>.end()) {
          it->second(params_in);
        } else {
          if constexpr (std::is_same_v<OutT, half>) {
            printf("half version gemm1 No matching kernel configuration found \n");
          } else {
            printf("bfloat version gemm1 No matching kernel configuration found, using default settings \n");
          }
        }
      }
    }
  } else {
    const int64_t EM = sorted_token_ids.size(0);
    GemmParams<char, OutT> params_in(
        (const char*)input.data_ptr<ElemT>(),
        (const char*)b_qweight.data_ptr<ElemT>(),
        (OutT*)output_alias.data_ptr(),
        (float*)a_scale.data_ptr(),
        (float*)b_scale.data_ptr(),
        topk_weights_ptr,
        sorted_token_ids.data_ptr<int32_t>(),
        expert_ids.data_ptr<int32_t>(),
        num_pad,
        num_tokens_post_pad.data_ptr<int32_t>(),
        size_m,
        size_n,
        size_k,
        stride_asm,
        stride_ask,
        stride_bse,
        stride_bsn,
        stride_bsk,
        EM,
        top_k,
        real_topk,
        is_marlin,
        tensorwise_scale);

    if (config_m <= 512) {
      if constexpr (std::is_same_v<ElemT, at::Float8_e4m3fn>) {
        auto it = kernel_maps_gemm2_decode_fp8<OutT>.find(mode);
        if (it != kernel_maps_gemm2_decode_fp8<OutT>.end()) {
          it->second(params_in);
        } else {
          printf("bfloat version gemm2 No matching kernel configuration found, using default settings \n");
        }
      } else {
        auto it = kernel_maps_gemm2_decode<OutT>.find(mode);
        if (it != kernel_maps_gemm2_decode<OutT>.end()) {
          it->second(params_in);
        } else {
          if constexpr (std::is_same_v<OutT, half>) {
            printf("half version gemm2 No matching kernel configuration found, using default settings \n");
          } else {
            printf("bfloat version gemm2 No matching kernel configuration found, using default settings \n");
          }
        }
      }
    } else {
      if constexpr (std::is_same_v<ElemT, at::Float8_e4m3fn>) {
        auto it = kernel_maps_gemm2_prefill_fp8<OutT>.find(mode);
        if (it != kernel_maps_gemm2_prefill_fp8<OutT>.end()) {
          it->second(params_in);
        } else {
          printf("bfloat version gemm2 No matching kernel configuration found, using default settings \n");
        }
      } else {
        auto it = kernel_maps_gemm2_prefill<OutT>.find(mode);
        if (it != kernel_maps_gemm2_prefill<OutT>.end()) {
          it->second(params_in);
        } else {
          if constexpr (std::is_same_v<OutT, half>) {
            printf("half version gemm2  No matching kernel configuration found \n");
          } else {
            printf("bfloat version gemm2 No matching kernel configuration found, using default settings \n");
          }
        }
      }
    }
  }
}



template <typename OutT, typename ElemT>
static void moe_marlin_w8a8_dispatch_gemm_stages_n160(
    bool first_stage,
    const torch::Tensor& input,
    const torch::Tensor& b_qweight,
    torch::Tensor& output_alias,
    const torch::Tensor& a_scale,
    const torch::Tensor& b_scale,
    const float* topk_weights_ptr,
    const torch::Tensor& sorted_token_ids,
    const torch::Tensor& expert_ids,
    int num_pad,
    const torch::Tensor& num_tokens_post_pad,
    int size_m,
    int size_n,
    int size_k,
    int stride_asm,
    int stride_ask,
    int stride_bse,
    int stride_bsn,
    int stride_bsk,
    int64_t top_k,
    uint32_t real_topk,
    bool is_marlin,
    int64_t mode,
    int config_m,
    bool tensorwise_scale = false) {
  static_assert(
      std::is_same_v<ElemT, at::Float8_e4m3fn> || std::is_same_v<ElemT, int8_t>,
      "ElemT must be at::Float8_e4m3fn or int8_t");
  if (first_stage) {
    const int64_t EM = sorted_token_ids.size(0);
    GemmParams<char, OutT> params_in(
        (const char*)input.data_ptr<ElemT>(),
        (const char*)b_qweight.data_ptr<ElemT>(),
        (OutT*)output_alias.data_ptr(),
        (float*)a_scale.data_ptr(),
        (float*)b_scale.data_ptr(),
        topk_weights_ptr,
        sorted_token_ids.data_ptr<int32_t>(),
        expert_ids.data_ptr<int32_t>(),
        num_pad,
        num_tokens_post_pad.data_ptr<int32_t>(),
        size_m,
        size_n,
        size_k,
        stride_asm,
        stride_ask,
        stride_bse,
        stride_bsn,
        stride_bsk,
        EM,
        top_k,
        real_topk,
        is_marlin,
        tensorwise_scale);

    if (config_m <= 512) {
      if constexpr (std::is_same_v<ElemT, at::Float8_e4m3fn>) {

        // std::cout<<"gemm1******************mode "<<mode<<std::endl;
        auto it = kernel_maps_gemm1_decode_n160_fp8<OutT>.find(mode);
        if (it != kernel_maps_gemm1_decode_n160_fp8<OutT>.end()) {
          it->second(params_in);
        } else {
          std::cout<<"moe_marlin_w8a8_dispatch_gemm_stages<half, at::Float8_e4m3fn> mode \n"<<mode<<std::endl;

        }
      } 
      // else {
      //   auto it = kernel_maps_gemm1_n160_decode<OutT>.find(mode);
      //   if (it != kernel_maps_gemm1_n160_decode<OutT>.end()) {
      //     it->second(params_in);
      //   } else {
      //     if constexpr (std::is_same_v<OutT, half>) {
      //       printf("half version gemm1 No matching kernel configuration found, using default settings \n");
      //     } else {
      //       printf("bfloat version gemm1 No matching kernel configuration found, using default settings \n");
      //     }
      //   }
      // }
    } else {
      if constexpr (std::is_same_v<ElemT, at::Float8_e4m3fn>) {
        auto it = kernel_maps_gemm1_prefill_n160_fp8<OutT>.find(mode);
        if (it != kernel_maps_gemm1_prefill_n160_fp8<OutT>.end()) {
          it->second(params_in);
        } else {
          printf("bfloat version gemm1 No matching kernel configuration found, using default settings \n");
        }
      } 
      // else {
      //   auto it = kernel_maps_gemm1_prefill_n160<OutT>.find(mode);
      //   if (it != kernel_maps_gemm1_prefill_n160<OutT>.end()) {
      //     it->second(params_in);
      //   } else {
      //     if constexpr (std::is_same_v<OutT, half>) {
      //       printf("half version gemm1 No matching kernel configuration found \n");
      //     } else {
      //       printf("bfloat version gemm1 No matching kernel configuration found, using default settings \n");
      //     }
      //   }
      // }
    }
  } else {
    const int64_t EM = sorted_token_ids.size(0);
    GemmParams<char, OutT> params_in(
        (const char*)input.data_ptr<ElemT>(),
        (const char*)b_qweight.data_ptr<ElemT>(),
        (OutT*)output_alias.data_ptr(),
        (float*)a_scale.data_ptr(),
        (float*)b_scale.data_ptr(),
        topk_weights_ptr,
        sorted_token_ids.data_ptr<int32_t>(),
        expert_ids.data_ptr<int32_t>(),
        num_pad,
        num_tokens_post_pad.data_ptr<int32_t>(),
        size_m,
        size_n,
        size_k,
        stride_asm,
        stride_ask,
        stride_bse,
        stride_bsn,
        stride_bsk,
        EM,
        top_k,
        real_topk,
        is_marlin,
        tensorwise_scale);

    if (config_m <= 512 ) {
      if constexpr (std::is_same_v<ElemT, at::Float8_e4m3fn>) {
        auto it = kernel_maps_gemm2_decode_n160_fp8<OutT>.find(mode);
        if (it != kernel_maps_gemm2_decode_n160_fp8<OutT>.end()) {
          it->second(params_in);
        } else {
          printf("bfloat version gemm2 No matching kernel configuration found, using default settings \n");
        }
      } 
      // else {
      //   auto it = kernel_maps_gemm2_n160_decode<OutT>.find(mode);
      //   if (it != kernel_maps_gemm2_n160_decode<OutT>.end()) {
      //     it->second(params_in);
      //   } else {
      //     if constexpr (std::is_same_v<OutT, half>) {
      //       printf("half version gemm2 No matching kernel configuration found, using default settings \n");
      //     } else {
      //       printf("bfloat version gemm2 No matching kernel configuration found, using default settings \n");
      //     }
      //   }
      // }
    } else {
      if constexpr (std::is_same_v<ElemT, at::Float8_e4m3fn>) {
        auto it = kernel_maps_gemm2_prefill_n160_fp8<OutT>.find(mode);
        if (it != kernel_maps_gemm2_prefill_n160_fp8<OutT>.end()) {
          it->second(params_in);
        } else {
          printf("bfloat version gemm2 No matching kernel configuration found, using default settings \n");
        }
      } 
      // else {
      //   auto it = kernel_maps_gemm2_prefill_n160<OutT>.find(mode);
      //   if (it != kernel_maps_gemm2_prefill_n160<OutT>.end()) {
      //     it->second(params_in);
      //   } else {
      //     if constexpr (std::is_same_v<OutT, half>) {
      //       printf("half version gemm2  No matching kernel configuration found \n");
      //     } else {
      //       printf("bfloat version gemm2 No matching kernel configuration found, using default settings \n");
      //     }
      //   }
      // }
    }
  }
}

















// 模板抽象
// BLOCK_MNK: [16, 128, 128]
// WARP_MNK: [16, 32, 64]
// MMA_MNK: [16, 16, 32]
// BLOCK_MNK / WARP_MNK = [1, 4, 2]  代表warp在MN方向的排布 block_k/warp_k代表stage
// WARP_MNK / MMA_MNK = [1, 2, 2]    代表warp在MNK方向重复计算的次数 会分配额外的寄存器
static torch::Tensor moe_c_moe_gemm_marlin_w8a8_impl(torch::Tensor input,
  torch::Tensor b_qweight,
  torch::Tensor output,
  torch::Tensor a_scale,
  torch::Tensor b_scale,
  std::optional<torch::Tensor> topk_weights,
  torch::Tensor sorted_token_ids, 
  torch::Tensor expert_ids,
  torch::Tensor num_tokens_post_pad, 
  int64_t top_k, // gemm1为topk  gemm2为1  因为gemm1输入为[m, k]  gemm2输入为[m*topk, k]
  int64_t mode,
  int64_t delta,
  int64_t config_m,
  bool tensorwise_scale
  ) {

    const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
    const int size_m = input.size(0); 
    const int EXPERTS = b_qweight.size(0);
    const int size_k = input.size(1);
    // std::cout<<"size_k"<<size_k<<std::endl;
    const int size_n = b_qweight.size(2)*b_qweight.size(1)/size_k;
    const int topk_size = output.size(1); // 输出为[m, topk_size, n]
    const int stride_asm = a_scale.stride(0);
    const int stride_ask = a_scale.stride(1);
    const int stride_bse = b_scale.stride(0);
    const int stride_bsn = b_scale.stride(1); 
    const int stride_bsk = b_scale.stride(2);
  // #if defined(__gfx938__)
    // std::cout<<"***********************************gfx938****************************************\n";
  // #endif
    
    const uint32_t real_topk = delta;
    
    constexpr int GROUP_N = 1;
    constexpr int GROUP_K = 1;
    bool is_marlin = true; // weight为[E, N, K]时 代表不进行重排
    bool first_stage = true;
    torch::Tensor output_alias = output.alias();

    //printf("size_n: %d stride_bse: %d stride_bsn: %d stride_bsk: %d\n", size_n, stride_bse, stride_bsn, stride_bsk);

    const float* topk_weights_ptr; // 第一阶段这里为null
    if (topk_weights.has_value()){
      topk_weights_ptr = (const float*)topk_weights.value().data_ptr();
      first_stage = false;
    }

    int num_pad = 0;
    if (output.scalar_type() == at::ScalarType::BFloat16){
    if (input.scalar_type() == at::ScalarType::Char){
      if(first_stage){
     
        int64_t EM = sorted_token_ids.size(0); // 一维线性化的token id
        GemmParams<char,bhalf_t> params_in(
          (const char*)input.data_ptr<int8_t>(), 
          (const char*)b_qweight.data_ptr<int8_t>(), 
          (bhalf_t*)output_alias.data_ptr(),
          (float*)a_scale.data_ptr(), 
          (float*)b_scale.data_ptr(),  
          topk_weights_ptr,
          sorted_token_ids.data_ptr<int32_t>(),
          expert_ids.data_ptr<int32_t>(), 
          num_pad, //num_tokens_post_pad[0].item<int>(), //这里获取值 会造成device->host的拷贝和一部分空泡
          num_tokens_post_pad.data_ptr<int32_t>(),
          size_m,
          size_n,
          size_k,
          stride_asm,
          stride_ask,
          stride_bse,
          stride_bsn,
          stride_bsk,
          EM,
          top_k,
          real_topk,
          is_marlin,
          tensorwise_scale
        );

        if(config_m <= 512){
          auto it = kernel_maps_gemm1_decode<bhalf_t>.find(mode);
          if (it != kernel_maps_gemm1_decode<bhalf_t>.end()) {
              it->second(params_in);
          } else {
              printf("bfloat version gemm1 No matching kernel configuration found, using default settings \n");
              // launch_moe_w8a8_first_stage_prefill<16, 64, 128, 16, 16, 64>(params_in);
          }
        }else{ //decode 
          // std::cout<<"***************************************decode \n" << mode;
          
          auto it = kernel_maps_gemm1_prefill<bhalf_t>.find(mode);
          // printf()
          if ( it != kernel_maps_gemm1_prefill<bhalf_t>.end()) {
              it->second(params_in);
              
          } else {
              printf("bfloat version gemm1 No matching kernel configuration found, using default settings \n");
              // launch_moe_w8a8_first_stage_decode<32, 256, 64, 32, 128, 64, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<64, 64, 128, 64, 64, 128, 1>(params_in);
              // launch_moe_w8a8_first_stage_decode<48, 16, 256, 48, 16, 128, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 32, 64, 16, 16, 64, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<32, 64, 128, 32, 64, 64, 1>(params_in);
          }
        }
        
        
      }else{ //gemm2
        int64_t EM = sorted_token_ids.size(0); // 一维线性化的token id
          
        // 使用int8类型处理
        GemmParams<char,bhalf_t> params_in(
          (const char*)input.data_ptr<int8_t>(), 
          (const char*)b_qweight.data_ptr<int8_t>(), 
          (bhalf_t*)output_alias.data_ptr(),
          (float*)a_scale.data_ptr(), 
          (float*)b_scale.data_ptr(),  
          topk_weights_ptr,
          sorted_token_ids.data_ptr<int32_t>(),
          expert_ids.data_ptr<int32_t>(), 
          num_pad, //num_tokens_post_pad[0].item<int>(),
          num_tokens_post_pad.data_ptr<int32_t>(),
          size_m,
          size_n,
          size_k,
          stride_asm,
          stride_ask,
          stride_bse,
          stride_bsn,
          stride_bsk,
          EM,
          top_k,
          real_topk,
          is_marlin,
          tensorwise_scale
        );

        if(config_m <= 512 ){
          auto it = kernel_maps_gemm2_decode<bhalf_t>.find(mode);
          if (it != kernel_maps_gemm2_decode<bhalf_t>.end()) {
              it->second(params_in);
          } else {
              printf("bfloat version gemm2 No matching kernel configuration found, using default settings \n");
              // launch_moe_w8a8_second_stage_prefill<16, 64, 128, 16, 16, 64>(params_in);
          }
        }else
        {
          // std::cout<<"mode: "<<mode<<std::endl;
          
          auto it = kernel_maps_gemm2_prefill<bhalf_t>.find(mode);
          if (    it != kernel_maps_gemm2_prefill<bhalf_t>.end()) {
            // printf("**********************************%d",mode);
              it->second(params_in);
          } else {

              printf("bfloat version gemm2 No matching kernel configuration found, using default settings \n");
              // launch_moe_w8a8_second_stage_decode<16, 64, 128, 16, 16, 128, 2>(params_in); 
              // launch_moe_w8a8_second_stage_decode<32, 1024, 64, 16, 128, 64, 1>(params_in);
              // launch_moe_w8a8_second_stage_decode<32, 128, 64, 32, 32, 64, 2>(params_in);
              // launch_moe_w8a8_second_stage_decode<64, 64, 64, 64, 32, 64, 2>(params_in); 
              // launch_moe_w8a8_second_stage_decode<32, 256, 128, 32, 128, 128, 1>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 64, 128, 16, 16, 64, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 128, 128, 16, 16, 64, 2>(params_in);

          }

        }
              // hipDeviceSynchronize();

      }
    } else {
      // TORCH_CHECK(false, "moe_w8a8_gemm only supports int8");
    }
    
  }
  else if (output.scalar_type() == at::ScalarType::Half)
  {
    // printf("********************************************halfhalfhalf");
    if (input.scalar_type() == at::ScalarType::Char){
      if(first_stage){
     
        int64_t EM = sorted_token_ids.size(0); // 一维线性化的token id
        GemmParams<char,half> params_in(
          (const char*)input.data_ptr<int8_t>(), 
          (const char*)b_qweight.data_ptr<int8_t>(), 
          (half*)output_alias.data_ptr(),
          (float*)a_scale.data_ptr(), 
          (float*)b_scale.data_ptr(),  
          topk_weights_ptr,
          sorted_token_ids.data_ptr<int32_t>(),
          expert_ids.data_ptr<int32_t>(), 
          num_pad, //num_tokens_post_pad[0].item<int>(), //这里获取值 会造成device->host的拷贝和一部分空泡
          num_tokens_post_pad.data_ptr<int32_t>(),
          size_m,
          size_n,
          size_k,
          stride_asm,
          stride_ask,
          stride_bse,
          stride_bsn,
          stride_bsk,
          EM,
          top_k,
          real_topk,
          is_marlin,
          tensorwise_scale
        );

        if(config_m <= 512){
          auto it = kernel_maps_gemm1_decode<half>.find(mode);
          if (it != kernel_maps_gemm1_decode<half>.end()) {
              it->second(params_in);
          } else {
              printf("half version gemm1 No matching kernel configuration found, using default settings \n");
              // launch_moe_w8a8_first_stage_prefill<16, 64, 128, 16, 16, 64>(params_in);
          }
        }else{ //decode 
          // std::cout<<"***************************************decode \n" << mode;
          auto it = kernel_maps_gemm1_prefill<half>.find(mode);
          // printf()
          if ( it != kernel_maps_gemm1_prefill<half>.end()) {
              it->second(params_in);
              
          } else {
              printf("half version gemm1 No matching kernel configuration found \n");
              // launch_moe_w8a8_first_stage_decode<16, 16, 512, 16, 16, 128, 4>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 32, 64, 16, 16, 64, 2>(params_in);

          }
        }
        
        
      }else{ //gemm2
        int64_t EM = sorted_token_ids.size(0); // 一维线性化的token id
          
        // 使用int8类型处理
        GemmParams<char,half> params_in(
          (const char*)input.data_ptr<int8_t>(), 
          (const char*)b_qweight.data_ptr<int8_t>(), 
          (half*)output_alias.data_ptr(),
          (float*)a_scale.data_ptr(), 
          (float*)b_scale.data_ptr(),  
          topk_weights_ptr,
          sorted_token_ids.data_ptr<int32_t>(),
          expert_ids.data_ptr<int32_t>(), 
          num_pad, //num_tokens_post_pad[0].item<int>(),
          num_tokens_post_pad.data_ptr<int32_t>(),
          size_m,
          size_n,
          size_k,
          stride_asm,
          stride_ask,
          stride_bse,
          stride_bsn,
          stride_bsk,
          EM,
          top_k,
          real_topk,
          is_marlin,
          tensorwise_scale
        );

        if(config_m <= 512 ){
          auto it = kernel_maps_gemm2_decode<half>.find(mode);
          if (it != kernel_maps_gemm2_decode<half>.end()) {
              it->second(params_in);
          } else {
              printf("half version gemm2 No matching kernel configuration found, using default settings \n");
              // launch_moe_w8a8_first_stage_prefill<16, 64, 128, 16, 16, 64>(params_in);
          }
        }else
        {
          // std::cout<<"*************************float16"<<std::endl;
          // mode =86;
          auto it = kernel_maps_gemm2_prefill<half>.find(mode);
          if (  it != kernel_maps_gemm2_prefill<half>.end()) {
            // printf("**********************************%d",mode);
              it->second(params_in);
          } else {

              printf("half version gemm2  No matching kernel configuration found \n");
              // launch_moe_w8a8_second_stage_decode<16, 64, 128, 16, 16, 128, 2>(params_in); 

              // launch_moe_w8a8_second_stage_decode<16, 256, 64, 16, 32, 64, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 64, 128, 16, 16, 64, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 128, 128, 16, 16, 64, 2>(params_in);

          }

        }
              // hipDeviceSynchronize();

      }
    } else {
      TORCH_CHECK(false, "moe_w8a8_gemm only supports int8");
    }
  }

    return output;
}

torch::Tensor moe_c_moe_gemm_marlin_w8a8(torch::Tensor input,
  torch::Tensor b_qweight,
  torch::Tensor output,
  torch::Tensor a_scale,
  torch::Tensor b_scale,
  std::optional<torch::Tensor> topk_weights,
  torch::Tensor sorted_token_ids, 
  torch::Tensor expert_ids,
  torch::Tensor num_tokens_post_pad, 
  int64_t top_k,
  int64_t mode,
  int64_t delta,
  int64_t config_m
  ) {
    return moe_c_moe_gemm_marlin_w8a8_impl(input, b_qweight, output, a_scale,
      b_scale, topk_weights, sorted_token_ids, expert_ids, num_tokens_post_pad,
      top_k, mode, delta, config_m, false);
}

torch::Tensor moe_c_moe_gemm_marlin_w8a8_tensorwise(torch::Tensor input,
  torch::Tensor b_qweight,
  torch::Tensor output,
  torch::Tensor a_scale,
  torch::Tensor b_scale,
  std::optional<torch::Tensor> topk_weights,
  torch::Tensor sorted_token_ids, 
  torch::Tensor expert_ids,
  torch::Tensor num_tokens_post_pad, 
  int64_t top_k,
  int64_t mode,
  int64_t delta,
  int64_t config_m
  ) {
    return moe_c_moe_gemm_marlin_w8a8_impl(input, b_qweight, output, a_scale,
      b_scale, topk_weights, sorted_token_ids, expert_ids, num_tokens_post_pad,
      top_k, mode, delta, config_m, true);
}


// 模板抽象
// BLOCK_MNK: [16, 128, 128]
// WARP_MNK: [16, 32, 64]
// MMA_MNK: [16, 16, 32]
// BLOCK_MNK / WARP_MNK = [1, 4, 2]  代表warp在MN方向的排布 block_k/warp_k代表stage
// WARP_MNK / MMA_MNK = [1, 2, 2]    代表warp在MNK方向重复计算的次数 会分配额外的寄存器
static torch::Tensor moe_c_moe_gemm_marlin_w8a8_fp8_impl(torch::Tensor input,
  torch::Tensor b_qweight,
  torch::Tensor output,
  torch::Tensor a_scale,
  torch::Tensor b_scale,
  std::optional<torch::Tensor> topk_weights,
  torch::Tensor sorted_token_ids, 
  torch::Tensor expert_ids,
  torch::Tensor num_tokens_post_pad, 
  int64_t top_k, // gemm1为topk  gemm2为1  因为gemm1输入为[m, k]  gemm2输入为[m*topk, k]
  int64_t mode,
  int64_t delta,
  int64_t config_m,
  bool tensorwise_scale
  ) {


    const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
    const int size_m = input.size(0); 
    const int EXPERTS = b_qweight.size(0);
    const int size_k = input.size(1);
    // std::cout<<"size_k"<<size_k<<std::endl;
    const int size_n = b_qweight.size(2)*b_qweight.size(1)/size_k;
    const int topk_size = output.size(1); // 输出为[m, topk_size, n]
    const int stride_asm = a_scale.stride(0);
    const int stride_ask = a_scale.stride(1);
    const int stride_bse = b_scale.stride(0);
    const int stride_bsn = b_scale.stride(1); 
    const int stride_bsk = b_scale.stride(2);
    
    const uint32_t real_topk = delta;
    
    constexpr int GROUP_N = 1;
    constexpr int GROUP_K = 1;
    bool is_marlin = true; // weight为[E, N, K]时 代表不进行重排
    bool first_stage = true;
    torch::Tensor output_alias = output.alias();

    //printf("size_n: %d stride_bse: %d stride_bsn: %d stride_bsk: %d\n", size_n, stride_bse, stride_bsn, stride_bsk);

    const float* topk_weights_ptr; // 第一阶段这里为null
    if (topk_weights.has_value()){
      topk_weights_ptr = (const float*)topk_weights.value().data_ptr();
      first_stage = false;
    }

    int num_pad = 0;

    if (output.scalar_type() == at::ScalarType::BFloat16){
    if (input.scalar_type() == at::ScalarType::Float8_e4m3fn){
      if(EXPERTS == 288){
        // std::cout<<"moe_marlin_w8a8_dispatch_gemm_stages<half, at::Float8_e4m3fn> mode \n"<<mode<<std::endl;

        
          moe_marlin_w8a8_dispatch_gemm_stages_n160<bhalf_t, at::Float8_e4m3fn>(
          first_stage,
          input,
          b_qweight,
          output_alias,
          a_scale,
          b_scale,
          topk_weights_ptr,
          sorted_token_ids,
          expert_ids,
          num_pad,
          num_tokens_post_pad,
          size_m,
          size_n,
          size_k,
          stride_asm,
          stride_ask,
          stride_bse,
          stride_bsn,
          stride_bsk,
          top_k,
          real_topk,
          is_marlin,
          mode,
          config_m,
          tensorwise_scale);
      }
      else{
        // std::cout<<"moe_marlin_w8a8_dispatch_gemm_stages_n160<bhalf_t, at::Float8_e4m3fn>\n";
        moe_marlin_w8a8_dispatch_gemm_stages<bhalf_t, at::Float8_e4m3fn>(
          first_stage,
          input,
          b_qweight,
          output_alias,
          a_scale,
          b_scale,
          topk_weights_ptr,
          sorted_token_ids,
          expert_ids,
          num_pad,
          num_tokens_post_pad,
          size_m,
          size_n,
          size_k,
          stride_asm,
          stride_ask,
          stride_bse,
          stride_bsn,
          stride_bsk,
          top_k,
          real_topk,
          is_marlin,
          mode,
          config_m,
          tensorwise_scale);
      }
      

          
        
        
       
    } else {
      // TORCH_CHECK(false, "moe_w8a8_gemm only supports fp8");
    }
    
  }
  else if (output.scalar_type() == at::ScalarType::Half)
  {
    // printf("********************************************halfhalfhalf");
    if (input.scalar_type() == at::ScalarType::Float8_e4m3fn){
      if(EXPERTS == 288){

        
          moe_marlin_w8a8_dispatch_gemm_stages_n160<half, at::Float8_e4m3fn>(
          first_stage,
          input,
          b_qweight,
          output_alias,
          a_scale,
          b_scale,
          topk_weights_ptr,
          sorted_token_ids,
          expert_ids,
          num_pad,
          num_tokens_post_pad,
          size_m,
          size_n,
          size_k,
          stride_asm,
          stride_ask,
          stride_bse,
          stride_bsn,
          stride_bsk,
          top_k,
          real_topk,
          is_marlin,
          mode,
          config_m,
          tensorwise_scale);
      }
      else{
         
        moe_marlin_w8a8_dispatch_gemm_stages<half, at::Float8_e4m3fn>(
          first_stage,
          input,
          b_qweight,
          output_alias,
          a_scale,
          b_scale,
          topk_weights_ptr,
          sorted_token_ids,
          expert_ids,
          num_pad,
          num_tokens_post_pad,
          size_m,
          size_n,
          size_k,
          stride_asm,
          stride_ask,
          stride_bse,
          stride_bsn,
          stride_bsk,
          top_k,
          real_topk,
          is_marlin,
          mode,
          config_m,
          tensorwise_scale);
      }
          

        
              // hipDeviceSynchronize();

      
    } else {
      // TORCH_CHECK(false, "moe_w8a8_gemm only supports fp8");
    }
    
  }



    return output;
}

torch::Tensor moe_c_moe_gemm_marlin_w8a8_fp8(torch::Tensor input,
  torch::Tensor b_qweight,
  torch::Tensor output,
  torch::Tensor a_scale,
  torch::Tensor b_scale,
  std::optional<torch::Tensor> topk_weights,
  torch::Tensor sorted_token_ids, 
  torch::Tensor expert_ids,
  torch::Tensor num_tokens_post_pad, 
  int64_t top_k,
  int64_t mode,
  int64_t delta,
  int64_t config_m
  ) {
    return moe_c_moe_gemm_marlin_w8a8_fp8_impl(input, b_qweight, output,
      a_scale, b_scale, topk_weights, sorted_token_ids, expert_ids,
      num_tokens_post_pad, top_k, mode, delta, config_m, false);
}

torch::Tensor moe_c_moe_gemm_marlin_w8a8_fp8_tensorwise(torch::Tensor input,
  torch::Tensor b_qweight,
  torch::Tensor output,
  torch::Tensor a_scale,
  torch::Tensor b_scale,
  std::optional<torch::Tensor> topk_weights,
  torch::Tensor sorted_token_ids, 
  torch::Tensor expert_ids,
  torch::Tensor num_tokens_post_pad, 
  int64_t top_k,
  int64_t mode,
  int64_t delta,
  int64_t config_m
  ) {
    return moe_c_moe_gemm_marlin_w8a8_fp8_impl(input, b_qweight, output,
      a_scale, b_scale, topk_weights, sorted_token_ids, expert_ids,
      num_tokens_post_pad, top_k, mode, delta, config_m, true);
}



// 模板抽象
// BLOCK_MNK: [16, 128, 128]
// WARP_MNK: [16, 32, 64]
// MMA_MNK: [16, 16, 32]
// BLOCK_MNK / WARP_MNK = [1, 4, 2]  代表warp在MN方向的排布 block_k/warp_k代表stage
// WARP_MNK / MMA_MNK = [1, 2, 2]    代表warp在MNK方向重复计算的次数 会分配额外的寄存器
torch::Tensor moe_c_moe_gemm_marlin_w4a8(torch::Tensor input,
  torch::Tensor b_qweight,
  torch::Tensor output,
  torch::Tensor a_scale,
  torch::Tensor b_scale,
  std::optional<torch::Tensor> topk_weights,
  torch::Tensor sorted_token_ids, 
  torch::Tensor expert_ids,
  torch::Tensor num_tokens_post_pad, 
  int64_t top_k, // gemm1为topk  gemm2为1  因为gemm1输入为[m, k]  gemm2输入为[m*topk, k]
  int64_t mode,
  int64_t delta,
  int64_t config_m
  ) {

    const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
    
    const int size_m = input.size(0); 
    const int EXPERTS = b_qweight.size(0);
    const int size_k = input.size(1);
    // std::cout<<"size_k"<<size_k<<std::endl;
    const int size_n = b_qweight.size(2)*b_qweight.size(1)/size_k * 2;
    const int topk_size = output.size(1); // 输出为[m, topk_size, n]
    const int stride_asm = a_scale.stride(0);
    const int stride_ask = a_scale.stride(1);
    const int stride_bse = b_scale.stride(0);
    const int stride_bsn = b_scale.stride(1); 
    const int stride_bsk = b_scale.stride(2);
  
    
    const uint32_t real_topk = delta;
    
    constexpr int GROUP_N = 1;
    constexpr int GROUP_K = 1;
    bool is_marlin = true; // weight为[E, N, K]时 代表不进行重排
    bool first_stage = true;
    torch::Tensor output_alias = output.alias();

    //printf("size_n: %d stride_bse: %d stride_bsn: %d stride_bsk: %d\n", size_n, stride_bse, stride_bsn, stride_bsk);

    const float* topk_weights_ptr; // 第一阶段这里为null
    if (topk_weights.has_value()){
      topk_weights_ptr = (const float*)topk_weights.value().data_ptr();
      first_stage = false;
    }

    int num_pad = 0;
    if(stride_bse == 256){
      if (output.scalar_type() == at::ScalarType::BFloat16){
      if (input.scalar_type() == at::ScalarType::Char){
      if(first_stage){
     
        int64_t EM = sorted_token_ids.size(0); // 一维线性化的token id
        GemmParams_w4a8<char,bhalf_t> params_in(
          (const char*)input.data_ptr<int8_t>(), 
          (const char*)b_qweight.data_ptr<int8_t>(), 
          (bhalf_t*)output_alias.data_ptr(),
          (float*)a_scale.data_ptr(), 
          (float*)b_scale.data_ptr(),  
          topk_weights_ptr,
          sorted_token_ids.data_ptr<int32_t>(),
          expert_ids.data_ptr<int32_t>(), 
          num_pad, //num_tokens_post_pad[0].item<int>(), //这里获取值 会造成device->host的拷贝和一部分空泡
          num_tokens_post_pad.data_ptr<int32_t>(),
          size_m,
          stride_bse,
          size_k,
          stride_asm,
          stride_ask,
          stride_bse,
          stride_bsn,
          stride_bsk,
          EM,
          top_k,
          real_topk,
          is_marlin
        );
        
        if(false){
          auto it = kernel_maps_gemm1_decode_w4a8<bhalf_t>.find(mode);
          if (it != kernel_maps_gemm1_decode_w4a8<bhalf_t>.end()) {
              it->second(params_in);
          } else {
              printf("Path1 bfloat version decode gemm1 No matching kernel configuration found, using default settings mode %d \n",mode);
              // launch_moe_w8a8_first_stage_prefill<16, 64, 128, 16, 16, 64>(params_in);
          }
        }else{ //decode 
          // std::cout<<"***************************************decode \n" << mode;
          
          auto it = kernel_maps_gemm1_prefill_w4a8_gemm1n256<bhalf_t>.find(mode);
          // printf()
          if ( it != kernel_maps_gemm1_prefill_w4a8_gemm1n256<bhalf_t>.end()) {
              it->second(params_in);
              
          } else {
              printf("Path1 bfloat version prefill gemm1 No matching kernel configuration found, using default settings mode %d size_m %d\n",mode,size_m);
              // launch_moe_w8a8_first_stage_decode<32, 256, 64, 32, 128, 64, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<64, 64, 128, 64, 64, 128, 1>(params_in);
              // launch_moe_w8a8_first_stage_decode<48, 16, 256, 48, 16, 128, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 32, 64, 16, 16, 64, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<32, 64, 128, 32, 64, 64, 1>(params_in);
          }
        }
        
        
      }else{ //gemm2
        int64_t EM = sorted_token_ids.size(0); // 一维线性化的token id
          
        // 使用int8类型处理
        GemmParams_w4a8<char,bhalf_t> params_in(
          (const char*)input.data_ptr<int8_t>(), 
          (const char*)b_qweight.data_ptr<int8_t>(), 
          (bhalf_t*)output_alias.data_ptr(),
          (float*)a_scale.data_ptr(), 
          (float*)b_scale.data_ptr(),  
          topk_weights_ptr,
          sorted_token_ids.data_ptr<int32_t>(),
          expert_ids.data_ptr<int32_t>(), 
          num_pad, //num_tokens_post_pad[0].item<int>(),
          num_tokens_post_pad.data_ptr<int32_t>(),
          size_m,
          stride_bse,
          size_k,
          stride_asm,
          stride_ask,
          stride_bse,
          stride_bsn,
          stride_bsk,
          EM,
          top_k,
          real_topk,
          is_marlin
        );

        if(false /* size_m <= 512 * real_topk */){
          auto it = kernel_maps_gemm2_decode_w4a8<bhalf_t>.find(mode);
          if (it != kernel_maps_gemm2_decode_w4a8<bhalf_t>.end()) {
              it->second(params_in);
          } else {
              printf("Path1 bfloat version decode gemm2 No matching kernel configuration found, using default settings mode %d size_m %d\n",mode,size_m);
              // launch_moe_w8a8_second_stage_prefill<16, 64, 128, 16, 16, 64>(params_in);
          }
        }else
        {
          // std::cout<<"mode: "<<mode<<std::endl;
          
          auto it = kernel_maps_gemm2_prefill_w4a8<bhalf_t>.find(mode);
          if (    it != kernel_maps_gemm2_prefill_w4a8<bhalf_t>.end()) {
            // printf("**********************************%d",mode);
              it->second(params_in);
          } else {

              printf("Path1 bfloat version prefill gemm2 No matching kernel configuration found, using default settings mode %d size_m %d\n",mode,size_m);
              // launch_moe_w8a8_second_stage_decode<16, 64, 128, 16, 16, 128, 2>(params_in); 
              // launch_moe_w8a8_second_stage_decode<32, 1024, 64, 16, 128, 64, 1>(params_in);
              // launch_moe_w8a8_second_stage_decode<32, 128, 64, 32, 32, 64, 2>(params_in);
              // launch_moe_w8a8_second_stage_decode<64, 64, 64, 64, 32, 64, 2>(params_in); 
              // launch_moe_w8a8_second_stage_decode<32, 256, 128, 32, 128, 128, 1>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 64, 128, 16, 16, 64, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 128, 128, 16, 16, 64, 2>(params_in);

          }

        }
              // hipDeviceSynchronize();

      }
    } else {
      // TORCH_CHECK(false, "moe_w8a8_gemm only supports int8");
    }
    
  }
  else if (output.scalar_type() == at::ScalarType::Half)
  {
    // printf("********************************************halfhalfhalf");
    if (input.scalar_type() == at::ScalarType::Char){
      if(first_stage){
     
        int64_t EM = sorted_token_ids.size(0); // 一维线性化的token id
        GemmParams_w4a8<char,half> params_in(
          (const char*)input.data_ptr<int8_t>(), 
          (const char*)b_qweight.data_ptr<int8_t>(), 
          (half*)output_alias.data_ptr(),
          (float*)a_scale.data_ptr(), 
          (float*)b_scale.data_ptr(),  
          topk_weights_ptr,
          sorted_token_ids.data_ptr<int32_t>(),
          expert_ids.data_ptr<int32_t>(), 
          num_pad, //num_tokens_post_pad[0].item<int>(), //这里获取值 会造成device->host的拷贝和一部分空泡
          num_tokens_post_pad.data_ptr<int32_t>(),
          size_m,
          stride_bse,
          size_k,
          stride_asm,
          stride_ask,
          stride_bse,
          stride_bsn,
          stride_bsk,
          EM,
          top_k,
          real_topk,
          is_marlin
        );

        if(false/* size_m <= 512 */){
          auto it = kernel_maps_gemm1_decode_w4a8<half>.find(mode);
          if (it != kernel_maps_gemm1_decode_w4a8<half>.end()) {
              it->second(params_in);
          } else {
              printf("Path1 half version decode gemm1 No matching kernel configuration found, using default settings mode %d size_m %d\n",mode,size_m);
              // launch_moe_w8a8_first_stage_prefill<16, 64, 128, 16, 16, 64>(params_in);
          }
        }else{ //decode 
          // std::cout<<"***************************************decode \n" << mode;
          auto it = kernel_maps_gemm1_prefill_w4a8_gemm1n256<half>.find(mode);
          // printf()
          if ( it != kernel_maps_gemm1_prefill_w4a8_gemm1n256<half>.end()) {
              it->second(params_in);
              
          } else {
              printf("Path1 half version prefill gemm1 No matching kernel configuration found, using default settings mode %d size_m %d\n",mode,size_m);
              // launch_moe_w8a8_first_stage_decode<16, 16, 512, 16, 16, 128, 4>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 32, 64, 16, 16, 64, 2>(params_in);

          }
        }
        
        
      }else{ //gemm2
        int64_t EM = sorted_token_ids.size(0); // 一维线性化的token id
          
        // 使用int8类型处理
        GemmParams_w4a8<char,half> params_in(
          (const char*)input.data_ptr<int8_t>(), 
          (const char*)b_qweight.data_ptr<int8_t>(), 
          (half*)output_alias.data_ptr(),
          (float*)a_scale.data_ptr(), 
          (float*)b_scale.data_ptr(),  
          topk_weights_ptr,
          sorted_token_ids.data_ptr<int32_t>(),
          expert_ids.data_ptr<int32_t>(), 
          num_pad, //num_tokens_post_pad[0].item<int>(),
          num_tokens_post_pad.data_ptr<int32_t>(),
          size_m,
          stride_bse,
          size_k,
          stride_asm,
          stride_ask,
          stride_bse,
          stride_bsn,
          stride_bsk,
          EM,
          top_k,
          real_topk,
          is_marlin
        );

        if(false/* size_m <= 512 * real_topk */){
          auto it = kernel_maps_gemm2_decode_w4a8<half>.find(mode);
          if (it != kernel_maps_gemm2_decode_w4a8<half>.end()) {
              it->second(params_in);
          } else {
              printf("Path1 half version decode gemm2 No matching kernel configuration found, using default settings mode %d size_m %d\n",mode,size_m);
              // launch_moe_w8a8_first_stage_prefill<16, 64, 128, 16, 16, 64>(params_in);
          }
        }else
        {
          // std::cout<<"*************************float16"<<std::endl;
          // mode =86;
          auto it = kernel_maps_gemm2_prefill_w4a8<half>.find(mode);
          if (  it != kernel_maps_gemm2_prefill_w4a8<half>.end()) {
            // printf("**********************************%d",mode);
              it->second(params_in);
          } else {

              printf("Path1 half version prefill gemm2 No matching kernel configuration found, using default settings mode %d size_m %d\n",mode,size_m);
              // launch_moe_w8a8_second_stage_decode<16, 64, 128, 16, 16, 128, 2>(params_in); 

              // launch_moe_w8a8_second_stage_decode<16, 256, 64, 16, 32, 64, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 64, 128, 16, 16, 64, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 128, 128, 16, 16, 64, 2>(params_in);

          }

        }
              // hipDeviceSynchronize();

      }
    } else {
      TORCH_CHECK(false, "moe_w8a8_gemm only supports int8");
    }
  }
    }
    else{
        if (output.scalar_type() == at::ScalarType::BFloat16){
        if (input.scalar_type() == at::ScalarType::Char){
        if(first_stage){
     
        int64_t EM = sorted_token_ids.size(0); // 一维线性化的token id
        GemmParams_w4a8<char,bhalf_t> params_in(
          (const char*)input.data_ptr<int8_t>(), 
          (const char*)b_qweight.data_ptr<int8_t>(), 
          (bhalf_t*)output_alias.data_ptr(),
          (float*)a_scale.data_ptr(), 
          (float*)b_scale.data_ptr(),  
          topk_weights_ptr,
          sorted_token_ids.data_ptr<int32_t>(),
          expert_ids.data_ptr<int32_t>(), 
          num_pad, //num_tokens_post_pad[0].item<int>(), //这里获取值 会造成device->host的拷贝和一部分空泡
          num_tokens_post_pad.data_ptr<int32_t>(),
          size_m,
          stride_bse,
          size_k,
          stride_asm,
          stride_ask,
          stride_bse,
          stride_bsn,
          stride_bsk,
          EM,
          top_k,
          real_topk,
          is_marlin
        );
        
        if(config_m <= 512){
          auto it = kernel_maps_gemm1_decode_w4a8<bhalf_t>.find(mode);
          if (it != kernel_maps_gemm1_decode_w4a8<bhalf_t>.end()) {
              it->second(params_in);
          } else {
              printf("Path2 bfloat version decode gemm1 No matching kernel configuration found, using default settings mode %d size_m %d\n",mode,size_m);
              // launch_moe_w8a8_first_stage_prefill<16, 64, 128, 16, 16, 64>(params_in);
          }
        }else{ //decode 
          // std::cout<<"***************************************decode \n" << mode;
          
          auto it = kernel_maps_gemm1_prefill_w4a8<bhalf_t>.find(mode);
          // printf()
          if ( it != kernel_maps_gemm1_prefill_w4a8<bhalf_t>.end()) {
              it->second(params_in);
              
          } else {
              printf("Path2 bfloat version prefill gemm1 No matching kernel configuration found, using default settings mode %d size_m %d\n",mode,size_m);
              // launch_moe_w8a8_first_stage_decode<32, 256, 64, 32, 128, 64, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<64, 64, 128, 64, 64, 128, 1>(params_in);
              // launch_moe_w8a8_first_stage_decode<48, 16, 256, 48, 16, 128, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 32, 64, 16, 16, 64, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<32, 64, 128, 32, 64, 64, 1>(params_in);
          }
        }
        
        
      }else{ //gemm2
        int64_t EM = sorted_token_ids.size(0); // 一维线性化的token id
          
        // 使用int8类型处理
        GemmParams_w4a8<char,bhalf_t> params_in(
          (const char*)input.data_ptr<int8_t>(), 
          (const char*)b_qweight.data_ptr<int8_t>(), 
          (bhalf_t*)output_alias.data_ptr(),
          (float*)a_scale.data_ptr(), 
          (float*)b_scale.data_ptr(),  
          topk_weights_ptr,
          sorted_token_ids.data_ptr<int32_t>(),
          expert_ids.data_ptr<int32_t>(), 
          num_pad, //num_tokens_post_pad[0].item<int>(),
          num_tokens_post_pad.data_ptr<int32_t>(),
          size_m,
          stride_bse,
          size_k,
          stride_asm,
          stride_ask,
          stride_bse,
          stride_bsn,
          stride_bsk,
          EM,
          top_k,
          real_topk,
          is_marlin
        );

        if( config_m <= 512 ){
          auto it = kernel_maps_gemm2_decode_w4a8<bhalf_t>.find(mode);
          if (it != kernel_maps_gemm2_decode_w4a8<bhalf_t>.end()) {
              it->second(params_in);
          } else {
              printf("Path2 bfloat version decode gemm2 No matching kernel configuration found, using default settings mode %d size_m %d\n",mode,size_m);
              // launch_moe_w8a8_second_stage_prefill<16, 64, 128, 16, 16, 64>(params_in);
          }
        }else
        {
          // std::cout<<"mode: "<<mode<<std::endl;
          
          auto it = kernel_maps_gemm2_prefill_w4a8<bhalf_t>.find(mode);
          if (    it != kernel_maps_gemm2_prefill_w4a8<bhalf_t>.end()) {
            // printf("**********************************%d",mode);
              it->second(params_in);
          } else {

              printf("Path2 bfloat version prefill gemm2 No matching kernel configuration found, using default settings mode %d size_m %d\n",mode,size_m);
              // launch_moe_w8a8_second_stage_decode<16, 64, 128, 16, 16, 128, 2>(params_in); 
              // launch_moe_w8a8_second_stage_decode<32, 1024, 64, 16, 128, 64, 1>(params_in);
              // launch_moe_w8a8_second_stage_decode<32, 128, 64, 32, 32, 64, 2>(params_in);
              // launch_moe_w8a8_second_stage_decode<64, 64, 64, 64, 32, 64, 2>(params_in); 
              // launch_moe_w8a8_second_stage_decode<32, 256, 128, 32, 128, 128, 1>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 64, 128, 16, 16, 64, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 128, 128, 16, 16, 64, 2>(params_in);

          }

        }
              // hipDeviceSynchronize();

      }
    } else {
      // TORCH_CHECK(false, "moe_w8a8_gemm only supports int8");
    }
    
  }
  else if (output.scalar_type() == at::ScalarType::Half)
  {
    // printf("********************************************halfhalfhalf");
    if (input.scalar_type() == at::ScalarType::Char){
      if(first_stage){
     
        int64_t EM = sorted_token_ids.size(0); // 一维线性化的token id
        GemmParams_w4a8<char,half> params_in(
          (const char*)input.data_ptr<int8_t>(), 
          (const char*)b_qweight.data_ptr<int8_t>(), 
          (half*)output_alias.data_ptr(),
          (float*)a_scale.data_ptr(), 
          (float*)b_scale.data_ptr(),  
          topk_weights_ptr,
          sorted_token_ids.data_ptr<int32_t>(),
          expert_ids.data_ptr<int32_t>(), 
          num_pad, //num_tokens_post_pad[0].item<int>(), //这里获取值 会造成device->host的拷贝和一部分空泡
          num_tokens_post_pad.data_ptr<int32_t>(),
          size_m,
          stride_bse,
          size_k,
          stride_asm,
          stride_ask,
          stride_bse,
          stride_bsn,
          stride_bsk,
          EM,
          top_k,
          real_topk,
          is_marlin
        );

        if(false/* size_m <= 512 */){
          auto it = kernel_maps_gemm1_decode_w4a8<half>.find(mode);
          if (it != kernel_maps_gemm1_decode_w4a8<half>.end()) {
              it->second(params_in);
          } else {
              printf("Path2 half version decode gemm1 No matching kernel configuration found, using default settings mode %d size_m %d\n",mode,size_m);
              // launch_moe_w8a8_first_stage_prefill<16, 64, 128, 16, 16, 64>(params_in);
          }
        }else{ //decode 
          // std::cout<<"***************************************decode \n" << mode;
          auto it = kernel_maps_gemm1_prefill_w4a8<half>.find(mode);
          // printf()
          if ( it != kernel_maps_gemm1_prefill_w4a8<half>.end()) {
              it->second(params_in);
              
          } else {
              printf("Path2 half version prefill gemm1 No matching kernel configuration found, using default settings mode %d size_m %d\n",mode,size_m);
              // launch_moe_w8a8_first_stage_decode<16, 16, 512, 16, 16, 128, 4>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 32, 64, 16, 16, 64, 2>(params_in);

          }
        }
        
        
      }else{ //gemm2
        int64_t EM = sorted_token_ids.size(0); // 一维线性化的token id
          
        // 使用int8类型处理
        GemmParams_w4a8<char,half> params_in(
          (const char*)input.data_ptr<int8_t>(), 
          (const char*)b_qweight.data_ptr<int8_t>(), 
          (half*)output_alias.data_ptr(),
          (float*)a_scale.data_ptr(), 
          (float*)b_scale.data_ptr(),  
          topk_weights_ptr,
          sorted_token_ids.data_ptr<int32_t>(),
          expert_ids.data_ptr<int32_t>(), 
          num_pad, //num_tokens_post_pad[0].item<int>(),
          num_tokens_post_pad.data_ptr<int32_t>(),
          size_m,
          stride_bse,
          size_k,
          stride_asm,
          stride_ask,
          stride_bse,
          stride_bsn,
          stride_bsk,
          EM,
          top_k,
          real_topk,
          is_marlin
        );

        if(false/* size_m <= 512 * real_topk */){
          auto it = kernel_maps_gemm2_decode_w4a8<half>.find(mode);
          if (it != kernel_maps_gemm2_decode_w4a8<half>.end()) {
              it->second(params_in);
          } else {
              printf("Path2 half version decode gemm2 No matching kernel configuration found, using default settings mode %d size_m %d\n",mode,size_m);
              // launch_moe_w8a8_first_stage_prefill<16, 64, 128, 16, 16, 64>(params_in);
          }
        }else
        {
          // std::cout<<"*************************float16"<<std::endl;
          // mode =86;
          auto it = kernel_maps_gemm2_prefill_w4a8<half>.find(mode);
          if (  it != kernel_maps_gemm2_prefill_w4a8<half>.end()) {
            // printf("**********************************%d",mode);
              it->second(params_in);
          } else {

              printf("Path2 half version prefill gemm2 No matching kernel configuration found, using default settings mode %d size_m %d\n",mode,size_m);
              // launch_moe_w8a8_second_stage_decode<16, 64, 128, 16, 16, 128, 2>(params_in); 

              // launch_moe_w8a8_second_stage_decode<16, 256, 64, 16, 32, 64, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 64, 128, 16, 16, 64, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 128, 128, 16, 16, 64, 2>(params_in);

          }

        }
              // hipDeviceSynchronize();

      }
    } else {
      TORCH_CHECK(false, "moe_w8a8_gemm only supports int8");
    }
  }
}

    return output;
}

torch::Tensor moe_c_moe_gemm_marlin_w4a16(torch::Tensor input,
  torch::Tensor b_qweight,
  torch::Tensor output,
  torch::Tensor b_scale,
  torch::Tensor b_zeros,
  std::optional<torch::Tensor> topk_weights,
  torch::Tensor sorted_token_ids, 
  torch::Tensor expert_ids,
  torch::Tensor num_tokens_post_pad, 
  int64_t top_k, // gemm1为topk  gemm2为1  因为gemm1输入为[m, k]  gemm2输入为[m*topk, k]
  int64_t mode,
  int64_t delta
  ) {
    const int size_m = input.size(0); 
    const int EXPERTS = b_qweight.size(0);
    const int size_k = input.size(1);
    // std::cout<<"size_k"<<size_k<<std::endl;
    const int size_n = b_scale.size(1);
    const int topk_size = output.size(1); // 输出为[m, topk_size, n]
    // const int stride_asm = a_scale.stride(0);
    // const int stride_ask = a_scale.stride(1);
    // const int stride_bse = b_scale.size(1);
    // const int stride_bsn = b_scale.size(2); 
    // const int stride_bsk = b_scale.size(2);
    // const uint32_t* b_zeros_ptr;
    // if (b_zeros.has_value())
    // b_zeros_ptr = (const uint32_t*)b_zeros.value().data_ptr<uint8_t>();
    // 单行printf打印所有步长变量，带标签便于识别
    // printf("stride_asm: %d, stride_ask: %d, stride_bse: %d, stride_bsn: %d, stride_bsk: %d\n", 
    //    stride_asm, stride_ask, stride_bse, stride_bsn, stride_bsk);
    constexpr int GROUP_N = 1;
    constexpr int GROUP_K = 1;
    bool is_marlin = true; // weight为[E, N, K]时 代表不进行重排
    bool first_stage = true;
    torch::Tensor output_alias = output.alias();

    //printf("size_n: %d stride_bse: %d stride_bsn: %d stride_bsk: %d\n", size_n, stride_bse, stride_bsn, stride_bsk);

    const float* topk_weights_ptr; // 第一阶段这里为null
    if (topk_weights.has_value()){
      topk_weights_ptr = (const float*)topk_weights.value().data_ptr();
      first_stage = false;
    }

    int num_pad = 0;

    if (input.scalar_type() == at::ScalarType::Half){
      if(first_stage){
     
        int64_t EM = sorted_token_ids.size(0); // 一维线性化的token id
        GemmParams_w4a16<half> params_in(
          (half*)input.data_ptr<at::Half>(), 
          (uint32_t*)b_qweight.data_ptr<uint32_t>(), 
          (half*)output_alias.data_ptr<at::Half>(),
          reinterpret_cast<uint32_t*>(b_zeros.data_ptr<uint8_t>()), 
          (half*)b_scale.data_ptr<at::Half>(),  
          topk_weights_ptr,
          sorted_token_ids.data_ptr<int32_t>(),
          expert_ids.data_ptr<int32_t>(), 
          num_pad, //num_tokens_post_pad[0].item<int>(), //这里获取值 会造成device->host的拷贝和一部分空泡
          num_tokens_post_pad.data_ptr<int32_t>(),
          size_m,
          size_n,
          size_k,
          // stride_asm,
          // stride_ask,
          // stride_bse,
          // stride_bsn,
          // stride_bsk,
          EM,
          top_k,
          delta,
          is_marlin
        );

        if(mode >= 500){
          // auto it = kernel_maps_gemm1_prefill.find(mode);
          // if (it != kernel_maps_gemm1_prefill.end()) {
          //     it->second(params_in);
          // } else {
          //     printf("gemm1 No matching kernel configuration found, using default settings \n");
          //     // launch_moe_w8a8_first_stage_prefill<16, 64, 128, 16, 16, 64>(params_in);
          // }
        }else{ //decode 
          // std::cout<<"***************************************decode \n" << mode;
          auto it = kernel_maps_gemm1_decode_w4a16<half>.find(mode);
          // printf()
          if (it != kernel_maps_gemm1_decode_w4a16<half>.end() ) {
            float milliseconds = 0;
            cudaEvent_t start, stop;
            const char* find_best = std::getenv("WHICH_TO_TEST");
            if (find_best) {
              cudaEventCreate(&start);
              cudaEventCreate(&stop);
              cudaEventRecord(start);        // 记录开始
            }

            it->second(params_in);

            if (find_best) {
              cudaEventRecord(stop);         // 记录结束
              cudaEventSynchronize(stop);    // 等待 kernel 执行完成

              
              cudaEventElapsedTime(&milliseconds, start, stop); // 计算时间

              /* std::cout << "kernel 1 time: " << milliseconds << " ms" << std::endl; */
              
              cudaEventDestroy(start);
              cudaEventDestroy(stop);

              std::ofstream ofs("./w4a16_kernel_1_timecost", std::ios::app); // 追加写入
              if (ofs.is_open()) {
                  ofs << milliseconds << std::endl;
                  ofs.close();
              }
            }
              
          } else {
              // printf("gemm1 No matching kernel configuration found, using default settings \n");
              // launch_moe_w8a8_first_stage_decode<16, 64, 256, 16, 32, 128, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 32, 64, 16, 16, 64, 2>(params_in);
          }
        }
        
        
      }else{ //gemm2
        int64_t EM = sorted_token_ids.size(0); // 一维线性化的token id
        GemmParams_w4a16<half> params_in(
          (half*)input.data_ptr<at::Half>(), 
          (uint32_t*)b_qweight.data_ptr<uint32_t>(), 
          (half*)output_alias.data_ptr<at::Half>(),
          reinterpret_cast<uint32_t*>(b_zeros.data_ptr<uint8_t>()), 
          (half*)b_scale.data_ptr<at::Half>(),  
          topk_weights_ptr,
          sorted_token_ids.data_ptr<int32_t>(),
          expert_ids.data_ptr<int32_t>(), 
          num_pad, //num_tokens_post_pad[0].item<int>(), //这里获取值 会造成device->host的拷贝和一部分空泡
          num_tokens_post_pad.data_ptr<int32_t>(),
          size_m,
          size_n,
          size_k,
          // stride_asm,
          // stride_ask,
          // stride_bse,
          // stride_bsn,
          // stride_bsk,
          EM,
          top_k,
          delta,
          is_marlin
        );

        if(mode >= 500){
        //   auto it = kernel_maps_gemm2_prefill.find(mode);
        //   if (it != kernel_maps_gemm2_prefill.end()) {
            

        //     // launch_moe_w8a8_second_stage_decode<16, 256, 64, 16, 32, 64, 2>(params_in);
        //       it->second(params_in);
        //   } else {
        //       printf("gemm2 No matching kernel configuration found, using default settings \n");
        //       launch_moe_w8a8_second_stage_prefill<16, 64, 128, 16, 16, 64>(params_in);
        //   }
        }else
        {
        //   // std::cout<<"mode: "<<mode<<std::endl;
          
          auto it = kernel_maps_gemm2_decode_w4a16<half>.find(mode);
          if (it != kernel_maps_gemm2_decode_w4a16<half>.end() ) {
            float milliseconds = 0;
            cudaEvent_t start, stop;
            const char* find_best = std::getenv("WHICH_TO_TEST");
            if (find_best) {
              cudaEventCreate(&start);
              cudaEventCreate(&stop);
              cudaEventRecord(start);        // 记录开始
            }

            it->second(params_in);

            if (find_best) {
            cudaEventRecord(stop);         // 记录结束
            cudaEventSynchronize(stop);    // 等待 kernel 执行完成

            
            cudaEventElapsedTime(&milliseconds, start, stop); // 计算时间

            /* std::cout << "kernel 1 time: " << milliseconds << " ms" << std::endl; */
            
            cudaEventDestroy(start);
            cudaEventDestroy(stop);

            std::ofstream ofs("./w4a16_kernel_2_timecost", std::ios::app); // 追加写入
            if (ofs.is_open()) {
                ofs << milliseconds << std::endl;
                ofs.close();
              }
            }
          } else {

        //       // printf("gemm2 No matching kernel configuration found, using default settings \n");
        //       // launch_moe_w8a8_second_stage_decode<16, 64, 128, 16, 16, 128, 2>(params_in); 
              // launch_moe_w8a8_second_stage_decode<16, 256, 128, 16, 64, 128, 2>(params_in);
              // launch_moe_w8a8_second_stage_decode<16, 256, 64, 16, 32, 64, 2>(params_in);
        //       // launch_moe_w8a8_first_stage_decode<16, 64, 128, 16, 16, 64, 2>(params_in);
        //       // launch_moe_w8a8_first_stage_decode<16, 128, 128, 16, 16, 64, 2>(params_in);

          }

        }
              // hipDeviceSynchronize();

      }
    } else if (input.scalar_type() == at::ScalarType::BFloat16){
      if(first_stage){
     
        int64_t EM = sorted_token_ids.size(0); // 一维线性化的token id
        GemmParams_w4a16<__hip_bfloat16> params_in(
          (__hip_bfloat16*)input.data_ptr<at::BFloat16>(), 
          (uint32_t*)b_qweight.data_ptr<uint32_t>(), 
          (__hip_bfloat16*)output_alias.data_ptr<at::BFloat16>(),
          reinterpret_cast<uint32_t*>(b_zeros.data_ptr<uint8_t>()), 
          (__hip_bfloat16*)b_scale.data_ptr<at::BFloat16>(),  
          topk_weights_ptr,
          sorted_token_ids.data_ptr<int32_t>(),
          expert_ids.data_ptr<int32_t>(), 
          num_pad, //num_tokens_post_pad[0].item<int>(), //这里获取值 会造成device->host的拷贝和一部分空泡
          num_tokens_post_pad.data_ptr<int32_t>(),
          size_m,
          size_n,
          size_k,
          // stride_asm,
          // stride_ask,
          // stride_bse,
          // stride_bsn,
          // stride_bsk,
          EM,
          top_k,
          delta,
          is_marlin
        );

        if(mode >= 500){
          // auto it = kernel_maps_gemm1_prefill.find(mode);
          // if (it != kernel_maps_gemm1_prefill.end()) {
          //     it->second(params_in);
          // } else {
          //     printf("gemm1 No matching kernel configuration found, using default settings \n");
          //     // launch_moe_w8a8_first_stage_prefill<16, 64, 128, 16, 16, 64>(params_in);
          // }
        }else{ //decode 
          // std::cout<<"***************************************decode \n" << mode;
          auto it = kernel_maps_gemm1_decode_w4a16<__hip_bfloat16>.find(mode);
          // printf()
          if (it != kernel_maps_gemm1_decode_w4a16<__hip_bfloat16>.end() ) {
            float milliseconds = 0;
            cudaEvent_t start, stop;
            const char* find_best = std::getenv("WHICH_TO_TEST");
            if (find_best) {
              cudaEventCreate(&start);
              cudaEventCreate(&stop);
              cudaEventRecord(start);        // 记录开始
            }

            it->second(params_in);

            if (find_best) {
              cudaEventRecord(stop);         // 记录结束
              cudaEventSynchronize(stop);    // 等待 kernel 执行完成

              
              cudaEventElapsedTime(&milliseconds, start, stop); // 计算时间

              /* std::cout << "kernel 1 time: " << milliseconds << " ms" << std::endl; */
              
              cudaEventDestroy(start);
              cudaEventDestroy(stop);

              std::ofstream ofs("./w4a16_kernel_1_timecost", std::ios::app); // 追加写入
              if (ofs.is_open()) {
                  ofs << milliseconds << std::endl;
                  ofs.close();
              }
            }
              
          } else {
              // printf("gemm1 No matching kernel configuration found, using default settings \n");
              // launch_moe_w8a8_first_stage_decode<16, 64, 256, 16, 32, 128, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 32, 64, 16, 16, 64, 2>(params_in);
          }
        }
        
        
      }else{ //gemm2
        int64_t EM = sorted_token_ids.size(0); // 一维线性化的token id
        GemmParams_w4a16<__hip_bfloat16> params_in(
          (__hip_bfloat16*)input.data_ptr<at::BFloat16>(), 
          (uint32_t*)b_qweight.data_ptr<uint32_t>(), 
          (__hip_bfloat16*)output_alias.data_ptr<at::BFloat16>(),
          reinterpret_cast<uint32_t*>(b_zeros.data_ptr<uint8_t>()), 
          (__hip_bfloat16*)b_scale.data_ptr<at::BFloat16>(),  
          topk_weights_ptr,
          sorted_token_ids.data_ptr<int32_t>(),
          expert_ids.data_ptr<int32_t>(), 
          num_pad, //num_tokens_post_pad[0].item<int>(), //这里获取值 会造成device->host的拷贝和一部分空泡
          num_tokens_post_pad.data_ptr<int32_t>(),
          size_m,
          size_n,
          size_k,
          // stride_asm,
          // stride_ask,
          // stride_bse,
          // stride_bsn,
          // stride_bsk,
          EM,
          top_k,
          delta,
          is_marlin
        );

        if(mode >= 500){
        //   auto it = kernel_maps_gemm2_prefill.find(mode);
        //   if (it != kernel_maps_gemm2_prefill.end()) {
            

        //     // launch_moe_w8a8_second_stage_decode<16, 256, 64, 16, 32, 64, 2>(params_in);
        //       it->second(params_in);
        //   } else {
        //       printf("gemm2 No matching kernel configuration found, using default settings \n");
        //       launch_moe_w8a8_second_stage_prefill<16, 64, 128, 16, 16, 64>(params_in);
        //   }
        }else
        {
        //   // std::cout<<"mode: "<<mode<<std::endl;
          
          auto it = kernel_maps_gemm2_decode_w4a16<__hip_bfloat16>.find(mode);
          if (it != kernel_maps_gemm2_decode_w4a16<__hip_bfloat16>.end() ) {
            float milliseconds = 0;
            cudaEvent_t start, stop;
            const char* find_best = std::getenv("WHICH_TO_TEST");
            if (find_best) {
              cudaEventCreate(&start);
              cudaEventCreate(&stop);
              cudaEventRecord(start);        // 记录开始
            }

            it->second(params_in);

            if (find_best) {
            cudaEventRecord(stop);         // 记录结束
            cudaEventSynchronize(stop);    // 等待 kernel 执行完成

            
            cudaEventElapsedTime(&milliseconds, start, stop); // 计算时间

            /* std::cout << "kernel 1 time: " << milliseconds << " ms" << std::endl; */
            
            cudaEventDestroy(start);
            cudaEventDestroy(stop);

            std::ofstream ofs("./w4a16_kernel_2_timecost", std::ios::app); // 追加写入
            if (ofs.is_open()) {
                ofs << milliseconds << std::endl;
                ofs.close();
              }
            }
          } else {

        //       // printf("gemm2 No matching kernel configuration found, using default settings \n");
        //       // launch_moe_w8a8_second_stage_decode<16, 64, 128, 16, 16, 128, 2>(params_in); 
              // launch_moe_w8a8_second_stage_decode<16, 256, 128, 16, 64, 128, 2>(params_in);
              // launch_moe_w8a8_second_stage_decode<16, 256, 64, 16, 32, 64, 2>(params_in);
        //       // launch_moe_w8a8_first_stage_decode<16, 64, 128, 16, 16, 64, 2>(params_in);
        //       // launch_moe_w8a8_first_stage_decode<16, 128, 128, 16, 16, 64, 2>(params_in);

          }

        }
              // hipDeviceSynchronize();

      }
    } else {
      TORCH_CHECK(false, "moe_w8a8_gemm only supports int8");
    }

    return output;
}

torch::Tensor moe_c_moe_gemm_marlin_w8a16(torch::Tensor input,
  torch::Tensor b_qweight,
  torch::Tensor output,
  torch::Tensor b_scale,
  std::optional<torch::Tensor> topk_weights,
  torch::Tensor sorted_token_ids, 
  torch::Tensor expert_ids,
  torch::Tensor num_tokens_post_pad, 
  int64_t top_k, // gemm1为topk  gemm2为1  因为gemm1输入为[m, k]  gemm2输入为[m*topk, k]
  int64_t mode,
  int64_t delta
  ) {
    const int size_m = input.size(0); 
    const int EXPERTS = b_qweight.size(0);
    const int size_k = input.size(1);
    // std::cout<<"size_k"<<size_k<<std::endl;
    const int size_n = b_scale.size(1); //modifyy
    const int topk_size = output.size(1); // 输出为[m, topk_size, n]
    // const int stride_asm = a_scale.stride(0);
    // const int stride_ask = a_scale.stride(1);
    // const int stride_bse = b_scale.size(1);
    // const int stride_bsn = b_scale.size(2); 
    // const int stride_bsk = b_scale.size(2);
    // const uint32_t* b_zeros_ptr;
    // if (b_zeros.has_value())
    // b_zeros_ptr = (const uint32_t*)b_zeros.value().data_ptr<uint8_t>();
    // 单行printf打印所有步长变量，带标签便于识别
    // printf("stride_asm: %d, stride_ask: %d, stride_bse: %d, stride_bsn: %d, stride_bsk: %d\n", 
    //    stride_asm, stride_ask, stride_bse, stride_bsn, stride_bsk);
    constexpr int GROUP_N = 1;
    constexpr int GROUP_K = 1;
    bool is_marlin = true; // weight为[E, N, K]时 代表不进行重排
    bool first_stage = true;
    torch::Tensor output_alias = output.alias();

    //printf("size_n: %d stride_bse: %d stride_bsn: %d stride_bsk: %d\n", size_n, stride_bse, stride_bsn, stride_bsk);

    const float* topk_weights_ptr = nullptr; // 第一阶段这里为null
    if (topk_weights.has_value()){
      topk_weights_ptr = (const float*)topk_weights.value().data_ptr();
      first_stage = false;
    }

    int num_pad = 0;


  // #if (DEBUG_W8A8_PERCHANNEL)
    if (input.scalar_type() == at::ScalarType::Half){
      if(first_stage){
     
        int64_t EM = sorted_token_ids.size(0); // 一维线性化的token id
        GemmParams_w8a16<half> params_in(
          (half*)input.data_ptr<at::Half>(), 
          (uint32_t*)b_qweight.data_ptr<uint32_t>(), 
          (half*)output_alias.data_ptr<at::Half>(),
          (half*)b_scale.data_ptr<at::Half>(),  
          topk_weights_ptr,
          sorted_token_ids.data_ptr<int32_t>(),
          expert_ids.data_ptr<int32_t>(), 
          num_pad, //num_tokens_post_pad[0].item<int>(), //这里获取值 会造成device->host的拷贝和一部分空泡
          num_tokens_post_pad.data_ptr<int32_t>(),
          size_m,
          size_n,
          size_k,
          // stride_asm,
          // stride_ask,
          // stride_bse,
          // stride_bsn,
          // stride_bsk,
          EM,
          top_k,
          delta,
          is_marlin
        );

        if(mode >= 500){
          // auto it = kernel_maps_gemm1_prefill.find(mode);
          // if (it != kernel_maps_gemm1_prefill.end()) {
          //     it->second(params_in);
          // } else {
          //     printf("gemm1 No matching kernel configuration found, using default settings \n");
          //     // launch_moe_w8a8_first_stage_prefill<16, 64, 128, 16, 16, 64>(params_in);
          // }
        }else{ //decode 
          // std::cout<<"***************************************decode \n" << mode;
          auto it = kernel_maps_gemm1_decode_w8a16<half>.find(mode);
          // printf()
          if (it != kernel_maps_gemm1_decode_w8a16<half>.end() ) {
            float milliseconds = 0;
            cudaEvent_t start, stop;
            const char* find_best = std::getenv("WHICH_TO_TEST");
            if (find_best) {
              cudaEventCreate(&start);
              cudaEventCreate(&stop);
              cudaEventRecord(start);        // 记录开始
            }

            it->second(params_in);

            if (find_best) {
              cudaEventRecord(stop);         // 记录结束
              cudaEventSynchronize(stop);    // 等待 kernel 执行完成

              
              cudaEventElapsedTime(&milliseconds, start, stop); // 计算时间

              // std::cout << "8888kernel 1 time-----------------: " << milliseconds << " ms" << std::endl;
              
              cudaEventDestroy(start);
              cudaEventDestroy(stop);

              std::ofstream ofs("./w8a16_kernel_1_timecost", std::ios::app); // 追加写入
              if (ofs.is_open()) {
                  ofs << milliseconds << std::endl;
                  ofs.close();
              }
            }
              
          } else {
              // printf("gemm1 No matching kernel configuration found, using default settings \n");
              // launch_moe_w8a8_first_stage_decode<16, 64, 256, 16, 32, 128, 2>(params_in);
              // launch_moe_w8a8_first_stage_decode<16, 32, 64, 16, 16, 64, 2>(params_in);
          }
        }
        
        
      }else{ //gemm2
        int64_t EM = sorted_token_ids.size(0); // 一维线性化的token id
        GemmParams_w8a16<half> params_in(
          (half*)input.data_ptr<at::Half>(), 
          (uint32_t*)b_qweight.data_ptr<uint32_t>(), 
          (half*)output_alias.data_ptr<at::Half>(),
          (half*)b_scale.data_ptr<at::Half>(),  
          topk_weights_ptr,
          sorted_token_ids.data_ptr<int32_t>(),
          expert_ids.data_ptr<int32_t>(), 
          num_pad, //num_tokens_post_pad[0].item<int>(), //这里获取值 会造成device->host的拷贝和一部分空泡
          num_tokens_post_pad.data_ptr<int32_t>(),
          size_m,
          size_n,
          size_k,
          // stride_asm,
          // stride_ask,
          // stride_bse,
          // stride_bsn,
          // stride_bsk,
          EM,
          top_k,
          delta,
          is_marlin
        );

        if(mode >= 500){
        //   auto it = kernel_maps_gemm2_prefill.find(mode);
        //   if (it != kernel_maps_gemm2_prefill.end()) {
            

        //     // launch_moe_w8a8_second_stage_decode<16, 256, 64, 16, 32, 64, 2>(params_in);
        //       it->second(params_in);
        //   } else {
        //       printf("gemm2 No matching kernel configuration found, using default settings \n");
        //       launch_moe_w8a8_second_stage_prefill<16, 64, 128, 16, 16, 64>(params_in);
        //   }
        }else
        {
        //   // std::cout<<"mode: "<<mode<<std::endl;
          
          auto it = kernel_maps_gemm2_decode_w8a16<half>.find(mode);
          if (it != kernel_maps_gemm2_decode_w8a16<half>.end() ) {
            float milliseconds = 0;
            cudaEvent_t start, stop;
            const char* find_best = std::getenv("WHICH_TO_TEST");
            if (find_best) {
              cudaEventCreate(&start);
              cudaEventCreate(&stop);
              cudaEventRecord(start);        // 记录开始
            }

            it->second(params_in);

            if (find_best) {
            cudaEventRecord(stop);         // 记录结束
            cudaEventSynchronize(stop);    // 等待 kernel 执行完成

            
            cudaEventElapsedTime(&milliseconds, start, stop); // 计算时间

            // std::cout << "------------kernel 1 time g2-----------------: " << milliseconds << " ms" << std::endl;
            
            cudaEventDestroy(start);
            cudaEventDestroy(stop);

            std::ofstream ofs("./w8a16_kernel_2_timecost", std::ios::app); // 追加写入
            if (ofs.is_open()) {
                ofs << milliseconds << std::endl;
                ofs.close();
              }
            }
          } else {

        //       // printf("gemm2 No matching kernel configuration found, using default settings \n");
        //       // launch_moe_w8a8_second_stage_decode<16, 64, 128, 16, 16, 128, 2>(params_in); 
              // launch_moe_w8a8_second_stage_decode<16, 256, 128, 16, 64, 128, 2>(params_in);
              // launch_moe_w8a8_second_stage_decode<16, 256, 64, 16, 32, 64, 2>(params_in);
        //       // launch_moe_w8a8_first_stage_decode<16, 64, 128, 16, 16, 64, 2>(params_in);
        //       // launch_moe_w8a8_first_stage_decode<16, 128, 128, 16, 16, 64, 2>(params_in);

          }

        }
              // hipDeviceSynchronize();

      }
    } else if (input.scalar_type() == at::ScalarType::BFloat16){
      int64_t EM = sorted_token_ids.size(0); // 一维线性化的token id
      GemmParams_w8a16<__hip_bfloat16> params_in(
        (__hip_bfloat16*)input.data_ptr<at::BFloat16>(),
        (uint32_t*)b_qweight.data_ptr<uint32_t>(),
        (__hip_bfloat16*)output_alias.data_ptr<at::BFloat16>(),
        (__hip_bfloat16*)b_scale.data_ptr<at::BFloat16>(),
        topk_weights_ptr,
        sorted_token_ids.data_ptr<int32_t>(),
        expert_ids.data_ptr<int32_t>(),
        num_pad, //num_tokens_post_pad[0].item<int>(), //这里获取值 会造成device->host的拷贝和一部分空泡
        num_tokens_post_pad.data_ptr<int32_t>(),
        size_m,
        size_n,
        size_k,
        EM,
        top_k,
        delta,
        is_marlin
      );

      if(mode < 500){
        if(first_stage){
          auto it = kernel_maps_gemm1_decode_w8a16<__hip_bfloat16>.find(mode);
          if (it != kernel_maps_gemm1_decode_w8a16<__hip_bfloat16>.end() ) {
            it->second(params_in);
          }
        } else {
          auto it = kernel_maps_gemm2_decode_w8a16<__hip_bfloat16>.find(mode);
          if (it != kernel_maps_gemm2_decode_w8a16<__hip_bfloat16>.end() ) {
            it->second(params_in);
          }
        }
      }
    } else {
      TORCH_CHECK(false, "moe_w8a16_gemm only supports float16 and bfloat16 input");
    }
// #endif
    return output;
}

torch::Tensor moe_c_moe_w8a8_gemm_block_wise(torch::Tensor input, torch::Tensor a_scales,torch::Tensor output,
                             torch::Tensor b_qweight, torch::Tensor b_scales,
                             std::optional<torch::Tensor> b_qzeros,
                             std::optional<torch::Tensor> topk_weights,
                             torch::Tensor sorted_token_ids,
                             torch::Tensor expert_ids,
                             torch::Tensor num_tokens_post_pad, int64_t group_size_n, int64_t group_size_k, int64_t top_k,
                             int64_t BLOCK_SIZE_m, int64_t BLOCK_SIZE_n,
                             int64_t BLOCK_SIZE_k, int64_t kloops, int64_t nloops, int64_t bit ) {
  
  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
  const int size_m = input.size(0);
  const int size_n = b_qweight.size(1);
  const int size_k = input.size(1);
  
  int64_t BLOCK_SIZE_N = 64;
  // std::cout<<"BLOCK_SIZE_N IS : "<<BLOCK_SIZE_N<<std::endl;
  int64_t BLOCK_SIZE_K_MIN =4*1024/BLOCK_SIZE_m;
  int64_t BLOCK_SIZE_K_MAX =8*1024/BLOCK_SIZE_m;
  // std::cout<<"BLOCK_SIZE_K_MAX IS : "<<BLOCK_SIZE_K_MAX<<std::endl;
  int64_t BLOCK_SIZE_K = std::min(BLOCK_SIZE_K_MAX,size_k);
  // std::cout<<"BLOCK_SIZE_K IS : "<<BLOCK_SIZE_K<<std::endl;

  int BLOCK_SIZE_M_MAX = std::min(16, size_m);
  // std::cout<<"BLOCK_SIZE_M_MAX IS : "<<BLOCK_SIZE_M_MAX<<std::endl;
  int BLOCK_SIZE_N_MAX_roofline = 64;
  int BLOCK_SIZE_N_MAX = std::min(BLOCK_SIZE_N_MAX_roofline, size_n);
  // std::cout<<"BLOCK_SIZE_N_MAX IS : "<<BLOCK_SIZE_N_MAX<<std::endl;

  // BLOCK_SIZE_K = std::min(128, BLOCK_SIZE_K);
  BLOCK_SIZE_K = 128;
  // std::cout<<"BLOCK_SIZE_K IS : "<<BLOCK_SIZE_K<<std::endl;

  // int block_size_m_loops = 1;// std::min(1,BLOCK_SIZE_M_MAX/BLOCK_SIZE_m);
  int block_size_m_loops = 1;
  int block_size_k_loops = kloops;
  // int block_size_n_loops = BLOCK_SIZE_N_MAX/BLOCK_SIZE_N;
  int block_size_n_loops = nloops;

  // {//for debug
  // int* dev_d_w = nullptr;
  // hipMalloc((void**)&dev_d_w, 16*64*sizeof(int));
  // }//for debug
  int d_w_out[16*64]; 
  
  int64_t EM = sorted_token_ids.size(0);
  if (size_m <= BLOCK_SIZE_m) {
    EM = min(EM, size_m * BLOCK_SIZE_m * top_k);
  }
  const int num_token_blocks = (EM + BLOCK_SIZE_m*block_size_m_loops - 1) / (BLOCK_SIZE_m*block_size_m_loops);

  const uint32_t* b_qzeros_ptr;
  if (b_qzeros.has_value())
    b_qzeros_ptr = (const uint32_t*)b_qzeros.value().data_ptr<uint8_t>();
  const float* topk_weights_ptr;
  if (topk_weights.has_value())
    topk_weights_ptr = (const float*)topk_weights.value().data_ptr();

  int groups_per_block_row = BLOCK_SIZE_K / group_size_k;
  TORCH_CHECK(bit == 4 || bit == 8, "bit must be 4 or 8");
  TORCH_CHECK(size_k % BLOCK_SIZE_K == 0,
              "size_k must divisible by BLOCK_SIZE_K");
  TORCH_CHECK(BLOCK_SIZE_K % group_size_k == 0,
              "BLOCK_SIZE_K must divisible by group_size_k");
  TORCH_CHECK(BLOCK_SIZE_m <= 64, "BLOCK_SIZE_m must less or equal to 64");
  TORCH_CHECK(groups_per_block_row == 1 || groups_per_block_row == 2 ||
                  groups_per_block_row == 4 || groups_per_block_row == 8,
              "BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8]");
  
  bool use_atomic = (size_k != BLOCK_SIZE_K*block_size_k_loops);

  


  std::optional<torch::Tensor> output_fp32;

  if (use_atomic){

    output_fp32 = torch::zeros(output.sizes(),output.options().dtype(torch::kFloat32));
    // output_fp32->zero_(); 
  } 
  float milliseconds = 0;
  cudaEvent_t start, stop;
  const char* find_best = std::getenv("WHICH_TO_TEST");
  if (find_best) {
    
    cudaEventCreate(&start);
    cudaEventCreate(&stop);
    cudaEventRecord(start);        // 记录开始
  }
  // if (true/* input.scalar_type() == at::ScalarType::QInt8 */) {

  //   BIT_SWITCH(bit, BIT, [&]{
  //     TOPK_SWITCH(top_k, TOPK, [&]{
  //       BLOCK_M_SWITCH(BLOCK_SIZE_m, BLOCK_SIZE_M_, [&]{
  //         BLOCK_N_SWITCH(BLOCK_SIZE_N, BLOCK_SIZE_N_, [&]{
  //           BLOCK_K_SWITCH(BLOCK_SIZE_K, BLOCK_SIZE_K_, [&]{
  //             // BOOL_SWITCH(b_qzeros.has_value(), has_zp, [&]{
  //               BOOL_SWITCH(topk_weights.has_value(), mul_topk_weight, [&]{
  //                 GROUP_SIZE_N_SWITCH(group_size_n, GROUP_SIZE_N, [&]{
  //                   GROUP_SIZE_K_SWITCH(group_size_k, GROUP_SIZE_K, [&]{
  //                     BLOCK_SIZE_M_LOOPS_SWITCH(block_size_m_loops , BLOCK_SIZE_M_LOOPS, [&]{
  //                     BLOCK_SIZE_N_LOOPS_SWITCH(block_size_n_loops , BLOCK_SIZE_N_LOOPS, [&]{
  //                       BLOCK_SIZE_K_LOOPS_SWITCH(block_size_k_loops , BLOCK_SIZE_K_LOOPS, [&]{
  //                         BOOL_SWITCH(use_atomic , USE_ATOMIC, [&]{
  //                   run_moe_w8a8_gemm_block_wise<half, 8, TOPK, BLOCK_SIZE_M_, BLOCK_SIZE_N_, BLOCK_SIZE_K_, false, mul_topk_weight, GROUP_SIZE_N, GROUP_SIZE_K, BLOCK_SIZE_M_LOOPS, BLOCK_SIZE_N_LOOPS, BLOCK_SIZE_K_LOOPS, USE_ATOMIC,256>(
                    
  //                     (const uint32_t*)input.data_ptr<int8_t>(),
  //                     // (const half*)d_input,
  //                     (const float*)a_scales.data_ptr(),
  //                     use_atomic ?(float*)output_fp32->data_ptr():(float*)output.data_ptr(),
  //                     // (float*)output_fp32->data_ptr(),
  //                     (const uint32_t*)b_qweight.data_ptr<int8_t>(),
  //                     // (const uint32_t*)d_w_test,
  //                     ( int*)&d_w_out[0], /*for debug*/ /*使用时需修改为device端地址，这里仅为占位使用*/
  //                     // (const half*)b_scales.data_ptr<at::Half>(), 
  //                     (const float*)b_scales.data_ptr(),
  //                     // (const half*)d_scale, 
  //                     b_qzeros_ptr,
  //                     // (const uint32_t*)d_scale,
  //                     topk_weights_ptr, 
  //                     sorted_token_ids.data_ptr<int32_t>(),
  //                     expert_ids.data_ptr<int32_t>(), 
  //                     num_tokens_post_pad.data_ptr<int32_t>(), 
  //                     // num_tokens_post_pad_value<int32_t>(),
  //                     // num_tokens_post_pad_data_ptr[0],
  //                     num_token_blocks, 
  //                     size_m, 
  //                     size_n,
  //                     size_k
  //                     );
  //                     // printf("run_moe_w8a8_gemm_block_wise\n"); // kernel-1 mma    
  //                     });              
  //                     });
  //                   });
  //                 });
  //               });
  //             });
  //           });
  //         });
  //       });
  //     });
  //   });
  // });
  // } else {
  //   TORCH_CHECK(false, "moe_w8a8_gemm_block_wise only supports int8_t");
  // }

  if (find_best) {
  cudaEventRecord(stop);         // 记录结束
  cudaEventSynchronize(stop);    // 等待 kernel 执行完成

  
  cudaEventElapsedTime(&milliseconds, start, stop); // 计算时间

  /* std::cout << "kernel 1 time: " << milliseconds << " ms" << std::endl; */
  
  

  std::ofstream ofs("./w8a8_kernel_1_timecost", std::ios::app); // 追加写入
  if (ofs.is_open()) {
      ofs << milliseconds << std::endl;
      ofs.close();
  }
}
  
  

  if (use_atomic){
    // std::cout<<" fp32 convert to fp16 "<<std::endl;
    output.copy_(output_fp32->to(torch::kFloat16));
  }
  cudaEventDestroy(start);
  cudaEventDestroy(stop);
  
  // {//for debug
  // hipDeviceSynchronize();
  // hipMemcpy(&d_w_out[0], dev_d_w, 16*64 * sizeof(int), hipMemcpyDeviceToHost);
  // for(int i =0;i<16;i++){
  //   for(int j = 0;j<64 ;j++){
  //     std::cout<<(int)(d_w_out[i*64+j])<<"\t";
  //   }
  //   std::cout<<"|||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||"<<std::endl;
  // }
  // hipFree(dev_d_w);
  // }
  return output;
  
}

torch::Tensor moe_c_moe_w8a8_gemm_block_wise_kernel2(torch::Tensor input, torch::Tensor a_scales,torch::Tensor output,
                             torch::Tensor b_qweight, torch::Tensor b_scales,
                             std::optional<torch::Tensor> b_qzeros,
                             std::optional<torch::Tensor> topk_weights,
                             torch::Tensor sorted_token_ids,
                             torch::Tensor expert_ids,
                             torch::Tensor num_tokens_post_pad, int64_t group_size_n, int64_t group_size_k, int64_t top_k,
                             int64_t BLOCK_SIZE_m, int64_t BLOCK_SIZE_n,
                             int64_t BLOCK_SIZE_k, int64_t kloops, int64_t nloops, int64_t bit ) {
  
  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));

  const int size_m = input.size(0);
  const int size_n = b_qweight.size(1);
  const int size_k = input.size(1);
  
  int64_t BLOCK_SIZE_N = 64;
  // std::cout<<"BLOCK_SIZE_N IS : "<<BLOCK_SIZE_N<<std::endl;
  int64_t BLOCK_SIZE_K_MIN =4*1024/BLOCK_SIZE_m;
  int64_t BLOCK_SIZE_K_MAX =8*1024/BLOCK_SIZE_m;
  // std::cout<<"BLOCK_SIZE_K_MAX IS : "<<BLOCK_SIZE_K_MAX<<std::endl;
  int64_t BLOCK_SIZE_K = std::min(BLOCK_SIZE_K_MAX,size_k);
  // std::cout<<"BLOCK_SIZE_K IS : "<<BLOCK_SIZE_K<<std::endl;

  int BLOCK_SIZE_M_MAX = std::min(16, size_m);
  // std::cout<<"BLOCK_SIZE_M_MAX IS : "<<BLOCK_SIZE_M_MAX<<std::endl;
  int BLOCK_SIZE_N_MAX_roofline = 64;
  int BLOCK_SIZE_N_MAX = std::min(BLOCK_SIZE_N_MAX_roofline, size_n);
  // std::cout<<"BLOCK_SIZE_N_MAX IS : "<<BLOCK_SIZE_N_MAX<<std::endl;

  // BLOCK_SIZE_K = std::min(128, BLOCK_SIZE_K);
  BLOCK_SIZE_K = 128;
  // std::cout<<"BLOCK_SIZE_K IS : "<<BLOCK_SIZE_K<<std::endl;

  // int block_size_m_loops = 1;// std::min(1,BLOCK_SIZE_M_MAX/BLOCK_SIZE_m);
  int block_size_m_loops = 1;
  int block_size_k_loops =kloops;
  // int block_size_n_loops = BLOCK_SIZE_N_MAX/BLOCK_SIZE_N;
  int block_size_n_loops = nloops;

  // {//for debug
  // int* dev_d_w = nullptr;
  // hipMalloc((void**)&dev_d_w, 16*64*sizeof(int));
  // }//for debug
  int d_w_out[16*64]; 
  
  int64_t EM = sorted_token_ids.size(0);
  if (size_m <= BLOCK_SIZE_m) {
    EM = min(EM, size_m * BLOCK_SIZE_m * top_k);
  }
  const int num_token_blocks = (EM + BLOCK_SIZE_m*block_size_m_loops - 1) / (BLOCK_SIZE_m*block_size_m_loops);

  const uint32_t* b_qzeros_ptr;
  if (b_qzeros.has_value())
    b_qzeros_ptr = (const uint32_t*)b_qzeros.value().data_ptr<uint8_t>();
  const float* topk_weights_ptr;
  if (topk_weights.has_value())
    topk_weights_ptr = (const float*)topk_weights.value().data_ptr();

  int groups_per_block_row = BLOCK_SIZE_K / group_size_k;
  TORCH_CHECK(bit == 4 || bit == 8, "bit must be 4 or 8");
  TORCH_CHECK(size_k % BLOCK_SIZE_K == 0,
              "size_k must divisible by BLOCK_SIZE_K");
  TORCH_CHECK(BLOCK_SIZE_K % group_size_k == 0,
              "BLOCK_SIZE_K must divisible by group_size_k");
  TORCH_CHECK(BLOCK_SIZE_m <= 64, "BLOCK_SIZE_m must less or equal to 64");
  TORCH_CHECK(groups_per_block_row == 1 || groups_per_block_row == 2 ||
                  groups_per_block_row == 4 || groups_per_block_row == 8,
              "BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8]");
  
  bool use_atomic = (size_k != BLOCK_SIZE_K*block_size_k_loops);

  


  std::optional<torch::Tensor> output_fp32;

  if (use_atomic){

    output_fp32 = torch::zeros(output.sizes(),output.options().dtype(torch::kFloat32));
    // output_fp32->zero_(); 
  } 
  float milliseconds = 0;
  cudaEvent_t start, stop;
  const char* find_best = std::getenv("WHICH_TO_TEST");
    if (find_best) {
      
    cudaEventCreate(&start);
    cudaEventCreate(&stop);

    cudaEventRecord(start);        // 记录开始
    }
  // if (true/* input.scalar_type() == at::ScalarType::QInt8 */) {

  //   BIT_SWITCH(bit, BIT, [&]{
  //     TOPK_SWITCH(top_k, TOPK, [&]{
  //       BLOCK_M_SWITCH(BLOCK_SIZE_m, BLOCK_SIZE_M_, [&]{
  //         BLOCK_N_SWITCH(BLOCK_SIZE_N, BLOCK_SIZE_N_, [&]{
  //           BLOCK_K_SWITCH(BLOCK_SIZE_K, BLOCK_SIZE_K_, [&]{
  //             // BOOL_SWITCH(b_qzeros.has_value(), has_zp, [&]{
  //               BOOL_SWITCH(topk_weights.has_value(), mul_topk_weight, [&]{
  //                 GROUP_SIZE_N_SWITCH(group_size_n, GROUP_SIZE_N, [&]{
  //                   GROUP_SIZE_K_SWITCH(group_size_k, GROUP_SIZE_K, [&]{
  //                     BLOCK_SIZE_M_LOOPS_SWITCH(block_size_m_loops , BLOCK_SIZE_M_LOOPS, [&]{
  //                     BLOCK_SIZE_N_LOOPS_SWITCH(block_size_n_loops , BLOCK_SIZE_N_LOOPS, [&]{
  //                       BLOCK_SIZE_K_LOOPS_SWITCH(block_size_k_loops , BLOCK_SIZE_K_LOOPS, [&]{
  //                         BOOL_SWITCH(use_atomic , USE_ATOMIC, [&]{
  //                   run_moe_w8a8_gemm_block_wise_kernel2<half, 8, TOPK, BLOCK_SIZE_M_, BLOCK_SIZE_N_, BLOCK_SIZE_K_, false, mul_topk_weight, GROUP_SIZE_N, GROUP_SIZE_K, BLOCK_SIZE_M_LOOPS, BLOCK_SIZE_N_LOOPS, BLOCK_SIZE_K_LOOPS, USE_ATOMIC,256>(
                    
  //                     (const uint32_t*)input.data_ptr<int8_t>(),
  //                     // (const half*)d_input,
  //                     (const float*)a_scales.data_ptr(),
  //                     use_atomic ?(float*)output_fp32->data_ptr():(float*)output.data_ptr(),
  //                     // (float*)output_fp32->data_ptr(),
  //                     (const uint32_t*)b_qweight.data_ptr<int8_t>(),
  //                     // (const uint32_t*)d_w_test,
  //                     ( int*)&d_w_out[0], /*for debug*/ /*使用时需修改为device端地址，这里仅为占位使用*/
  //                     // (const half*)b_scales.data_ptr<at::Half>(), 
  //                     (const float*)b_scales.data_ptr(),
  //                     // (const half*)d_scale, 
  //                     b_qzeros_ptr,
  //                     // (const uint32_t*)d_scale,
  //                     topk_weights_ptr, 
  //                     sorted_token_ids.data_ptr<int32_t>(),
  //                     expert_ids.data_ptr<int32_t>(), 
  //                     num_tokens_post_pad.data_ptr<int32_t>(), 
  //                     // num_tokens_post_pad_value<int32_t>(),
  //                     // num_tokens_post_pad_data_ptr[0],
  //                     num_token_blocks, 
  //                     size_m, 
  //                     size_n,
  //                     size_k
  //                     );
  //                     // printf("run_moe_w8a8_gemm_block_wise\n"); // kernel-1 mma    
  //                     });              
  //                     });
  //                   });
  //                 });
  //               });
  //             });
  //           });
  //         });
  //       });
  //     });
  //   });
  // });
  // } else {
  //   TORCH_CHECK(false, "moe_w8a8_gemm_block_wise only supports int8_t");
  // }

  if (find_best) {
  cudaEventRecord(stop);         // 记录结束
  cudaEventSynchronize(stop);    // 等待 kernel 执行完成

  
  cudaEventElapsedTime(&milliseconds, start, stop); // 计算时间

  /* std::cout << "kernel 1 time: " << milliseconds << " ms" << std::endl; */
  
  
  

  

  std::ofstream ofs("./w8a8_kerne2_1_timecost", std::ios::app); // 追加写入
  if (ofs.is_open()) {
      ofs << milliseconds << std::endl;
      ofs.close();
  }
}
  
  

  if (use_atomic){
    // std::cout<<" fp32 convert to fp16 "<<std::endl;
    output.copy_(output_fp32->to(torch::kFloat16));
  }

  cudaEventDestroy(start);
  cudaEventDestroy(stop);

  // {//for debug
  // hipDeviceSynchronize();
  // hipMemcpy(&d_w_out[0], dev_d_w, 16*64 * sizeof(int), hipMemcpyDeviceToHost);
  // for(int i =0;i<16;i++){
  //   for(int j = 0;j<64 ;j++){
  //     std::cout<<(int)(d_w_out[i*64+j])<<"\t";
  //   }
  //   std::cout<<"|||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||"<<std::endl;
  // }
  // hipFree(dev_d_w);
  // }
  return output;
  
}


torch::Tensor moe_c_moe_w8a8_gemm_block_wise_fp8(torch::Tensor input, torch::Tensor a_scales,torch::Tensor output,
                             torch::Tensor b_qweight, torch::Tensor b_scales,
                             std::optional<torch::Tensor> b_qzeros,
                             std::optional<torch::Tensor> topk_weights,
                             torch::Tensor sorted_token_ids,
                             torch::Tensor expert_ids,
                             torch::Tensor num_tokens_post_pad, int64_t group_size_n, int64_t group_size_k, int64_t top_k,
                             int64_t BLOCK_SIZE_m, int64_t BLOCK_SIZE_n,
                             int64_t BLOCK_SIZE_k, int64_t kloops, int64_t nloops, int64_t bit ) {
  
  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
  const int size_m = input.size(0);
  const int size_n = b_qweight.size(1);
  const int size_k = input.size(1);
  
  int64_t BLOCK_SIZE_N = 64;
  // std::cout<<"BLOCK_SIZE_N IS : "<<BLOCK_SIZE_N<<std::endl;
  int64_t BLOCK_SIZE_K_MIN =4*1024/BLOCK_SIZE_m;
  int64_t BLOCK_SIZE_K_MAX =8*1024/BLOCK_SIZE_m;
  // std::cout<<"BLOCK_SIZE_K_MAX IS : "<<BLOCK_SIZE_K_MAX<<std::endl;
  int64_t BLOCK_SIZE_K = std::min(BLOCK_SIZE_K_MAX,size_k);
  // std::cout<<"BLOCK_SIZE_K IS : "<<BLOCK_SIZE_K<<std::endl;

  int BLOCK_SIZE_M_MAX = std::min(16, size_m);
  // std::cout<<"BLOCK_SIZE_M_MAX IS : "<<BLOCK_SIZE_M_MAX<<std::endl;
  int BLOCK_SIZE_N_MAX_roofline = 64;
  int BLOCK_SIZE_N_MAX = std::min(BLOCK_SIZE_N_MAX_roofline, size_n);
  // std::cout<<"BLOCK_SIZE_N_MAX IS : "<<BLOCK_SIZE_N_MAX<<std::endl;

  // BLOCK_SIZE_K = std::min(128, BLOCK_SIZE_K);
  BLOCK_SIZE_K = 128;
  // std::cout<<"BLOCK_SIZE_K IS : "<<BLOCK_SIZE_K<<std::endl;

  // int block_size_m_loops = 1;// std::min(1,BLOCK_SIZE_M_MAX/BLOCK_SIZE_m);
  int block_size_m_loops = 1;
  int block_size_k_loops = kloops;
  // int block_size_n_loops = BLOCK_SIZE_N_MAX/BLOCK_SIZE_N;
  int block_size_n_loops = nloops;

  // {//for debug
  // int* dev_d_w = nullptr;
  // hipMalloc((void**)&dev_d_w, 16*64*sizeof(int));
  // }//for debug
  int d_w_out[16*64]; 
  
  int64_t EM = sorted_token_ids.size(0);
  if (size_m <= BLOCK_SIZE_m) {
    EM = min(EM, size_m * BLOCK_SIZE_m * top_k);
  }
  const int num_token_blocks = (EM + BLOCK_SIZE_m*block_size_m_loops - 1) / (BLOCK_SIZE_m*block_size_m_loops);

  const uint32_t* b_qzeros_ptr;
  if (b_qzeros.has_value())
    b_qzeros_ptr = (const uint32_t*)b_qzeros.value().data_ptr<uint8_t>();
  const float* topk_weights_ptr;
  if (topk_weights.has_value())
    topk_weights_ptr = (const float*)topk_weights.value().data_ptr();

  int groups_per_block_row = BLOCK_SIZE_K / group_size_k;
  TORCH_CHECK(bit == 4 || bit == 8, "bit must be 4 or 8");
  TORCH_CHECK(size_k % BLOCK_SIZE_K == 0,
              "size_k must divisible by BLOCK_SIZE_K");
  TORCH_CHECK(BLOCK_SIZE_K % group_size_k == 0,
              "BLOCK_SIZE_K must divisible by group_size_k");
  TORCH_CHECK(BLOCK_SIZE_m <= 64, "BLOCK_SIZE_m must less or equal to 64");
  TORCH_CHECK(groups_per_block_row == 1 || groups_per_block_row == 2 ||
                  groups_per_block_row == 4 || groups_per_block_row == 8,
              "BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8]");
  
  bool use_atomic = (size_k != BLOCK_SIZE_K*block_size_k_loops);

  


  std::optional<torch::Tensor> output_fp32;

  if (use_atomic){

    output_fp32 = torch::zeros(output.sizes(),output.options().dtype(torch::kFloat32));
    // output_fp32->zero_(); 
  } 
  float milliseconds = 0;
  cudaEvent_t start, stop;
  const char* find_best = std::getenv("WHICH_TO_TEST");
  if (find_best) {
    
    cudaEventCreate(&start);
    cudaEventCreate(&stop);
    cudaEventRecord(start);        // 记录开始
  }

  // if (true/* input.scalar_type() == at::ScalarType::QInt8 */) {

  //   BIT_SWITCH(bit, BIT, [&]{
  //     TOPK_SWITCH(top_k, TOPK, [&]{
  //       BLOCK_M_SWITCH(BLOCK_SIZE_m, BLOCK_SIZE_M_, [&]{
  //         BLOCK_N_SWITCH(BLOCK_SIZE_N, BLOCK_SIZE_N_, [&]{
  //           BLOCK_K_SWITCH(BLOCK_SIZE_K, BLOCK_SIZE_K_, [&]{
  //             // BOOL_SWITCH(b_qzeros.has_value(), has_zp, [&]{
  //               BOOL_SWITCH(topk_weights.has_value(), mul_topk_weight, [&]{
  //                 GROUP_SIZE_N_SWITCH(group_size_n, GROUP_SIZE_N, [&]{
  //                   GROUP_SIZE_K_SWITCH(group_size_k, GROUP_SIZE_K, [&]{
  //                     BLOCK_SIZE_M_LOOPS_SWITCH(block_size_m_loops , BLOCK_SIZE_M_LOOPS, [&]{
  //                     BLOCK_SIZE_N_LOOPS_SWITCH(block_size_n_loops , BLOCK_SIZE_N_LOOPS, [&]{
  //                       BLOCK_SIZE_K_LOOPS_SWITCH(block_size_k_loops , BLOCK_SIZE_K_LOOPS, [&]{
  //                         BOOL_SWITCH(use_atomic , USE_ATOMIC, [&]{
  //                   run_moe_w8a8_gemm_block_wise_fp8<half, 8, TOPK, BLOCK_SIZE_M_, BLOCK_SIZE_N_, BLOCK_SIZE_K_, false, mul_topk_weight, GROUP_SIZE_N, GROUP_SIZE_K, BLOCK_SIZE_M_LOOPS, BLOCK_SIZE_N_LOOPS, BLOCK_SIZE_K_LOOPS, USE_ATOMIC,256>(
                    
  //                     (const uint32_t*)input.data_ptr<int8_t>(),
  //                     // (const half*)d_input,
  //                     (const float*)a_scales.data_ptr(),
  //                     use_atomic ?(float*)output_fp32->data_ptr():(float*)output.data_ptr(),
  //                     // (float*)output_fp32->data_ptr(),
  //                     (const uint32_t*)b_qweight.data_ptr<int8_t>(),
  //                     // (const uint32_t*)d_w_test,
  //                     ( int*)&d_w_out[0], /*for debug*/ /*使用时需修改为device端地址，这里仅为占位使用*/
  //                     // (const half*)b_scales.data_ptr<at::Half>(), 
  //                     (const float*)b_scales.data_ptr(),
  //                     // (const half*)d_scale, 
  //                     b_qzeros_ptr,
  //                     // (const uint32_t*)d_scale,
  //                     topk_weights_ptr, 
  //                     sorted_token_ids.data_ptr<int32_t>(),
  //                     expert_ids.data_ptr<int32_t>(), 
  //                     num_tokens_post_pad.data_ptr<int32_t>(), 
  //                     // num_tokens_post_pad_value<int32_t>(),
  //                     // num_tokens_post_pad_data_ptr[0],
  //                     num_token_blocks, 
  //                     size_m, 
  //                     size_n,
  //                     size_k
  //                     );
  //                     // printf("run_moe_w8a8_gemm_block_wise\n"); // kernel-1 mma    
  //                     });              
  //                     });
  //                   });
  //                 });
  //               });
  //             });
  //           });
  //         });
  //       });
  //     });
  //   });
  // });
  // } else {
  //   TORCH_CHECK(false, "moe_w8a8_gemm_block_wise only supports int8_t");
  // }


  if (find_best) {
  cudaEventRecord(stop);         // 记录结束
  cudaEventSynchronize(stop);    // 等待 kernel 执行完成

  
  cudaEventElapsedTime(&milliseconds, start, stop); // 计算时间

  /* std::cout << "kernel 1 time: " << milliseconds << " ms" << std::endl; */
  
  

  std::ofstream ofs("./w8a8_kernel_1_timecost", std::ios::app); // 追加写入
  if (ofs.is_open()) {
      ofs << milliseconds << std::endl;
      ofs.close();
  }
}
  
  

  if (use_atomic){
    // std::cout<<" fp32 convert to fp16 "<<std::endl;
    output.copy_(output_fp32->to(torch::kBFloat16));
  }
  cudaEventDestroy(start);
  cudaEventDestroy(stop);
  
  // {//for debug
  // hipDeviceSynchronize();
  // hipMemcpy(&d_w_out[0], dev_d_w, 16*64 * sizeof(int), hipMemcpyDeviceToHost);
  // for(int i =0;i<16;i++){
  //   for(int j = 0;j<64 ;j++){
  //     std::cout<<(int)(d_w_out[i*64+j])<<"\t";
  //   }
  //   std::cout<<"|||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||"<<std::endl;
  // }
  // hipFree(dev_d_w);
  // }
  return output;
  
}

torch::Tensor moe_c_moe_w8a8_gemm_block_wise_kernel2_fp8(torch::Tensor input, torch::Tensor a_scales,torch::Tensor output,
                             torch::Tensor b_qweight, torch::Tensor b_scales,
                             std::optional<torch::Tensor> b_qzeros,
                             std::optional<torch::Tensor> topk_weights,
                             torch::Tensor sorted_token_ids,
                             torch::Tensor expert_ids,
                             torch::Tensor num_tokens_post_pad, int64_t group_size_n, int64_t group_size_k, int64_t top_k,
                             int64_t BLOCK_SIZE_m, int64_t BLOCK_SIZE_n,
                             int64_t BLOCK_SIZE_k, int64_t kloops, int64_t nloops, int64_t bit ) {
  
  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));

  const int size_m = input.size(0);
  const int size_n = b_qweight.size(1);
  const int size_k = input.size(1);
  
  int64_t BLOCK_SIZE_N = 64;
  // std::cout<<"BLOCK_SIZE_N IS : "<<BLOCK_SIZE_N<<std::endl;
  int64_t BLOCK_SIZE_K_MIN =4*1024/BLOCK_SIZE_m;
  int64_t BLOCK_SIZE_K_MAX =8*1024/BLOCK_SIZE_m;
  // std::cout<<"BLOCK_SIZE_K_MAX IS : "<<BLOCK_SIZE_K_MAX<<std::endl;
  int64_t BLOCK_SIZE_K = std::min(BLOCK_SIZE_K_MAX,size_k);
  // std::cout<<"BLOCK_SIZE_K IS : "<<BLOCK_SIZE_K<<std::endl;

  int BLOCK_SIZE_M_MAX = std::min(16, size_m);
  // std::cout<<"BLOCK_SIZE_M_MAX IS : "<<BLOCK_SIZE_M_MAX<<std::endl;
  int BLOCK_SIZE_N_MAX_roofline = 64;
  int BLOCK_SIZE_N_MAX = std::min(BLOCK_SIZE_N_MAX_roofline, size_n);
  // std::cout<<"BLOCK_SIZE_N_MAX IS : "<<BLOCK_SIZE_N_MAX<<std::endl;

  // BLOCK_SIZE_K = std::min(128, BLOCK_SIZE_K);
  BLOCK_SIZE_K = 128;
  // std::cout<<"BLOCK_SIZE_K IS : "<<BLOCK_SIZE_K<<std::endl;

  // int block_size_m_loops = 1;// std::min(1,BLOCK_SIZE_M_MAX/BLOCK_SIZE_m);
  int block_size_m_loops = 1;
  int block_size_k_loops =kloops;
  // int block_size_n_loops = BLOCK_SIZE_N_MAX/BLOCK_SIZE_N;
  int block_size_n_loops = nloops;

  // {//for debug
  // int* dev_d_w = nullptr;
  // hipMalloc((void**)&dev_d_w, 16*64*sizeof(int));
  // }//for debug
  int d_w_out[16*64]; 
  
  int64_t EM = sorted_token_ids.size(0);
  if (size_m <= BLOCK_SIZE_m) {
    EM = min(EM, size_m * BLOCK_SIZE_m * top_k);
  }
  const int num_token_blocks = (EM + BLOCK_SIZE_m*block_size_m_loops - 1) / (BLOCK_SIZE_m*block_size_m_loops);

  const uint32_t* b_qzeros_ptr;
  if (b_qzeros.has_value())
    b_qzeros_ptr = (const uint32_t*)b_qzeros.value().data_ptr<uint8_t>();
  const float* topk_weights_ptr;
  if (topk_weights.has_value())
    topk_weights_ptr = (const float*)topk_weights.value().data_ptr();

  int groups_per_block_row = BLOCK_SIZE_K / group_size_k;
  TORCH_CHECK(bit == 4 || bit == 8, "bit must be 4 or 8");
  TORCH_CHECK(size_k % BLOCK_SIZE_K == 0,
              "size_k must divisible by BLOCK_SIZE_K");
  TORCH_CHECK(BLOCK_SIZE_K % group_size_k == 0,
              "BLOCK_SIZE_K must divisible by group_size_k");
  TORCH_CHECK(BLOCK_SIZE_m <= 64, "BLOCK_SIZE_m must less or equal to 64");
  TORCH_CHECK(groups_per_block_row == 1 || groups_per_block_row == 2 ||
                  groups_per_block_row == 4 || groups_per_block_row == 8,
              "BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8]");
  
  bool use_atomic = (size_k != BLOCK_SIZE_K*block_size_k_loops);

  


  std::optional<torch::Tensor> output_fp32;

  if (use_atomic){

    output_fp32 = torch::zeros(output.sizes(),output.options().dtype(torch::kFloat32));
    // output_fp32->zero_(); 
  } 
  float milliseconds = 0;
  cudaEvent_t start, stop;
  const char* find_best = std::getenv("WHICH_TO_TEST");
    if (find_best) {
      
    cudaEventCreate(&start);
    cudaEventCreate(&stop);

    cudaEventRecord(start);        // 记录开始
    }

  // if (true/* input.scalar_type() == at::ScalarType::QInt8 */) {

  //   BIT_SWITCH(bit, BIT, [&]{
  //     TOPK_SWITCH(top_k, TOPK, [&]{
  //       BLOCK_M_SWITCH(BLOCK_SIZE_m, BLOCK_SIZE_M_, [&]{
  //         BLOCK_N_SWITCH(BLOCK_SIZE_N, BLOCK_SIZE_N_, [&]{
  //           BLOCK_K_SWITCH(BLOCK_SIZE_K, BLOCK_SIZE_K_, [&]{
  //             // BOOL_SWITCH(b_qzeros.has_value(), has_zp, [&]{
  //               BOOL_SWITCH(topk_weights.has_value(), mul_topk_weight, [&]{
  //                 GROUP_SIZE_N_SWITCH(group_size_n, GROUP_SIZE_N, [&]{
  //                   GROUP_SIZE_K_SWITCH(group_size_k, GROUP_SIZE_K, [&]{
  //                     BLOCK_SIZE_M_LOOPS_SWITCH(block_size_m_loops , BLOCK_SIZE_M_LOOPS, [&]{
  //                     BLOCK_SIZE_N_LOOPS_SWITCH(block_size_n_loops , BLOCK_SIZE_N_LOOPS, [&]{
  //                       BLOCK_SIZE_K_LOOPS_SWITCH(block_size_k_loops , BLOCK_SIZE_K_LOOPS, [&]{
  //                         BOOL_SWITCH(use_atomic , USE_ATOMIC, [&]{
  //                   run_moe_w8a8_gemm_block_wise_kernel2_fp8<half, 8, TOPK, BLOCK_SIZE_M_, BLOCK_SIZE_N_, BLOCK_SIZE_K_, false, mul_topk_weight, GROUP_SIZE_N, GROUP_SIZE_K, BLOCK_SIZE_M_LOOPS, BLOCK_SIZE_N_LOOPS, BLOCK_SIZE_K_LOOPS, USE_ATOMIC,256>(
                    
  //                     (const uint32_t*)input.data_ptr<int8_t>(),
  //                     // (const half*)d_input,
  //                     (const float*)a_scales.data_ptr(),
  //                     use_atomic ?(float*)output_fp32->data_ptr():(float*)output.data_ptr(),
  //                     // (float*)output_fp32->data_ptr(),
  //                     (const uint32_t*)b_qweight.data_ptr<int8_t>(),
  //                     // (const uint32_t*)d_w_test,
  //                     ( int*)&d_w_out[0], /*for debug*/ /*使用时需修改为device端地址，这里仅为占位使用*/
  //                     // (const half*)b_scales.data_ptr<at::Half>(), 
  //                     (const float*)b_scales.data_ptr(),
  //                     // (const half*)d_scale, 
  //                     b_qzeros_ptr,
  //                     // (const uint32_t*)d_scale,
  //                     topk_weights_ptr, 
  //                     sorted_token_ids.data_ptr<int32_t>(),
  //                     expert_ids.data_ptr<int32_t>(), 
  //                     num_tokens_post_pad.data_ptr<int32_t>(), 
  //                     // num_tokens_post_pad_value<int32_t>(),
  //                     // num_tokens_post_pad_data_ptr[0],
  //                     num_token_blocks, 
  //                     size_m, 
  //                     size_n,
  //                     size_k
  //                     );
  //                     // printf("run_moe_w8a8_gemm_block_wise\n"); // kernel-1 mma    
  //                     });              
  //                     });
  //                   });
  //                 });
  //               });
  //             });
  //           });
  //         });
  //       });
  //     });
  //   });
  // });
  // } else {
  //   TORCH_CHECK(false, "moe_w8a8_gemm_block_wise only supports int8_t");
  // }

  if (find_best) {
  cudaEventRecord(stop);         // 记录结束
  cudaEventSynchronize(stop);    // 等待 kernel 执行完成

  
  cudaEventElapsedTime(&milliseconds, start, stop); // 计算时间

  /* std::cout << "kernel 1 time: " << milliseconds << " ms" << std::endl; */
  
  
  

  

  std::ofstream ofs("./w8a8_kerne2_1_timecost", std::ios::app); // 追加写入
  if (ofs.is_open()) {
      ofs << milliseconds << std::endl;
      ofs.close();
  }
}
  
  

  if (use_atomic){
    // std::cout<<" fp32 convert to fp16 "<<std::endl;
    output.copy_(output_fp32->to(torch::kBFloat16));
  }

  cudaEventDestroy(start);
  cudaEventDestroy(stop);

  // {//for debug
  // hipDeviceSynchronize();
  // hipMemcpy(&d_w_out[0], dev_d_w, 16*64 * sizeof(int), hipMemcpyDeviceToHost);
  // for(int i =0;i<16;i++){
  //   for(int j = 0;j<64 ;j++){
  //     std::cout<<(int)(d_w_out[i*64+j])<<"\t";
  //   }
  //   std::cout<<"|||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||"<<std::endl;
  // }
  // hipFree(dev_d_w);
  // }
  return output;
}

torch::Tensor moe_c_moe_w8a16_gemm_awq(torch::Tensor input, torch::Tensor output,
                             torch::Tensor b_qweight, torch::Tensor b_scales,
                             std::optional<torch::Tensor> b_qzeros,
                             std::optional<torch::Tensor> topk_weights,
                             torch::Tensor sorted_token_ids,
                             torch::Tensor expert_ids,
                             torch::Tensor num_tokens_post_pad, int64_t top_k,
                             int64_t BLOCK_SIZE_m, int64_t BLOCK_SIZE_n,
                             int64_t BLOCK_SIZE_k, int64_t bit) {
 
  // const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
  // auto options = torch::TensorOptions().dtype(input.dtype()).device(input.device());

  const int size_m = input.size(0);
  const int size_n = b_qweight.size(1);
  const int size_k = input.size(1);
  const int group_size = size_k / b_scales.size(2);
  /*经验值4-8个block，lds为64k，左矩阵BM*BK*2 范围为8k-16k， 所以BM*BN应在4k-8k*/
  // std::cout<<"size_m IS : "<<size_m<<std::endl;
  // std::cout<<"size_n IS : "<<size_n<<std::endl;
  // std::cout<<"size_k IS : "<<size_k<<std::endl;

  // std::cout<<"BLOCK_SIZE_m IS : "<<BLOCK_SIZE_m<<std::endl;
  // std::cout<<"BLOCK_SIZE_n IS : "<<BLOCK_SIZE_n<<std::endl;
  // std::cout<<"BLOCK_SIZE_k IS : "<<BLOCK_SIZE_k<<std::endl;
  int64_t BLOCK_SIZE_N = std::min(64, size_n);
  // std::cout<<"BLOCK_SIZE_N IS : "<<BLOCK_SIZE_N<<std::endl;
  int64_t BLOCK_SIZE_K_MIN =4*1024/BLOCK_SIZE_m;
  int64_t BLOCK_SIZE_K_MAX =8*1024/BLOCK_SIZE_m;
  // std::cout<<"BLOCK_SIZE_K_MAX IS : "<<BLOCK_SIZE_K_MAX<<std::endl;
  int64_t BLOCK_SIZE_K = std::min(BLOCK_SIZE_K_MAX,size_k);
  // std::cout<<"BLOCK_SIZE_K IS : "<<BLOCK_SIZE_K<<std::endl;

  int BLOCK_SIZE_M_MAX = std::min(16, size_m);
  // std::cout<<"BLOCK_SIZE_M_MAX IS : "<<BLOCK_SIZE_M_MAX<<std::endl;
  int BLOCK_SIZE_N_MAX_roofline = 256;
  int BLOCK_SIZE_N_MAX = std::min(BLOCK_SIZE_N_MAX_roofline, size_n);
  // std::cout<<"BLOCK_SIZE_N_MAX IS : "<<BLOCK_SIZE_N_MAX<<std::endl;

  BLOCK_SIZE_K = std::min(128, BLOCK_SIZE_K);
  // std::cout<<"BLOCK_SIZE_K IS : "<<BLOCK_SIZE_K<<std::endl;

  int block_size_m_loops = 1;// std::min(1,BLOCK_SIZE_M_MAX/BLOCK_SIZE_m);
  int block_size_k_loops = std::min(size_k/BLOCK_SIZE_K, 2);
  int block_size_n_loops = BLOCK_SIZE_N_MAX/BLOCK_SIZE_N;

  half_t * d_w_out =nullptr;
  int64_t EM = sorted_token_ids.size(0);
  
  if (size_m <= BLOCK_SIZE_m) {
    EM = min(EM, size_m * BLOCK_SIZE_m * top_k);
  }
  const int num_token_blocks = (EM + BLOCK_SIZE_m*block_size_m_loops - 1) / (BLOCK_SIZE_m*block_size_m_loops);

  const uint32_t* b_qzeros_ptr;
  if (b_qzeros.has_value())
    b_qzeros_ptr = (const uint32_t*)b_qzeros.value().data_ptr<uint8_t>();
  const float* topk_weights_ptr;
  if (topk_weights.has_value())
    topk_weights_ptr = (const float*)topk_weights.value().data_ptr();

  int groups_per_block_row = BLOCK_SIZE_K / group_size;
  TORCH_CHECK(bit == 4 || bit == 8, "bit must be 4 or 8");
  TORCH_CHECK(size_k % BLOCK_SIZE_K == 0,
              "size_k must divisible by BLOCK_SIZE_K");
  TORCH_CHECK(BLOCK_SIZE_K % group_size == 0,
              "BLOCK_SIZE_K must divisible by group_size");
  TORCH_CHECK(BLOCK_SIZE_m <= 64, "BLOCK_SIZE_m must less or equal to 64");
  TORCH_CHECK(groups_per_block_row == 1 || groups_per_block_row == 2 ||
                  groups_per_block_row == 4 || groups_per_block_row == 8,
              "BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8]");
  
  bool use_atomic = (size_k != BLOCK_SIZE_K*block_size_k_loops);

  // std::cout<<"size_m is : "<<size_m<<std::endl;
  // std::cout<<"size_n is : "<<size_n<<std::endl;
  // std::cout<<"size_k is : "<<size_k<<std::endl;
  // std::cout<<"group_size : "<<group_size<<std::endl;
  // std::cout<<"BLOCK_SIZE_m IS : "<<BLOCK_SIZE_m<<std::endl;
  // std::cout<<"BLOCK_SIZE_n IS : "<<BLOCK_SIZE_n<<std::endl;
  // std::cout<<"BLOCK_SIZE_k IS : "<<BLOCK_SIZE_k<<std::endl;

  // std::cout<<"BLOCK_SIZE_M IS : "<<BLOCK_SIZE_M<<std::endl;
  // std::cout<<"BLOCK_SIZE_N IS : "<<BLOCK_SIZE_N<<std::endl;
  // std::cout<<"BLOCK_SIZE_K IS : "<<BLOCK_SIZE_K<<std::endl;
  // std::cout<<"block_size_m_loops IS : "<<block_size_m_loops<<std::endl;
  // std::cout<<"block_size_k_loops IS : "<<block_size_k_loops<<std::endl;
  // std::cout<<"block_size_n_loops IS : "<<block_size_n_loops<<std::endl;



  std::optional<torch::Tensor> output_fp32;

  if (use_atomic){

    output_fp32 = torch::zeros(output.sizes(),output.options().dtype(torch::kFloat32));
    // output_fp32->zero_(); 
  } 
  //  hipDeviceSynchronize();
  // hipEventRecord(stop, 0 );
  // hipEventSynchronize( stop );

  // float ave_time;
  // hipEventElapsedTime( &ave_time,start, stop );
  // printf( "Time to generate: %9f ms\n", ave_time );

  // std::cout<<"input scalar type is :"<<input.scalar_type()<<std::endl;
  // std::cout<<"block_size_n_loops IS : "<<block_size_n_loops<<std::endl;

  // if (input.scalar_type() == at::ScalarType::Half) {

  //   BIT_SWITCH(bit, BIT, [&]{
  //     TOPK_SWITCH(top_k, TOPK, [&]{
  //       BLOCK_M_SWITCH(BLOCK_SIZE_m, BLOCK_SIZE_M_, [&]{
  //         BLOCK_N_SWITCH(BLOCK_SIZE_N, BLOCK_SIZE_N_, [&]{
  //           BLOCK_K_SWITCH(BLOCK_SIZE_K, BLOCK_SIZE_K_, [&]{
  //             // BOOL_SWITCH(b_qzeros.has_value(), has_zp, [&]{
  //               BOOL_SWITCH(topk_weights.has_value(), mul_topk_weight, [&]{
  //                 GROUP_SIZE_K_SWITCH(group_size, GROUP_SIZE_K, [&]{
  //                   BLOCK_SIZE_M_LOOPS_SWITCH(block_size_m_loops , BLOCK_SIZE_M_LOOPS, [&]{
  //                     BLOCK_SIZE_N_LOOPS_SWITCH(block_size_n_loops , BLOCK_SIZE_N_LOOPS, [&]{
  //                       BLOCK_SIZE_K_LOOPS_SWITCH(block_size_k_loops , BLOCK_SIZE_K_LOOPS, [&]{
  //                         BOOL_SWITCH(use_atomic , USE_ATOMIC, [&]{
  //                   run_moe_wna16_gemm_awq<half, 8, TOPK, BLOCK_SIZE_M_, BLOCK_SIZE_N_, BLOCK_SIZE_K_, true, mul_topk_weight, GROUP_SIZE_K, BLOCK_SIZE_M_LOOPS, BLOCK_SIZE_N_LOOPS, BLOCK_SIZE_K_LOOPS, USE_ATOMIC,256>(
  //                     (const half*)input.data_ptr<at::Half>(),
  //                     // (const half*)d_input,
  //                     use_atomic ?(float*)output_fp32->data_ptr():(float*)output.data_ptr(),
  //                     // (float*)output_fp32->data_ptr(),
  //                     (const uint32_t*)b_qweight.data_ptr<uint8_t>(),
  //                     // (const uint32_t*)d_w_test,
  //                     ( half_t*) d_w_out, /*for debug*/
  //                     (const half*)b_scales.data_ptr<at::Half>(), 
  //                     // (const half*)d_scale, 
  //                     b_qzeros_ptr,
  //                     // (const uint32_t*)d_scale,
  //                     topk_weights_ptr, 
  //                     sorted_token_ids.data_ptr<int32_t>(),
  //                     expert_ids.data_ptr<int32_t>(), 
  //                     num_tokens_post_pad.data_ptr<int32_t>(), 
  //                     // num_tokens_post_pad_value<int32_t>(),
  //                     // num_tokens_post_pad_data_ptr[0],
  //                     num_token_blocks, 
  //                     size_m, 
  //                     size_n,
  //                     size_k
  //                     ); // kernel-1 mma    
  //                     // run_moe_wna16_gemm_blockwise_<half, BIT, TOPK, BLOCK_SIZE_M_, BLOCK_SIZE_N_, BLOCK_SIZE_K_, true, mul_topk_weight, GROUP_SIZE, BLOCK_SIZE_M_LOOPS, BLOCK_SIZE_N_LOOPS, BLOCK_SIZE_K_LOOPS, USE_ATOMIC,256>(
  //                     // (const half*)input.data_ptr<at::Half>(),
  //                     // // (const half*)d_input,
  //                     // use_atomic ?(float*)output_fp32->data_ptr():(float*)output.data_ptr(),
  //                     // // (float*)output_fp32->data_ptr(),
  //                     // (const uint32_t*)b_qweight.data_ptr<uint8_t>(),
  //                     // // (const uint32_t*)d_w_test,
  //                     // ( half_t*) d_w_out, /*for debug*/
  //                     // (const half*)b_scales.data_ptr<at::Half>(), 
  //                     // // (const half*)d_scale, 
  //                     // b_qzeros_ptr,
  //                     // // (const uint32_t*)d_scale,
  //                     // topk_weights_ptr, 
  //                     // sorted_token_ids.data_ptr<int32_t>(),
  //                     // expert_ids.data_ptr<int32_t>(), 
  //                     // num_tokens_post_pad.data_ptr<int32_t>(), 
  //                     // // num_tokens_post_pad_value<int32_t>(),
  //                     // // num_tokens_post_pad_data_ptr[0],
  //                     // num_token_blocks, 
  //                     // size_m, 
  //                     // size_n,
  //                     // size_k
  //                     // ); // kernel-1 mma    
  //                   });              
  //                   });
  //                 });
  //               });
  //             });
  //           });
  //         });
  //       });
  //     });
  //   });
  // });
  // } 
  
  // else {
  //   TORCH_CHECK(false, "moe_w8a16_gemm_awq only supports float16");
  // }

  // hipEvent_t  start, stop;
  // hipEventCreate(&start);
  // hipEventCreate(&stop);
  // hipEventRecord(start, 0);
  
  if (use_atomic){
    // std::cout<<" fp32 convert to fp16 "<<std::endl;
    output.copy_(output_fp32->to(torch::kFloat16));
  }

  // hipDeviceSynchronize();
  // hipEventRecord(stop, 0 );
  // hipEventSynchronize( stop );

  // float ave_time;
  // hipEventElapsedTime( &ave_time,start, stop );
  // printf( "Time to generate: %9f ms\n", ave_time );

  return output;
  
}


torch::Tensor moe_c_moe_w8a16_gemm_block_wise(torch::Tensor input, torch::Tensor output,
                             torch::Tensor b_qweight, torch::Tensor b_scales,
                             std::optional<torch::Tensor> b_qzeros,
                             std::optional<torch::Tensor> topk_weights,
                             torch::Tensor sorted_token_ids,
                             torch::Tensor expert_ids,
                             torch::Tensor num_tokens_post_pad, int64_t group_size_n, int64_t group_size_k, int64_t top_k,
                             int64_t BLOCK_SIZE_m, int64_t BLOCK_SIZE_n,
                             int64_t BLOCK_SIZE_k, int64_t bit) {
  // const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
  // auto options = torch::TensorOptions().dtype(input.dtype()).device(input.device());
  
  

  const int size_m = input.size(0);
  const int size_n = b_qweight.size(1);
  const int size_k = input.size(1);
  // const int group_size = size_k / b_scales.size(2);

  /*经验值4-8个block，lds为64k，左矩阵BM*BK*2 范围为8k-16k， 所以BM*BN应在4k-8k*/

  int64_t BLOCK_SIZE_N = std::min(64, size_n);
  // std::cout<<"BLOCK_SIZE_N IS : "<<BLOCK_SIZE_N<<std::endl;
  int64_t BLOCK_SIZE_K_MIN =4*1024/BLOCK_SIZE_m;
  int64_t BLOCK_SIZE_K_MAX =8*1024/BLOCK_SIZE_m;
  // std::cout<<"BLOCK_SIZE_K_MAX IS : "<<BLOCK_SIZE_K_MAX<<std::endl;
  int64_t BLOCK_SIZE_K = std::min(BLOCK_SIZE_K_MAX,size_k);
  // std::cout<<"BLOCK_SIZE_K IS : "<<BLOCK_SIZE_K<<std::endl;

  int BLOCK_SIZE_M_MAX = std::min(16, size_m);
  // std::cout<<"BLOCK_SIZE_M_MAX IS : "<<BLOCK_SIZE_M_MAX<<std::endl;
  int BLOCK_SIZE_N_MAX_roofline = 256;
  int BLOCK_SIZE_N_MAX = std::min(BLOCK_SIZE_N_MAX_roofline, size_n);
  // std::cout<<"BLOCK_SIZE_N_MAX IS : "<<BLOCK_SIZE_N_MAX<<std::endl;

  BLOCK_SIZE_K = std::min(128, BLOCK_SIZE_K);
  // std::cout<<"BLOCK_SIZE_K IS : "<<BLOCK_SIZE_K<<std::endl;

  int block_size_m_loops = 1;// std::min(1,BLOCK_SIZE_M_MAX/BLOCK_SIZE_m);
  int block_size_k_loops = std::min(size_k/BLOCK_SIZE_K, 2);
  int block_size_n_loops = BLOCK_SIZE_N_MAX/BLOCK_SIZE_N;

  half_t * d_w_out =nullptr;
  int64_t EM = sorted_token_ids.size(0);
  
  if (size_m <= BLOCK_SIZE_m) {
    EM = min(EM, size_m * BLOCK_SIZE_m * top_k);
  }
  const int num_token_blocks = (EM + BLOCK_SIZE_m*block_size_m_loops - 1) / (BLOCK_SIZE_m*block_size_m_loops);

  const uint32_t* b_qzeros_ptr;
  if (b_qzeros.has_value())
    b_qzeros_ptr = (const uint32_t*)b_qzeros.value().data_ptr<uint8_t>();
  const float* topk_weights_ptr;
  if (topk_weights.has_value())
    topk_weights_ptr = (const float*)topk_weights.value().data_ptr();

  int groups_per_block_row = BLOCK_SIZE_K / group_size_k;
  TORCH_CHECK(bit == 4 || bit == 8, "bit must be 4 or 8");
  TORCH_CHECK(size_k % BLOCK_SIZE_K == 0,
              "size_k must divisible by BLOCK_SIZE_K");
  TORCH_CHECK(BLOCK_SIZE_K % group_size_k == 0,
              "BLOCK_SIZE_K must divisible by group_size_k");
  TORCH_CHECK(BLOCK_SIZE_m <= 64, "BLOCK_SIZE_m must less or equal to 64");
  TORCH_CHECK(groups_per_block_row == 1 || groups_per_block_row == 2 ||
                  groups_per_block_row == 4 || groups_per_block_row == 8,
              "BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8]");
  
  bool use_atomic = (size_k != BLOCK_SIZE_K*block_size_k_loops);



  std::optional<torch::Tensor> output_fp32;

  if (use_atomic){

    output_fp32 = torch::zeros(output.sizes(),output.options().dtype(torch::kFloat32));
    // output_fp32->zero_(); 
  } 

  // if (input.scalar_type() == at::ScalarType::Half) {

  //   BIT_SWITCH(bit, BIT, [&]{
  //     TOPK_SWITCH(top_k, TOPK, [&]{
  //       BLOCK_M_SWITCH(BLOCK_SIZE_m, BLOCK_SIZE_M_, [&]{
  //         BLOCK_N_SWITCH(BLOCK_SIZE_N, BLOCK_SIZE_N_, [&]{
  //           BLOCK_K_SWITCH(BLOCK_SIZE_K, BLOCK_SIZE_K_, [&]{
  //             // BOOL_SWITCH(b_qzeros.has_value(), has_zp, [&]{
  //               BOOL_SWITCH(topk_weights.has_value(), mul_topk_weight, [&]{
  //                 GROUP_SIZE_N_SWITCH(group_size_n, GROUP_SIZE_N, [&]{
  //                   GROUP_SIZE_K_SWITCH(group_size_k, GROUP_SIZE_K, [&]{
  //                     BLOCK_SIZE_M_LOOPS_SWITCH(block_size_m_loops , BLOCK_SIZE_M_LOOPS, [&]{
  //                     BLOCK_SIZE_N_LOOPS_SWITCH(block_size_n_loops , BLOCK_SIZE_N_LOOPS, [&]{
  //                       BLOCK_SIZE_K_LOOPS_SWITCH(block_size_k_loops , BLOCK_SIZE_K_LOOPS, [&]{
  //                         BOOL_SWITCH(use_atomic , USE_ATOMIC, [&]{
  //                   run_moe_wna16_gemm_block_wise<half, 8, TOPK, BLOCK_SIZE_M_, BLOCK_SIZE_N_, BLOCK_SIZE_K_, false, mul_topk_weight, GROUP_SIZE_N, GROUP_SIZE_K, BLOCK_SIZE_M_LOOPS, BLOCK_SIZE_N_LOOPS, BLOCK_SIZE_K_LOOPS, USE_ATOMIC,256>(
  //                     (const half*)input.data_ptr<at::Half>(),
  //                     // (const half*)d_input,
  //                     use_atomic ?(float*)output_fp32->data_ptr():(float*)output.data_ptr(),
  //                     // (float*)output_fp32->data_ptr(),
  //                     (const uint32_t*)b_qweight.data_ptr<int8_t>(),
  //                     // (const uint32_t*)d_w_test,
  //                     ( half_t*) d_w_out, /*for debug*/
  //                     // (const half*)b_scales.data_ptr<at::Half>(), 
  //                     (const float*)b_scales.data_ptr(),
  //                     // (const half*)d_scale, 
  //                     b_qzeros_ptr,
  //                     // (const uint32_t*)d_scale,
  //                     topk_weights_ptr, 
  //                     sorted_token_ids.data_ptr<int32_t>(),
  //                     expert_ids.data_ptr<int32_t>(), 
  //                     num_tokens_post_pad.data_ptr<int32_t>(), 
  //                     // num_tokens_post_pad_value<int32_t>(),
  //                     // num_tokens_post_pad_data_ptr[0],
  //                     num_token_blocks, 
  //                     size_m, 
  //                     size_n,
  //                     size_k
  //                     ); // kernel-1 mma    
  //                     });              
  //                     });
  //                   });
  //                 });
  //               });
  //             });
  //           });
  //         });
  //       });
  //     });
  //   });
  // });
  // } else {
  //   TORCH_CHECK(false, "moe_w8a16_gemm_block_wise only supports float16");
  // }
  // // half_t* tmp = reinterpret_cast<half_t*>(output.data_ptr());
  // // half_t* host_tmp = new half[1];  // 仅拷贝第一个值
  // // hipMemcpy(host_tmp, tmp, sizeof(half_t), hipMemcpyDeviceToHost);

  // // float first_value = __half2float(host_tmp[0]);
  // // std::cout << "first value: " << first_value << std::endl;
  // // delete[] host_tmp;

  // if (use_atomic){
  //   // std::cout<<" fp32 convert to fp16 "<<std::endl;
  //   output.copy_(output_fp32->to(torch::kFloat16));
  // }
 
  return output;
  
}

torch::Tensor moe_c_moe_wna16_gemm_base(torch::Tensor input, torch::Tensor output,
                             torch::Tensor b_qweight, torch::Tensor b_scales,
                             std::optional<torch::Tensor> b_qzeros,
                             std::optional<torch::Tensor> topk_weights,
                             torch::Tensor sorted_token_ids,
                             torch::Tensor expert_ids,
                             torch::Tensor num_tokens_post_pad, int64_t top_k,
                             int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
                             int64_t BLOCK_SIZE_K, int64_t bit) {
  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
  auto options =
      torch::TensorOptions().dtype(input.dtype()).device(input.device());

  const int size_m = input.size(0);
  const int size_n = b_qweight.size(1);
  const int size_k = input.size(1);
  const int group_size = size_k / b_scales.size(2);
  
  // BLOCK_SIZE_K = std::min(group_size, BLOCK_SIZE_K);
  BLOCK_SIZE_K = 64;
  BLOCK_SIZE_N = 256; 
  //std::cout<<"BLOCK_SIZE_N IS : "<<BLOCK_SIZE_N<<std::endl;

  // std::cout<<"group_size : "<<group_size<<std::endl;
  // std::cout<<"BLOCK_SIZE_K IS : "<<BLOCK_SIZE_K<<std::endl;
  // std::cout<<"group_size : "<<group_size<<std::endl;
  int64_t EM = sorted_token_ids.size(0);
  if (size_m <= BLOCK_SIZE_M) {
    EM = min(EM, size_m * BLOCK_SIZE_M * top_k);
  }
  const int num_token_blocks = (EM + BLOCK_SIZE_M - 1) / BLOCK_SIZE_M;

  const uint32_t* b_qzeros_ptr;
  if (b_qzeros.has_value())
    b_qzeros_ptr = (const uint32_t*)b_qzeros.value().data_ptr<uint8_t>();
  const float* topk_weights_ptr;
  if (topk_weights.has_value())
    topk_weights_ptr = (const float*)topk_weights.value().data_ptr();

  int groups_per_block_row = BLOCK_SIZE_K / group_size;

  TORCH_CHECK(bit == 4 || bit == 8, "bit must be 4 or 8");
  TORCH_CHECK(size_k % BLOCK_SIZE_K == 0,
              "size_k must divisible by BLOCK_SIZE_K");
  TORCH_CHECK(BLOCK_SIZE_K % group_size == 0,
              "BLOCK_SIZE_K must divisible by group_size");
  TORCH_CHECK(BLOCK_SIZE_M <= 64, "BLOCK_SIZE_M must less or equal to 64");
  TORCH_CHECK(groups_per_block_row == 1 || groups_per_block_row == 2 ||
                  groups_per_block_row == 4 || groups_per_block_row == 8,
              "BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8]");
  
  torch::Tensor output_fp32 = torch::empty(output.sizes(),output.options().dtype(torch::kFloat32));
  //   if (input.scalar_type() == at::ScalarType::Half) {
  //    half_t * d_w_out =nullptr;
  //     TOPK_SWITCH(top_k, TOPK, [&]{
  //       BLOCK_M_SWITCH(BLOCK_SIZE_M, BLOCK_SIZE_M_, [&]{
  //         BLOCK_K_SWITCH(BLOCK_SIZE_K, BLOCK_SIZE_K_, [&]{
  //             // BOOL_SWITCH(b_qzeros.has_value(), has_zp, [&]{
  //               BOOL_SWITCH(topk_weights.has_value(), mul_topk_weight, [&]{
  //                 GROUP_SIZE_SWITCH(group_size, GROUP_SIZE, [&]{
  //                   // run_moe_wna16_gemm<half, bit, top_k, BLOCK_SIZE_M_, BLOCK_SIZE_N, BLOCK_SIZE_K, true, mul_topk_weight, group_size>(
  //                     run_moe_wna16_gemm_base<half, 4, TOPK, BLOCK_SIZE_M_, 256, BLOCK_SIZE_K_, true, mul_topk_weight, GROUP_SIZE>(
  //                     (const half*)input.data_ptr<at::Half>(),
  //                     (float*)output_fp32.data_ptr(),
  //                     (const uint32_t*)b_qweight.data_ptr<uint8_t>(),
  //                     (const half*)b_scales.data_ptr<at::Half>(), 
  //                     b_qzeros_ptr,
  //                     topk_weights_ptr, 
  //                     sorted_token_ids.data_ptr<int32_t>(),
  //                     expert_ids.data_ptr<int32_t>(), 
  //                     num_tokens_post_pad.data_ptr<int32_t>(), 
  //                     num_token_blocks, 
  //                     size_m, 
  //                     size_n,
  //                     size_k
  //                   );
  //                 });
  //                 });
  //             });
  //         });
  //     });
  // } 
  // else if (input.scalar_type() == at::ScalarType::BFloat16) {
  //    __hip_bfloat16 * d_w_out =nullptr;
  //     TOPK_SWITCH(top_k, TOPK, [&]{
  //       BLOCK_M_SWITCH(BLOCK_SIZE_M, BLOCK_SIZE_M_, [&]{
  //         BLOCK_K_SWITCH(BLOCK_SIZE_K, BLOCK_SIZE_K_, [&]{
  //             // BOOL_SWITCH(b_qzeros.has_value(), has_zp, [&]{
  //               BOOL_SWITCH(topk_weights.has_value(), mul_topk_weight, [&]{
  //                 GROUP_SIZE_SWITCH(group_size, GROUP_SIZE, [&]{
  //                   // run_moe_wna16_gemm<half, bit, top_k, BLOCK_SIZE_M_, BLOCK_SIZE_N, BLOCK_SIZE_K, true, mul_topk_weight, group_size>(
  //                     run_moe_wna16_gemm_base<__hip_bfloat16, 4, TOPK, BLOCK_SIZE_M_, 256, BLOCK_SIZE_K_, true, mul_topk_weight, GROUP_SIZE>(
  //                     (const __hip_bfloat16*)input.data_ptr<at::BFloat16>(),
  //                     (float*)output_fp32.data_ptr(),
  //                     (const uint32_t*)b_qweight.data_ptr<uint8_t>(),
  //                     (const __hip_bfloat16*)b_scales.data_ptr<at::BFloat16>(), 
  //                     b_qzeros_ptr,
  //                     topk_weights_ptr, 
  //                     sorted_token_ids.data_ptr<int32_t>(),
  //                     expert_ids.data_ptr<int32_t>(), 
  //                     num_tokens_post_pad.data_ptr<int32_t>(), 
  //                     num_token_blocks, 
  //                     size_m, 
  //                     size_n,
  //                     size_k
  //                   );
  //                 });
  //                 });
  //             });
  //         });
  //     });
  // } 
  // else {
  //   TORCH_CHECK(false, "moe_wna16_gemm not supports");
  // }
  if (input.scalar_type() == at::ScalarType::Half) {
    output.copy_(output_fp32.to(torch::kFloat16));
  } else if (input.scalar_type() == at::ScalarType::BFloat16) {
    output.copy_(output_fp32.to(torch::kBFloat16));  // 转换为 BF16
  }
  return output;
}

torch::Tensor moe_c_moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
                             torch::Tensor b_qweight, torch::Tensor b_scales,
                             std::optional<torch::Tensor> b_qzeros,
                             std::optional<torch::Tensor> topk_weights,
                             torch::Tensor sorted_token_ids,
                             torch::Tensor expert_ids,
                             torch::Tensor num_tokens_post_pad, int64_t top_k,
                             int64_t BLOCK_SIZE_m, int64_t BLOCK_SIZE_n,
                             int64_t BLOCK_SIZE_k, int64_t kloops, int64_t nloops, int64_t bit) {
 
  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
  // auto options = torch::TensorOptions().dtype(input.dtype()).device(input.device());
  const int size_m = input.size(0);
  const int size_n = b_qweight.size(3);
  const int size_k = input.size(1);
  const int group_size = size_k / b_scales.size(2);

  int64_t BLOCK_SIZE_N = std::min(64, size_n);
  // std::cout<<"BLOCK_SIZE_N IS : "<<BLOCK_SIZE_N<<std::endl;
  int64_t BLOCK_SIZE_K_MIN =4*1024/BLOCK_SIZE_m;
  int64_t BLOCK_SIZE_K_MAX =8*1024/BLOCK_SIZE_m;
  // std::cout<<"BLOCK_SIZE_K_MAX IS : "<<BLOCK_SIZE_K_MAX<<std::endl;
  int64_t BLOCK_SIZE_K = std::min(BLOCK_SIZE_K_MAX,size_k);
  // std::cout<<"BLOCK_SIZE_K IS : "<<BLOCK_SIZE_K<<std::endl;

  // int BLOCK_SIZE_M_MAX = std::min(32, size_m);
  // std::cout<<"BLOCK_SIZE_M_MAX IS : "<<BLOCK_SIZE_M_MAX<<std::endl;
  int BLOCK_SIZE_N_MAX_roofline = 256;
  int BLOCK_SIZE_N_MAX = std::min(BLOCK_SIZE_N_MAX_roofline, size_n);
  // std::cout<<"BLOCK_SIZE_N_MAX IS : "<<BLOCK_SIZE_N_MAX<<std::endl;

  BLOCK_SIZE_K = std::min(128, BLOCK_SIZE_K);
  // std::cout<<"BLOCK_SIZE_K IS : "<<BLOCK_SIZE_K<<std::endl;
  // printf("__hip_bfloat16_size == %d \n",sizeof(__hip_bfloat16));
  int block_size_m_loops = 1;// std::min(1,BLOCK_SIZE_M_MAX/BLOCK_SIZE_m);
  int block_size_k_loops = kloops;
  int block_size_n_loops = nloops;
  // half_t * d_w_out_half =nullptr;
  // __hip_bfloat16 * d_w_out_bf =nullptr;
  // float * float_d_out = nullptr;
  // hipMalloc(&d_w_out_half, 64 * 8 * sizeof(half_t));
  // hipMalloc(&d_w_out_bf, 64 * 8 * sizeof(__hip_bfloat16));
  // hipMalloc(&float_d_out, 64 * 8 * sizeof(float));
  int64_t EM = sorted_token_ids.size(0);
  
  if (size_m <= BLOCK_SIZE_m) {
    EM = min(EM, size_m * BLOCK_SIZE_m * top_k);
  }
  const int num_token_blocks = (EM + BLOCK_SIZE_m*block_size_m_loops - 1) / (BLOCK_SIZE_m*block_size_m_loops);


  const uint32_t* b_qzeros_ptr;
  if (b_qzeros.has_value())
    b_qzeros_ptr = (const uint32_t*)b_qzeros.value().data_ptr<uint8_t>();
  const float* topk_weights_ptr;
  if (topk_weights.has_value())
    topk_weights_ptr = (const float*)topk_weights.value().data_ptr();

  int groups_per_block_row = BLOCK_SIZE_K / group_size;
  TORCH_CHECK(bit == 4 || bit == 8, "bit must be 4 or 8");
  TORCH_CHECK(size_k % BLOCK_SIZE_K == 0,
              "size_k must divisible by BLOCK_SIZE_K");
  TORCH_CHECK(BLOCK_SIZE_K % group_size == 0,
              "BLOCK_SIZE_K must divisible by group_size");
  TORCH_CHECK(BLOCK_SIZE_m <= 64, "BLOCK_SIZE_m must less or equal to 64");
  TORCH_CHECK(groups_per_block_row == 1 || groups_per_block_row == 2 ||
                  groups_per_block_row == 4 || groups_per_block_row == 8,
              "BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8]");
  
  bool use_atomic = (size_k != BLOCK_SIZE_K*block_size_k_loops);

  std::optional<torch::Tensor> output_fp32;

  if (use_atomic){

    output_fp32 = torch::zeros(output.sizes(),output.options().dtype(torch::kFloat32));
    // output_fp32->zero_(); 
  } 

  float milliseconds = 0;
  cudaEvent_t start, stop;
  const char* find_best = std::getenv("WHICH_TO_TEST");
  if (find_best) {
    cudaEventCreate(&start);
    cudaEventCreate(&stop);
    cudaEventRecord(start);        // 记录开始
  }

  // if (input.scalar_type() == at::ScalarType::Half) {
  //   BIT_SWITCH(bit, BIT, [&]{
  //     TOPK_SWITCH(top_k, TOPK, [&]{
  //       BLOCK_M_SWITCH(BLOCK_SIZE_m, BLOCK_SIZE_M_, [&]{
  //         BLOCK_N_SWITCH(BLOCK_SIZE_N, BLOCK_SIZE_N_, [&]{
  //           BLOCK_K_SWITCH(BLOCK_SIZE_K, BLOCK_SIZE_K_, [&]{
  //             // BOOL_SWITCH(b_qzeros.has_value(), has_zp, [&]{
  //               BOOL_SWITCH(topk_weights.has_value(), mul_topk_weight, [&]{
  //                 GROUP_SIZE_K_SWITCH(group_size, GROUP_SIZE_K, [&]{
  //                   BLOCK_SIZE_M_LOOPS_SWITCH(block_size_m_loops , BLOCK_SIZE_M_LOOPS, [&]{
  //                     BLOCK_SIZE_N_LOOPS_SWITCH(block_size_n_loops , BLOCK_SIZE_N_LOOPS, [&]{
  //                       BLOCK_SIZE_K_LOOPS_SWITCH(block_size_k_loops , BLOCK_SIZE_K_LOOPS, [&]{
  //                         BOOL_SWITCH(use_atomic , USE_ATOMIC, [&]{
  //                   run_moe_wna16_gemm<half, 4, TOPK, BLOCK_SIZE_M_, BLOCK_SIZE_N_, BLOCK_SIZE_K_, true, mul_topk_weight, GROUP_SIZE_K, BLOCK_SIZE_M_LOOPS, BLOCK_SIZE_N_LOOPS, BLOCK_SIZE_K_LOOPS, USE_ATOMIC,256>(
  //                     (const half*)input.data_ptr<at::Half>(),
  //                     // (const half*)d_input,
  //                     use_atomic ?(float*)output_fp32->data_ptr():(float*)output.data_ptr(),
  //                     // (float*)output_fp32->data_ptr(),
  //                     (const uint32_t*)b_qweight.data_ptr<uint32_t>(),
  //                     // (const uint32_t*)d_w_test,
  //                     // ( half_t*) d_w_out_half, /*for debug*/
  //                     // ( float*) float_d_out, /*for debug*/
  //                     (const half*)b_scales.data_ptr<at::Half>(), 
  //                     // (const half*)d_scale, 
  //                     b_qzeros_ptr,
  //                     // (const uint32_t*)d_scale,
  //                     topk_weights_ptr, 
  //                     sorted_token_ids.data_ptr<int32_t>(),
  //                     expert_ids.data_ptr<int32_t>(), 
  //                     num_tokens_post_pad.data_ptr<int32_t>(), 
  //                     // num_tokens_post_pad_value<int32_t>(),
  //                     // num_tokens_post_pad_data_ptr[0],
  //                     num_token_blocks, 
  //                     size_m, 
  //                     size_n,
  //                     size_k
  //                     ); // kernel-1 mma    
  //                   });              
  //                   });
  //                 });
  //               });
  //             });
  //           });
  //         });
  //       });
  //     });
  //   });
  // });
  // } 
  //   else if (input.scalar_type() == at::ScalarType::BFloat16) {
  //     BIT_SWITCH(bit, BIT, [&]{
  //     TOPK_SWITCH(top_k, TOPK, [&]{
  //       BLOCK_M_SWITCH(BLOCK_SIZE_m, BLOCK_SIZE_M_, [&]{
  //         BLOCK_N_SWITCH(BLOCK_SIZE_N, BLOCK_SIZE_N_, [&]{
  //           BLOCK_K_SWITCH(BLOCK_SIZE_K, BLOCK_SIZE_K_, [&]{
  //             // BOOL_SWITCH(b_qzeros.has_value(), has_zp, [&]{
  //               BOOL_SWITCH(topk_weights.has_value(), mul_topk_weight, [&]{
  //                 GROUP_SIZE_K_SWITCH(group_size, GROUP_SIZE_K, [&]{
  //                   BLOCK_SIZE_M_LOOPS_SWITCH(block_size_m_loops , BLOCK_SIZE_M_LOOPS, [&]{
  //                     BLOCK_SIZE_N_LOOPS_SWITCH(block_size_n_loops , BLOCK_SIZE_N_LOOPS, [&]{
  //                       BLOCK_SIZE_K_LOOPS_SWITCH(block_size_k_loops , BLOCK_SIZE_K_LOOPS, [&]{
  //                         BOOL_SWITCH(use_atomic , USE_ATOMIC, [&]{
  //                   run_moe_wna16_gemm<__hip_bfloat16, 4, TOPK, BLOCK_SIZE_M_, BLOCK_SIZE_N_, BLOCK_SIZE_K_, true, mul_topk_weight, GROUP_SIZE_K, BLOCK_SIZE_M_LOOPS, BLOCK_SIZE_N_LOOPS, BLOCK_SIZE_K_LOOPS, USE_ATOMIC,256>(
  //                     (const __hip_bfloat16*)input.data_ptr<at::BFloat16>(),
  //                     // (const half*)d_input,
  //                     use_atomic ?(float*)output_fp32->data_ptr():(float*)output.data_ptr(),
  //                     // (float*)output_fp32->data_ptr(),
  //                     (const uint32_t*)b_qweight.data_ptr<uint32_t>(),
  //                     // (const uint32_t*)d_w_test,
  //                     // ( __hip_bfloat16*) d_w_out_bf, /*for debug*/
  //                     // ( float*) float_d_out, /*for debug*/
  //                     (const __hip_bfloat16*)b_scales.data_ptr<at::BFloat16>(), 
  //                     // (const half*)d_scale, 
  //                     b_qzeros_ptr,
  //                     // (const uint32_t*)d_scale,
  //                     topk_weights_ptr, 
  //                     sorted_token_ids.data_ptr<int32_t>(),
  //                     expert_ids.data_ptr<int32_t>(), 
  //                     num_tokens_post_pad.data_ptr<int32_t>(), 
  //                     // num_tokens_post_pad_value<int32_t>(),
  //                     // num_tokens_post_pad_data_ptr[0],
  //                     num_token_blocks, 
  //                     size_m, 
  //                     size_n,
  //                     size_k
  //                     ); // kernel-1 mma    
  //                   });              
  //                   });
  //                 });
  //               });
  //             });
  //           });
  //         });
  //       });
  //     });
  //   });
  // });
  // } 
  // else {
  //   TORCH_CHECK(false, "moe_w8a16_gemm_awq only supports float16");
  // }

  if (find_best) {
    cudaEventRecord(stop);         // 记录结束
    cudaEventSynchronize(stop);    // 等待 kernel 执行完成

    
    cudaEventElapsedTime(&milliseconds, start, stop); // 计算时间

    /* std::cout << "kernel 1 time: " << milliseconds << " ms" << std::endl; */
    
    cudaEventDestroy(start);
    cudaEventDestroy(stop);

    std::ofstream ofs("./w4a16_kernel_1_timecost", std::ios::app); // 追加写入
    if (ofs.is_open()) {
        ofs << milliseconds << std::endl;
        ofs.close();
    }
  }

  if (use_atomic){
      if (input.scalar_type() == at::ScalarType::Half) {
        output.copy_(output_fp32->to(torch::kFloat16));
      } else if (input.scalar_type() == at::ScalarType::BFloat16) {
        output.copy_(output_fp32->to(torch::kBFloat16));  // 转换为 BF16
      }
  }
  return output;
}

torch::Tensor moe_c_moe_wna16_gemm_2(torch::Tensor input, torch::Tensor output,
                             torch::Tensor b_qweight, torch::Tensor b_scales,
                             std::optional<torch::Tensor> b_qzeros,
                             std::optional<torch::Tensor> topk_weights,
                             torch::Tensor sorted_token_ids,
                             torch::Tensor expert_ids,
                             torch::Tensor num_tokens_post_pad, int64_t top_k,
                             int64_t BLOCK_SIZE_m, int64_t BLOCK_SIZE_n,
                             int64_t BLOCK_SIZE_k, int64_t kloops, int64_t nloops, int64_t bit) {
 
  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
  // auto options = torch::TensorOptions().dtype(input.dtype()).device(input.device());
  const int size_m = input.size(0);
  const int size_n = b_qweight.size(1);
  const int size_k = input.size(1);
  const int group_size = size_k / b_scales.size(2);

  int64_t BLOCK_SIZE_N = std::min(64, size_n);
  // std::cout<<"BLOCK_SIZE_N IS : "<<BLOCK_SIZE_N<<std::endl;
  int64_t BLOCK_SIZE_K_MIN =4*1024/BLOCK_SIZE_m;
  int64_t BLOCK_SIZE_K_MAX =8*1024/BLOCK_SIZE_m;
  // std::cout<<"BLOCK_SIZE_K_MAX IS : "<<BLOCK_SIZE_K_MAX<<std::endl;
  int64_t BLOCK_SIZE_K = std::min(BLOCK_SIZE_K_MAX,size_k);
  // std::cout<<"BLOCK_SIZE_K IS : "<<BLOCK_SIZE_K<<std::endl;

  // int BLOCK_SIZE_M_MAX = std::min(32, size_m);
  // std::cout<<"BLOCK_SIZE_M_MAX IS : "<<BLOCK_SIZE_M_MAX<<std::endl;
  int BLOCK_SIZE_N_MAX_roofline = 256;
  int BLOCK_SIZE_N_MAX = std::min(BLOCK_SIZE_N_MAX_roofline, size_n);
  // std::cout<<"BLOCK_SIZE_N_MAX IS : "<<BLOCK_SIZE_N_MAX<<std::endl;

  BLOCK_SIZE_K = std::min(128, BLOCK_SIZE_K);

  int block_size_m_loops = 1;// std::min(1,BLOCK_SIZE_M_MAX/BLOCK_SIZE_m);
  int block_size_k_loops = kloops;
  int block_size_n_loops = nloops;
  int64_t EM = sorted_token_ids.size(0);
  
  if (size_m <= BLOCK_SIZE_m) {
    EM = min(EM, size_m * BLOCK_SIZE_m * top_k);
  }
  const int num_token_blocks = (EM + BLOCK_SIZE_m*block_size_m_loops - 1) / (BLOCK_SIZE_m*block_size_m_loops);

  const uint32_t* b_qzeros_ptr;
  if (b_qzeros.has_value())
    b_qzeros_ptr = (const uint32_t*)b_qzeros.value().data_ptr<uint8_t>();
  const float* topk_weights_ptr;
  if (topk_weights.has_value())
    topk_weights_ptr = (const float*)topk_weights.value().data_ptr();

  int groups_per_block_row = BLOCK_SIZE_K / group_size;
  TORCH_CHECK(bit == 4 || bit == 8, "bit must be 4 or 8");
  TORCH_CHECK(size_k % BLOCK_SIZE_K == 0,
              "size_k must divisible by BLOCK_SIZE_K");
  TORCH_CHECK(BLOCK_SIZE_K % group_size == 0,
              "BLOCK_SIZE_K must divisible by group_size");
  TORCH_CHECK(BLOCK_SIZE_m <= 64, "BLOCK_SIZE_m must less or equal to 64");
  TORCH_CHECK(groups_per_block_row == 1 || groups_per_block_row == 2 ||
                  groups_per_block_row == 4 || groups_per_block_row == 8,
              "BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8]");
  
  bool use_atomic = (size_k != BLOCK_SIZE_K*block_size_k_loops);

  std::optional<torch::Tensor> output_fp32;

  if (use_atomic){

    output_fp32 = torch::zeros(output.sizes(),output.options().dtype(torch::kFloat32));
    // output_fp32->zero_(); 
  } 

  float milliseconds = 0;
  cudaEvent_t start, stop;
  const char* find_best = std::getenv("WHICH_TO_TEST");
  if (find_best) {
    cudaEventCreate(&start);
    cudaEventCreate(&stop);
    cudaEventRecord(start);        // 记录开始
  }

  // if (input.scalar_type() == at::ScalarType::Half) {
  //   BIT_SWITCH(bit, BIT, [&]{
  //     TOPK_SWITCH(top_k, TOPK, [&]{
  //       BLOCK_M_SWITCH(BLOCK_SIZE_m, BLOCK_SIZE_M_, [&]{
  //         BLOCK_N_SWITCH(BLOCK_SIZE_N, BLOCK_SIZE_N_, [&]{
  //           BLOCK_K_SWITCH(BLOCK_SIZE_K, BLOCK_SIZE_K_, [&]{
  //             // BOOL_SWITCH(b_qzeros.has_value(), has_zp, [&]{
  //               BOOL_SWITCH(topk_weights.has_value(), mul_topk_weight, [&]{
  //                 GROUP_SIZE_K_SWITCH(group_size, GROUP_SIZE_K, [&]{
  //                   BLOCK_SIZE_M_LOOPS_SWITCH(block_size_m_loops , BLOCK_SIZE_M_LOOPS, [&]{
  //                     BLOCK_SIZE_N_LOOPS_SWITCH(block_size_n_loops , BLOCK_SIZE_N_LOOPS, [&]{
  //                       BLOCK_SIZE_K_LOOPS_SWITCH(block_size_k_loops , BLOCK_SIZE_K_LOOPS, [&]{
  //                         BOOL_SWITCH(use_atomic , USE_ATOMIC, [&]{
  //                   run_moe_wna16_gemm_2<half, 4, TOPK, BLOCK_SIZE_M_, BLOCK_SIZE_N_, BLOCK_SIZE_K_, true, mul_topk_weight, GROUP_SIZE_K, BLOCK_SIZE_M_LOOPS, BLOCK_SIZE_N_LOOPS, BLOCK_SIZE_K_LOOPS, USE_ATOMIC,256>(
  //                     (const half*)input.data_ptr<at::Half>(),
  //                     // (const half*)d_input,
  //                     use_atomic ?(float*)output_fp32->data_ptr():(float*)output.data_ptr(),
  //                     // (float*)output_fp32->data_ptr(),
  //                     (const uint32_t*)b_qweight.data_ptr<uint8_t>(),
  //                     // (const uint32_t*)d_w_test,
  //                     // ( half_t*) d_w_out_half, /*for debug*/
  //                     // ( float*) float_d_out, /*for debug*/
  //                     (const half*)b_scales.data_ptr<at::Half>(), 
  //                     // (const half*)d_scale, 
  //                     b_qzeros_ptr,
  //                     // (const uint32_t*)d_scale,
  //                     topk_weights_ptr, 
  //                     sorted_token_ids.data_ptr<int32_t>(),
  //                     expert_ids.data_ptr<int32_t>(), 
  //                     num_tokens_post_pad.data_ptr<int32_t>(), 
  //                     // num_tokens_post_pad_value<int32_t>(),
  //                     // num_tokens_post_pad_data_ptr[0],
  //                     num_token_blocks, 
  //                     size_m, 
  //                     size_n,
  //                     size_k
  //                     ); // kernel-1 mma    
  //                   });              
  //                   });
  //                 });
  //               });
  //             });
  //           });
  //         });
  //       });
  //     });
  //   });
  // });
  // } 
  //   else if (input.scalar_type() == at::ScalarType::BFloat16) {
  //     BIT_SWITCH(bit, BIT, [&]{
  //     TOPK_SWITCH(top_k, TOPK, [&]{
  //       BLOCK_M_SWITCH(BLOCK_SIZE_m, BLOCK_SIZE_M_, [&]{
  //         BLOCK_N_SWITCH(BLOCK_SIZE_N, BLOCK_SIZE_N_, [&]{
  //           BLOCK_K_SWITCH(BLOCK_SIZE_K, BLOCK_SIZE_K_, [&]{
  //             // BOOL_SWITCH(b_qzeros.has_value(), has_zp, [&]{
  //               BOOL_SWITCH(topk_weights.has_value(), mul_topk_weight, [&]{
  //                 GROUP_SIZE_K_SWITCH(group_size, GROUP_SIZE_K, [&]{
  //                   BLOCK_SIZE_M_LOOPS_SWITCH(block_size_m_loops , BLOCK_SIZE_M_LOOPS, [&]{
  //                     BLOCK_SIZE_N_LOOPS_SWITCH(block_size_n_loops , BLOCK_SIZE_N_LOOPS, [&]{
  //                       BLOCK_SIZE_K_LOOPS_SWITCH(block_size_k_loops , BLOCK_SIZE_K_LOOPS, [&]{
  //                         BOOL_SWITCH(use_atomic , USE_ATOMIC, [&]{
  //                   run_moe_wna16_gemm_2<__hip_bfloat16, 4, TOPK, BLOCK_SIZE_M_, BLOCK_SIZE_N_, BLOCK_SIZE_K_, true, mul_topk_weight, GROUP_SIZE_K, BLOCK_SIZE_M_LOOPS, BLOCK_SIZE_N_LOOPS, BLOCK_SIZE_K_LOOPS, USE_ATOMIC,256>(
  //                     (const __hip_bfloat16*)input.data_ptr<at::BFloat16>(),
  //                     // (const half*)d_input,
  //                     use_atomic ?(float*)output_fp32->data_ptr():(float*)output.data_ptr(),
  //                     // (float*)output_fp32->data_ptr(),
  //                     (const uint32_t*)b_qweight.data_ptr<uint8_t>(),
  //                     // (const uint32_t*)d_w_test,
  //                     // ( __hip_bfloat16*) d_w_out_bf, /*for debug*/
  //                     // ( float*) float_d_out, /*for debug*/
  //                     (const __hip_bfloat16*)b_scales.data_ptr<at::BFloat16>(), 
  //                     // (const half*)d_scale, 
  //                     b_qzeros_ptr,
  //                     // (const uint32_t*)d_scale,
  //                     topk_weights_ptr, 
  //                     sorted_token_ids.data_ptr<int32_t>(),
  //                     expert_ids.data_ptr<int32_t>(), 
  //                     num_tokens_post_pad.data_ptr<int32_t>(), 
  //                     // num_tokens_post_pad_value<int32_t>(),
  //                     // num_tokens_post_pad_data_ptr[0],
  //                     num_token_blocks, 
  //                     size_m, 
  //                     size_n,
  //                     size_k
  //                     ); // kernel-1 mma    
  //                   });              
  //                   });
  //                 });
  //               });
  //             });
  //           });
  //         });
  //       });
  //     });
  //   });
  // });
  // } 
  // else {
  //   TORCH_CHECK(false, "moe_w8a16_gemm_awq only supports float16");
  // }

  if (find_best) {
    cudaEventRecord(stop);         // 记录结束
    cudaEventSynchronize(stop);    // 等待 kernel 执行完成

    
    cudaEventElapsedTime(&milliseconds, start, stop); // 计算时间

    /* std::cout << "kernel 1 time: " << milliseconds << " ms" << std::endl; */
    
    cudaEventDestroy(start);
    cudaEventDestroy(stop);

    std::ofstream ofs("./w4a16_kernel_2_timecost", std::ios::app); // 追加写入
    if (ofs.is_open()) {
        ofs << milliseconds << std::endl;
        ofs.close();
    }
  }

  if (use_atomic){
      if (input.scalar_type() == at::ScalarType::Half) {
        output.copy_(output_fp32->to(torch::kFloat16));
      } else if (input.scalar_type() == at::ScalarType::BFloat16) {
        output.copy_(output_fp32->to(torch::kBFloat16));  // 转换为 BF16
      }
  }
  return output;
}

void moe_c_silu_and_mul(torch::Tensor& out,    // [..., d]
                  torch::Tensor& input,  // [..., 2 * d]
                  int64_t rows_per_block = 1,
                  int64_t vec_size = 2)
{
  moe_c::silu_and_mul(out, input, static_cast<int>(rows_per_block),
                      static_cast<int>(vec_size));
}