#pragma once
// SPDX-License-Identifier: MIT
// Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
#include "aiter_enum.h"
#include <torch/extension.h>

void biased_grouped_topk(torch::Tensor& gating_output,   // [num_tokens, num_experts]
                         torch::Tensor& correction_bias, // [num_expert]
                         torch::Tensor& topk_weights,    // [num_tokens, topk]
                         torch::Tensor& topk_ids,        // [num_tokens, topk]
                         int num_expert_group,
                         int topk_group,
                         bool renormalize,
                         const float routed_scaling_factor = 1.);

void grouped_topk(torch::Tensor& gating_output, // [num_tokens, num_experts]
                  torch::Tensor& topk_weights,  // [num_tokens, topk]
                  torch::Tensor& topk_ids,      // [num_tokens, topk]
                  int num_expert_group,
                  int topk_grp,
                  bool need_renorm,
                  bool is_softmax                   = true,
                  const float routed_scaling_factor = 1.);

std::vector<at::Tensor> moe_fused_gate(at::Tensor& input,
                                       at::Tensor& bias,
                                       at::Tensor& topk_weights,
                                       at::Tensor& topk_ids,
                                       int64_t num_expert_group,
                                       int64_t topk_group,
                                       int64_t topk,
                                       int64_t n_share_experts_fusion,
                                       double routed_scaling_factor);

void moe_align_block_size(torch::Tensor topk_ids,
                          int64_t num_experts,
                          int64_t block_size,
                          torch::Tensor sorted_token_ids,
                          torch::Tensor experts_ids,
                          torch::Tensor token_nums,
                          torch::Tensor num_tokens_post_pad);

void fmoe(torch::Tensor& out,               // [token_cnt, dim]
          torch::Tensor& input,             // [token_cnt, dim] M,K
          torch::Tensor& gate,              // [expert, hidden_dim, dim] N,K
          torch::Tensor& down,              // [expert, hidden_dim, dim]
          torch::Tensor& sorted_token_ids,  // [max_num_tokens_padded]
          torch::Tensor& sorted_weights,    // [max_num_tokens_padded]
          torch::Tensor& sorted_expert_ids, // [max_num_m_blocks]
          torch::Tensor& num_valid_ids,     // [1]
          uint32_t topk                     //
);

void fmoe_int8_g1u0(torch::Tensor& out,               // [token_cnt, dim]
                    torch::Tensor& input,             // [token_cnt, dim] M,K
                    torch::Tensor& gate,              // [expert, hidden_dim, dim] N,K
                    torch::Tensor& down,              // [expert, hidden_dim, dim]
                    torch::Tensor& sorted_token_ids,  // [max_num_tokens_padded]
                    torch::Tensor& sorted_weights,    // [max_num_tokens_padded]
                    torch::Tensor& sorted_expert_ids, // [max_num_m_blocks]
                    torch::Tensor& num_valid_ids,     // [1]
                    uint32_t topk,                    //
                    torch::Tensor& input_scale,       // [token_cnt, 1]
                    torch::Tensor& fc1_scale,         // [expert, 1, hidden_dim]
                    torch::Tensor& fc2_scale,         // [expert, 1, dim]
                    torch::Tensor& fc2_smooth_scale,  // [expert, 1, hidden_dim]
                    ActivationType activation = ActivationType::Silu);

void fmoe_g1u1(
    torch::Tensor& out,               // [token_cnt, dim]
    torch::Tensor& input,             // [token_cnt, dim] M,K
    torch::Tensor& gate,              // [expert, hidden_dim*2, dim] N,K
    torch::Tensor& down,              // [expert, hidden_dim, dim]
    torch::Tensor& sorted_token_ids,  // [max_num_tokens_padded]
    torch::Tensor& sorted_weights,    // [max_num_tokens_padded]
    torch::Tensor& sorted_expert_ids, // [max_num_m_blocks]
    torch::Tensor& num_valid_ids,     // [1]
    uint32_t topk,                    //
    torch::Tensor& input_scale,       // [token_cnt, 1]
    torch::Tensor& fc1_scale,         // [expert, 1, hidden_dim]
    torch::Tensor& fc2_scale,         // [expert, 1, dim]
    std::string& kernel_name,
    std::optional<torch::Tensor> fc2_smooth_scale = std::nullopt, // [expert, 1, hidden_dim]
    ActivationType activation                     = ActivationType::Silu);

void fmoe_g1u1_tkw1(
    torch::Tensor& out,               // [token_cnt, dim]
    torch::Tensor& input,             // [token_cnt, dim] M,K
    torch::Tensor& gate,              // [expert, hidden_dim*2, dim] N,K
    torch::Tensor& down,              // [expert, hidden_dim, dim]
    torch::Tensor& sorted_token_ids,  // [max_num_tokens_padded]
    torch::Tensor& sorted_weights,    // [max_num_tokens_padded]
    torch::Tensor& sorted_expert_ids, // [max_num_m_blocks]
    torch::Tensor& num_valid_ids,     // [1]
    uint32_t topk,                    //
    torch::Tensor& input_scale,       // [token_cnt, 1]
    torch::Tensor& fc1_scale,         // [expert, 1, hidden_dim]
    torch::Tensor& fc2_scale,         // [expert, 1, dim]
    std::string& kernel_name,
    std::optional<torch::Tensor> fc2_smooth_scale = std::nullopt, // [expert, 1, hidden_dim]
    ActivationType activation                     = ActivationType::Silu);

