# SPDX-License-Identifier: MIT
# Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.

import torch
from torch import Tensor
from typing import Optional,List
from ..jit.core import (
    compile_ops,
)
from .enum import ActivationType, Enum, QuantType



@compile_ops("module_moe_c_kernel")
def moe_c_moe_gemm_marlin_w8a8( 
    input: torch.Tensor,
    b_qweight : torch.Tensor,
    output : torch.Tensor,
    a_scale: torch.Tensor,
    b_scale : torch.Tensor,
    topk_weights : Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids : torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    top_k : int,
    mode :int,
    delta: int)-> torch.Tensor: 
    """
    ---------------------------------------------------------------
    # MoE 场景下 8bit 量化的 GEMM 计算（Marlin 优化版）   

    ## 关键前置条件
    必须配合对应的权重 Shuffle 函数使用，否则会导致计算结果完全错误：
    - GEMM1 场景：使用 ops.marlin_weights 处理权重
    - GEMM2 场景：使用 ops.marlin_weights_ours 处理权重




    
    ---------------------------------------------------------------
    """
    
    pass


@compile_ops("module_moe_c_kernel")
def moe_c_moe_gemm_marlin_w4a8( 
    input: torch.Tensor,
    b_qweight : torch.Tensor,
    output : torch.Tensor,
    a_scale: torch.Tensor,
    b_scale : torch.Tensor,
    topk_weights : Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids : torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    top_k : int,
    mode :int,
    delta: int)-> torch.Tensor: 
    """
    ---------------------------------------------------------------
    # MoE 场景下 8bit 量化的 GEMM 计算（Marlin 优化版）   

    ## 关键前置条件
    必须配合对应的权重 Shuffle 函数使用，否则会导致计算结果完全错误：
    - GEMM1 场景：使用 ops.marlin_weights 处理权重
    - GEMM2 场景：使用 ops.marlin_weights_ours 处理权重




    
    ---------------------------------------------------------------
    """
    
    pass
      
    

@compile_ops("module_moe_c_kernel")
def moe_c_moe_gemm_marlin_w8a8_fp8( 
    input: torch.Tensor,
    b_qweight : torch.Tensor,
    output : torch.Tensor,
    a_scale: torch.Tensor,
    b_scale : torch.Tensor,
    topk_weights : Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids : torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    top_k : int,
    mode :int,
    delta: int)-> torch.Tensor: 
    """
    ---------------------------------------------------------------
    # MoE 场景下 8bit 量化的 GEMM 计算（Marlin 优化版）   

    ## 关键前置条件
    必须配合对应的权重 Shuffle 函数使用，否则会导致计算结果完全错误：
    - GEMM1 场景：使用 ops.marlin_weights 处理权重
    - GEMM2 场景：使用 ops.marlin_weights_ours 处理权重




    
    ---------------------------------------------------------------
    """
    
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_gemm_marlin_w4a16( 
    input: torch.Tensor,
    b_qweight : torch.Tensor,
    output : torch.Tensor,
    b_scale: torch.Tensor,
    b_zeros : torch.Tensor,
    topk_weights : Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids : torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    top_k : int,
    mode :int,
    delta: int)-> torch.Tensor: 
    """
    ---------------------------------------------------------------
    # MoE 场景下 4bit 量化的 GEMM 计算（Marlin 优化版）   

    ## 关键前置条件
    必须配合对应的权重 Shuffle 函数使用，否则会导致计算结果完全错误：

    
    ---------------------------------------------------------------
    """
    
    pass