void fmoe_int8_g1u0_a16(torch::Tensor& out,               // [token_cnt, dim]
                        torch::Tensor& input,             // [token_cnt, dim] M,K
                        torch::Tensor& gate,              // [expert, inter_dim, dim] N,K
                        torch::Tensor& down,              // [expert, dim, inter_dim]
                        torch::Tensor& sorted_token_ids,  // [max_num_tokens_padded]
                        torch::Tensor& sorted_weights,    // [max_num_tokens_padded]
                        torch::Tensor& sorted_expert_ids, // [max_num_m_blocks]
                        torch::Tensor& num_valid_ids,     // [1]
                        uint32_t topk,                    //
                        torch::Tensor& fc1_scale,         // [expert, 1, hidden_dim]
                        torch::Tensor& fc2_scale,         // [expert, 1, dim]
                        torch::Tensor& fc1_smooth_scale,  // [expert, 1, hidden_dim]
                        torch::Tensor& fc2_smooth_scale   // [expert, 1, hidden_dim]
);
void fmoe_g1u1_a16(torch::Tensor& out,               // [token_cnt, dim]
                   torch::Tensor& input,             // [token_cnt, dim] M,K
                   torch::Tensor& gate,              // [expert, inter_dim, dim] N,K
                   torch::Tensor& down,              // [expert, dim, inter_dim]
                   torch::Tensor& sorted_token_ids,  // [max_num_tokens_padded]
                   torch::Tensor& sorted_weights,    // [max_num_tokens_padded]
                   torch::Tensor& sorted_expert_ids, // [max_num_m_blocks]
                   torch::Tensor& num_valid_ids,     // [1]
                   uint32_t topk,                    //
                   torch::Tensor& fc1_scale,         // [expert, 1, hidden_dim]
                   torch::Tensor& fc2_scale,         // [expert, 1, dim]
                   torch::Tensor& fc1_smooth_scale,  // [expert, 1, hidden_dim]
                   torch::Tensor& fc2_smooth_scale,  // [expert, 1, hidden_dim]
                   ActivationType activation);

void fmoe_fp8_blockscale_g1u1(
    torch::Tensor& out,               // [token_cnt, dim]
    torch::Tensor& input,             // [token_cnt, dim] M,K
    torch::Tensor& gate,              // [expert, inter_dim*2, dim] N,K
    torch::Tensor& down,              // [expert, dim, inter_dim]
    torch::Tensor& sorted_token_ids,  // [max_num_tokens_padded]
    torch::Tensor& sorted_weights,    // [max_num_tokens_padded]
    torch::Tensor& sorted_expert_ids, // [max_num_m_blocks]
    torch::Tensor& num_valid_ids,     // [1]
    uint32_t topk,                    //
    torch::Tensor& input_scale,       // [expert, 1, dim]
    torch::Tensor& fc1_scale,         // [expert, 1, inter_dim]
    torch::Tensor& fc2_scale,         // [expert, 1, dim]
    std::string& kernel_name,
    int fc_scale_blkn                             = 128,          // = 128,
    int fc_scale_blkk                             = 128,          // = 128
    std::optional<torch::Tensor> fc2_smooth_scale = std::nullopt, // [expert, 1, inter_dim]
    ActivationType activation                     = ActivationType::Silu);

void moe_stage1_g1u1(
    torch::Tensor& input,             // [token_cnt, model_dim] M,K
    torch::Tensor& w1,                // [expert, inter_dim*2, model_dim] N,K
    torch::Tensor& w2,                // [expert, model_dim, inter_dim]
    torch::Tensor& sorted_token_ids,  // [max_num_tokens_padded]
    torch::Tensor& sorted_expert_ids, // [max_num_m_blocks]
    torch::Tensor& num_valid_ids,     // [1]
    torch::Tensor& out,               // [token_cnt, topk, inter_dim]
    int inter_dim,
    std::string& kernelName,
    int block_m,
    int ksplit,
    ActivationType activation,
    QuantType quant_type,
    std::optional<torch::Tensor> a1_scale, // [token_cnt, 1], token scale
    std::optional<torch::Tensor> w1_scale, // [expert, 1, inter_dim], gate(up) scale
    std::optional<torch::Tensor> sorted_weights);

void topk_softmax_asm(torch::Tensor& topk_weights,         // [num_tokens, topk]
                      torch::Tensor& topk_indices,         // [num_tokens, topk]
                      torch::Tensor& token_expert_indices, // [num_tokens, topk]
                      torch::Tensor& gating_output,        // [num_tokens, num_experts]
                      bool need_renorm);

namespace aiter {

void topk_softmax(torch::Tensor& topk_weights,
                  torch::Tensor& topk_indices,
                  torch::Tensor& token_expert_indices,
                  torch::Tensor& gating_output,
                  bool need_renorm);

void moe_align_block_size(torch::Tensor topk_ids,
                          int64_t num_experts,
                          int64_t block_size,
                          torch::Tensor sorted_token_ids,
                          torch::Tensor experts_ids,
                          torch::Tensor token_nums,
                          torch::Tensor num_tokens_post_pad);

void moe_sum(torch::Tensor& input, torch::Tensor& output);

void topk_sigmoid(torch::Tensor topk_weights,   // [tokens, topk]
                  torch::Tensor topk_indices,   // [tokens, topk]
                  torch::Tensor gating_output); // [tokens, experts]

} // namespace aiter