@compile_ops("module_moe_c_kernel")
def moe_c_moe_w8a8_gemm_block_wise(
    input: torch.Tensor,
    a_scales: torch.Tensor,
    output: torch.Tensor,
    b_qweight: torch.Tensor,
    b_scales: torch.Tensor,
    b_qzeros: Optional[torch.Tensor],   
    topk_weights: Optional[torch.Tensor],   
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    group_size_n: int,  
    group_size_k: int,  
    top_k: int,   
    BLOCK_SIZE_m: int,  
    BLOCK_SIZE_n: int,  
    BLOCK_SIZE_k: int,  
    kloops: int,   
    nloops: int,   
    bit: int   
) -> torch.Tensor:  
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_w8a8_gemm_block_wise_kernel2(
    input: torch.Tensor,
    a_scales: torch.Tensor,
    output: torch.Tensor,
    b_qweight: torch.Tensor,
    b_scales: torch.Tensor,
    b_qzeros: Optional[torch.Tensor],
    topk_weights: Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    group_size_n: int,
    group_size_k: int,
    top_k: int,
    BLOCK_SIZE_m: int,
    BLOCK_SIZE_n: int,
    BLOCK_SIZE_k: int,
    kloops: int,
    nloops: int,
    bit: int
) -> torch.Tensor:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_w8a8_gemm_block_wise_fp8(
    input: torch.Tensor,
    a_scales: torch.Tensor,
    output: torch.Tensor,
    b_qweight: torch.Tensor,
    b_scales: torch.Tensor,
    b_qzeros: Optional[torch.Tensor],
    topk_weights: Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    group_size_n: int,
    group_size_k: int,
    top_k: int,
    BLOCK_SIZE_m: int,
    BLOCK_SIZE_n: int,
    BLOCK_SIZE_k: int,
    kloops: int,
    nloops: int,
    bit: int
) -> torch.Tensor:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_w8a8_gemm_block_wise_kernel2_fp8(
    input: torch.Tensor,
    a_scales: torch.Tensor,
    output: torch.Tensor,
    b_qweight: torch.Tensor,
    b_scales: torch.Tensor,
    b_qzeros: Optional[torch.Tensor],
    topk_weights: Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    group_size_n: int,
    group_size_k: int,
    top_k: int,
    BLOCK_SIZE_m: int,
    BLOCK_SIZE_n: int,
    BLOCK_SIZE_k: int,
    kloops: int,
    nloops: int,
    bit: int
) -> torch.Tensor:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_w8a16_gemm_awq(
    input: torch.Tensor,
    output: torch.Tensor,
    b_qweight: torch.Tensor,
    b_scales: torch.Tensor,
    b_qzeros: Optional[torch.Tensor],
    topk_weights: Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    top_k: int,
    BLOCK_SIZE_m: int,
    BLOCK_SIZE_n: int,
    BLOCK_SIZE_k: int,
    bit: int
) -> torch.Tensor:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_w8a16_gemm_block_wise(
    input: torch.Tensor,
    output: torch.Tensor,
    b_qweight: torch.Tensor,
    b_scales: torch.Tensor,
    b_qzeros: Optional[torch.Tensor],
    topk_weights: Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    group_size_n: int,
    group_size_k: int,
    top_k: int,
    BLOCK_SIZE_m: int,
    BLOCK_SIZE_n: int,
    BLOCK_SIZE_k: int,
    bit: int
) -> torch.Tensor:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_wna16_gemm_base(
    input: torch.Tensor,
    output: torch.Tensor,
    b_qweight: torch.Tensor,
    b_scales: torch.Tensor,
    b_qzeros: Optional[torch.Tensor],
    topk_weights: Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    top_k: int,
    BLOCK_SIZE_M: int,
    BLOCK_SIZE_N: int,
    BLOCK_SIZE_K: int,
    bit: int
) -> torch.Tensor:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_wna16_gemm(
    input: torch.Tensor,
    output: torch.Tensor,
    b_qweight: torch.Tensor,
    b_scales: torch.Tensor,
    b_qzeros: Optional[torch.Tensor],
    topk_weights: Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    top_k: int,
    BLOCK_SIZE_m: int,
    BLOCK_SIZE_n: int,
    BLOCK_SIZE_k: int,
    kloops: int,
    nloops: int,
    bit: int
) -> torch.Tensor:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_wna16_gemm_2(
    input: torch.Tensor,
    output: torch.Tensor,
    b_qweight: torch.Tensor,
    b_scales: torch.Tensor,
    b_qzeros: Optional[torch.Tensor],
    topk_weights: Optional[torch.Tensor],
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
    top_k: int,
    BLOCK_SIZE_m: int,
    BLOCK_SIZE_n: int,
    BLOCK_SIZE_k: int,
    kloops: int,
    nloops: int,
    bit: int
) -> torch.Tensor:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_topk_softmax(
    topk_weights: torch.Tensor,  # 移除 C++ 引用 &
    topk_indices: torch.Tensor,  # 移除 C++ 引用 &
    token_expert_indices: torch.Tensor,  # 移除 C++ 引用 &
    gating_output: torch.Tensor  # 移除 C++ 引用 &
) -> None:  # 替代 -> None (C++ 中的 void)
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_silu_and_mul( out : torch.Tensor,   
                  input : torch.Tensor) -> None:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_sum(
    input: torch.Tensor,  # 移除 C++ 引用 &
    output: torch.Tensor,  # 移除 C++ 引用 &
    topk_ids: torch.Tensor
) -> None:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_moe_align_block_size(
    topk_ids: torch.Tensor,
    num_experts: int,
    block_size: int,
    sorted_token_ids: torch.Tensor,
    experts_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor
) -> None:
    pass

@compile_ops("module_moe_c_kernel")
def moe_c_sgl_moe_align_block_size(
    topk_ids: torch.Tensor,
    num_experts: int,
    block_size: int,
    sorted_token_ids: torch.Tensor,
    experts_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor
) -> None:
    pass



