from utils import make_match_reference
from task import input_t, output_t
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple, List, Optional
import math

import triton
import triton.language as tl
from typing import Any, Callable, Dict, List, Optional, Tuple

import os 
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx942'
from torch.utils.cpp_extension import load_inline


CPP_WRAPPER = """
void hip_moe(torch::Tensor a, torch::Tensor b, torch::Tensor c, torch::Tensor topk_weight,
            torch::Tensor sorted_token, torch::Tensor expert_ids, int num_tokens_post_padded,
            int topk, int num_valid_tokens, bool MUL_ROUTED_WEIGHT);
"""

CUDA_SRC = """
#include <hip/hip_runtime.h>
#include <hip/hip_fp16.h>

#include <hip/hip_fp8.h>

#include <iostream>
#include <iomanip>

using half = __half;
using float16_t = _Float16;

#include <random>
#include <rocwmma/rocwmma.hpp>


#define FRAG_M 16
#define FRAG_N 16
#define FRAG_K 16

#define NUM_REPEAT_MN (4)

#define BLOCK_SIZE_M (FRAG_M * NUM_REPEAT_MN)
#define BLOCK_SIZE_N (FRAG_N * NUM_REPEAT_MN)
#define BLOCK_SIZE_K (FRAG_K * 8)

#include <fstream>
#include <cstring>
#include <initializer_list>
#include <cmath>
#include <vector>


constexpr const int BLOCK = 128;

__device__ __inline__ void __sync_warp(void) { __asm__ volatile("s_barrier" : : : "memory"); }

__global__ void hip_invoke_fused_moe_kernel(float16_t * input,
                                            float16_t ** weight_ptr,
                                            float16_t * output,
                                            float16_t * topk_weight,
                                            int * sorted_token_ids,
                                            int * expert_ids,
                                            int num_tokens_post_padded,
                                            int M,
                                            int N,
                                            int K,
                                            int top_k,
                                            int num_valid_tokens,
                                            bool MUL_ROUTED_WEIGHT) {
    int lane_id = threadIdx.x % warpSize;
    constexpr unsigned int ele_per_thread = FRAG_M * FRAG_N / warpSize;
   
    int c_row_base = blockIdx.x * BLOCK_SIZE_M;
    int c_col_base = blockIdx.y * BLOCK_SIZE_N;

    __shared__ int2 lhs_tile[BLOCK_SIZE_K / 4][BLOCK_SIZE_M];
    __shared__ int2 rhs_tile[BLOCK_SIZE_K / 4][BLOCK_SIZE_N];

    int2* lhs_tile_base_ptr = &(lhs_tile[0][0]);
    int2* rhs_tile_base_ptr = &(rhs_tile[0][0]);

    using float16x4 = __attribute__((__vector_size__(4 * sizeof(float16_t)))) float16_t;
    using floatx4 = __attribute__((__vector_size__(4 * sizeof(float)))) float;

    uint result_frag[2][4] = {0};

    if (c_row_base >= num_tokens_post_padded) {
        return;
    }

    int thread_id_in_a_block = threadIdx.x + threadIdx.y * blockDim.x;
    int row_for_input = thread_id_in_a_block / 32;
    int col_for_input = thread_id_in_a_block % 32;

    float16_t * input_ptrs[8];
    for(int i = 0; i < 8; ++i) {
        int idx = c_row_base + row_for_input * 8 + i;
        int offs_token = sorted_token_ids[idx];
        int input_offs_token = offs_token / top_k;

        if(offs_token < num_valid_tokens) {
            input_ptrs[i] = &input[input_offs_token * K];
        }
        else {
            input_ptrs[i] = nullptr;
        }
    }

    int off_experts = expert_ids[blockIdx.x];
    float16_t * cur_expert_ptr = weight_ptr[off_experts];
    for (int bi = 0; bi < K; bi += BLOCK) {

        if(threadIdx.y < 4) {
            // input : M x K
            constexpr int thr_per_row = 32;
            int s_col = (thread_id_in_a_block % thr_per_row);
            int s_row = (thread_id_in_a_block / thr_per_row);

            int g_col = bi + s_col * 4;
            float16_t tmp[8 * 4] = {0};

            for(int r = 0; r < 8; ++r) {
                if(input_ptrs[r] != nullptr) {
                    int2 * lhs_2i_ptr = (int2 *)(input_ptrs[r] + g_col);
                    auto ptr = tmp + r * 4;
                    int2 * ptr_2i_ptr = (int2 *) ptr;
                    *ptr_2i_ptr = *lhs_2i_ptr;
                }
            }

            for(int r = 0; r < 8; ++r) {
                int idx = s_col * BLOCK_SIZE_M + s_row * 8 + r;
                int2* ptr = (int2 *)(tmp + r * 4);

                int2* lhs_tile_2i_ptr = (int2 *)(lhs_tile_base_ptr + idx);
                * lhs_tile_2i_ptr = * ptr;
            }
        }
        else {
            thread_id_in_a_block = thread_id_in_a_block % 256;
            constexpr int thr_per_row = 8;
            int s_col = (thread_id_in_a_block % thr_per_row);
            int s_row = (thread_id_in_a_block / thr_per_row);

            int g_col = c_col_base + s_col * 8;
            int g_row = bi + s_row * 4;

            float16_t tmp0[4 * 8] = {0};
            float16_t tmp1[8 * 4] = {0};

            for(int kk = 0; kk < 4; ++kk) {
                int4* rhs_4i_ptr = (int4*)(cur_expert_ptr + (g_row + kk) * N + g_col);
                auto ptr = tmp0 + kk * 8;
                int4* ptr_4i_ptr = (int4*)ptr;
                *ptr_4i_ptr = *rhs_4i_ptr;
            }

            for(int tr = 0; tr < 4; ++tr) {
                for(int tc = 0; tc < 8; ++tc) {
                    int idx0 = tr * 8 + tc;
                    int idx1 = tc * 4 + tr;

                    tmp1[idx1] = tmp0[idx0];
                }
            }

            for(int c = 0; c < 8; ++c) {
                int2 * rhs_tile_2i_ptr = rhs_tile_base_ptr + s_row * BLOCK_SIZE_N + s_col * 8 + c;
                *rhs_tile_2i_ptr = *((int2*)(tmp1 + c * 4));
            }
        } // load_tile_b
        __syncthreads();

        for(int ii = 0; ii < BLOCK; ii += 16) {
            int i = bi + ii;

            float16_t v_a_frag[2][4] = {0};
            float16_t b_frag[4] = {0};

            {
                // // (K / 4) * M * 4
                int s_row = (lane_id / 16);
                int s_col = (lane_id % 16);
                int warp_col = threadIdx.y / 4;
                int2* lhs_tile_ptr = lhs_tile_base_ptr + (s_row + ii / 4) * BLOCK_SIZE_M + s_col + warp_col * 32;

                for(int ri = 0; ri < 2; ++ri) {
                    int2 * ptr = (int2*)v_a_frag[ri];
                    * ptr = lhs_tile_ptr[ri * 16];
                }
            } // load_tile_a

            {
                // (K / 4) *  N * 4
                int s_row = (lane_id / 16);
                int s_col = (lane_id % 16);
                int warp_col = threadIdx.y % 4;
                int2* rhs_tile_ptr = rhs_tile_base_ptr + (s_row + ii / 4) * BLOCK_SIZE_N + s_col + warp_col * 16;

                int2* b_ptr = (int2 *)b_frag;
                *b_ptr = * rhs_tile_ptr;
            } // load_tile_b

            for(int ri = 0; ri < 2; ++ri) {
                auto a_frag = v_a_frag[ri];

                float16x4 *a = reinterpret_cast<float16x4 *>(a_frag);
                float16x4 *b = reinterpret_cast<float16x4 *>(b_frag);

                floatx4 *d = reinterpret_cast<floatx4 *>(result_frag[ri]);
                __sync_warp();

                *d = __builtin_amdgcn_mfma_f32_16x16x16f16(
                    *a,
                    *b,
                    *d, 0, 0, 0);

            }
        } // for ii
    }// for bi

    float16_t moe_weight[8] = {0};
    float16_t * output_ptrs[8];
    int row_for_output = lane_id / 16;
    int col_for_output = lane_id % 16;

    int offs[8];
    for(int i = 0; i < 8; ++i) {
        int idx = c_row_base + (threadIdx.y / 4) * 32 + row_for_output * 4 + (i / 4) * 16 + i % 4;
        int offs_token = sorted_token_ids[idx];
        if(offs_token < num_valid_tokens) {
            if(MUL_ROUTED_WEIGHT) {
                moe_weight[i] = topk_weight[offs_token];
            }

            output_ptrs[i] = &output[offs_token * N];

            offs[i] = offs_token;
        }
        else {
            output_ptrs[i] = nullptr;
        }
    }

    if(MUL_ROUTED_WEIGHT) {
        auto result_ptr = (float*)&result_frag[0][0];
        for(int i = 0; i < 8; ++i) {
            result_ptr[i] = result_ptr[i] * moe_weight[i];
        }
    }

    // save the result.
    for(int ri = 0; ri < 2; ++ri) {
        floatx4 *d = reinterpret_cast<floatx4 *>(result_frag[ri]);
        for(int j = 0; j < ele_per_thread; ++j) {
            auto y = lane_id % 16;
            // auto x = lane_id / 16;
            int outIdx = ri * 4 + j;
            if(output_ptrs[outIdx] != nullptr) {
                auto ptr = output_ptrs[outIdx] + (threadIdx.y % 4) * 16 + c_col_base;
                *(ptr + y) = (*d)[j];
                // *(ptr + y) = moe_weight[outIdx];
                // *(ptr + y) = debug_v;
                // *(ptr + y) = *(input + sorted_token_ids[0] * K + thread_id_in_a_block % 64);

                // *(ptr + y) = input_ptrs[outIdx][thread_id_in_a_block];
                // *(ptr + y) = sorted_token_ids[0];
                // *(ptr + y) = num_tokens_post_padded;

                // *(ptr + y) = offs[outIdx];
                // *(ptr + y) = 3;

            }
        }
    }
}

void hip_moe(torch::Tensor a, torch::Tensor weight_vptr, torch::Tensor c, torch::Tensor topk_weight,
            torch::Tensor sorted_token, torch::Tensor expert_ids, int num_tokens_post_padded,
            int topk, int num_valid_tokens, bool MUL_ROUTED_WEIGHT) {
    int M = a.size(0);
    int K = a.size(1);
    int N = c.size(2);
    int sorted_token_num = sorted_token.size(0);
    
    //std::cout << "hip_moe" << std::endl;
    //std::cout << "MNK : " << M << " " << N << " " << K << std::endl;
    //std::cout << "num_tokens_post_padded : " << num_tokens_post_padded << std::endl;
    //std::cout << "topk : " << topk << std::endl << std::endl;
    
    #define CEILDIV(x, y) (((x) + (y)-1) / (y))
    dim3 grid(CEILDIV(sorted_token_num, BLOCK_SIZE_M), CEILDIV(N, BLOCK_SIZE_N));
    dim3 block(64, 8);
    
    float16_t * lhs_d = (float16_t *)a.data_ptr();
    float16_t ** rhs_d = (float16_t **)weight_vptr.data_ptr();
    float16_t * output_d = (float16_t *)c.data_ptr();
    float16_t * topk_weight_d = (float16_t *)topk_weight.data_ptr();
    int * sorted_token_d = (int *)sorted_token.data_ptr();
    int * expert_ids_d = (int *)expert_ids.data_ptr();
    
    hip_invoke_fused_moe_kernel<<<grid, block>>>(lhs_d, rhs_d, output_d, topk_weight_d,
                                            sorted_token_d, expert_ids_d, num_tokens_post_padded,
                                            M, N, K, topk, num_valid_tokens, MUL_ROUTED_WEIGHT);
}
"""

import os
os.environ["CXX"] = "clang++"
module = load_inline(
    name='hip_moe',
    cpp_sources=[CPP_WRAPPER],
    cuda_sources=[CUDA_SRC],
    functions=['hip_moe'],
    verbose=True,
    extra_cuda_cflags=["--offload-arch=gfx942", "-std=c++20"],
)

class Expert(nn.Module):
    def __init__(self, config: Dict, d_expert: Optional[int] = None):
        super().__init__()
        self.config = config
        self.act_fn = nn.SiLU()
        self.d_hidden: int = config["d_hidden"]
        self.d_expert: int = config["d_expert"] if d_expert is None else d_expert

        self.W_gate = nn.Linear(self.d_hidden, self.d_expert, bias=False).cuda()
        self.W_up = nn.Linear(self.d_hidden, self.d_expert, bias=False).cuda()
        self.W_down = nn.Linear(self.d_expert, self.d_hidden, bias=False).cuda()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate = self.act_fn(self.W_gate(x))
        out = self.W_down(gate * self.W_up(x))
        return out

class PengcuoSharedExpert(nn.Module):
    def __init__(self, config: Dict, d_expert: Optional[int] = None):
        super().__init__()
        self.config = config
        self.act_fn = nn.SiLU()
        self.d_hidden: int = config["d_hidden"]
        self.d_expert: int = config["d_expert"] if d_expert is None else d_expert

    def update_w(self, w_gate, w_up, w_down):
        self.w_gate = w_gate
        self.w_up = w_up
        self.w_down = w_down


    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate = self.act_fn(x @ self.w_gate)
        out = (gate * (x @ self.w_up)) @ self.w_down
        return out

class PengcuoMoEGate(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.top_k: int = config["n_experts_per_token"]
        self.num_experts: int = config["n_routed_experts"]
        self.d_hidden: int = config["d_hidden"]

    def update_w(self, w_g):
        self.w_g = w_g

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        logits = x @ self.w_g
        scores = logits.softmax(dim=-1)
        topk_scores, topk_indices = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)

        return topk_indices, topk_scores

@triton.jit
def moe_align_block_size_stage1(
    topk_ids_ptr,
    tokens_cnts_ptr,
    num_experts: tl.constexpr,
    numel: tl.constexpr,
    tokens_per_thread: tl.constexpr,
):
    pid = tl.program_id(0)

    start_idx = pid * tokens_per_thread

    off_c = (pid + 1) * num_experts

    for i in range(tokens_per_thread):
        if start_idx + i < numel:
            idx = tl.load(topk_ids_ptr + start_idx + i)
            token_cnt = tl.load(tokens_cnts_ptr + off_c + idx)
            tl.store(tokens_cnts_ptr + off_c + idx, token_cnt + 1)

@triton.jit
def moe_align_block_size_stage2(
    tokens_cnts_ptr,
    num_experts: tl.constexpr,
):
    pid = tl.program_id(0)

    last_cnt = 0
    for i in range(1, num_experts + 1):
        token_cnt = tl.load(tokens_cnts_ptr + i * num_experts + pid)
        last_cnt = last_cnt + token_cnt
        tl.store(tokens_cnts_ptr + i * num_experts + pid, last_cnt)


@triton.jit
def moe_align_block_size_stage3(
    total_tokens_post_pad_ptr,
    tokens_cnts_ptr,
    cumsum_ptr,
    num_experts: tl.constexpr,
    block_size: tl.constexpr,
):
    last_cumsum = 0
    off_cnt = num_experts * num_experts
    for i in range(1, num_experts + 1):
        token_cnt = tl.load(tokens_cnts_ptr + off_cnt + i - 1)
        last_cumsum = last_cumsum + tl.cdiv(token_cnt, block_size) * block_size
        tl.store(cumsum_ptr + i, last_cumsum)
    tl.store(total_tokens_post_pad_ptr, last_cumsum)


@triton.jit
def moe_align_block_size_stage4(
    topk_ids_ptr,
    sorted_token_ids_ptr,
    expert_ids_ptr,
    tokens_cnts_ptr,
    cumsum_ptr,
    num_experts: tl.constexpr,
    block_size: tl.constexpr,
    numel: tl.constexpr,
    tokens_per_thread: tl.constexpr,
):
    pid = tl.program_id(0)
    start_idx = tl.load(cumsum_ptr + pid)
    end_idx = tl.load(cumsum_ptr + pid + 1)

    for i in range(start_idx, end_idx, block_size):
        tl.store(expert_ids_ptr + i // block_size, pid)

    start_idx = pid * tokens_per_thread
    off_t = pid * num_experts

    for i in range(start_idx, tl.minimum(start_idx + tokens_per_thread, numel)):
        expert_id = tl.load(topk_ids_ptr + i)
        token_cnt = tl.load(tokens_cnts_ptr + off_t + expert_id)
        rank_post_pad = token_cnt + tl.load(cumsum_ptr + expert_id)
        tl.store(sorted_token_ids_ptr + rank_post_pad, i)
        tl.store(tokens_cnts_ptr + off_t + expert_id, token_cnt + 1)

def ceil_div(a, b):
    return (a + b - 1) // b

@triton.jit
def fused_moe_kernel_gate_up(
    a_ptr,
    b_ptr,
    w1_gate_ptr_array,
    w1_up_ptr_array,
    c_ptr,
    a_scale_ptr,
    b_scale_ptr,
    topk_weights_ptr,
    sorted_token_ids_ptr,
    expert_ids_ptr,
    num_tokens_post_padded_ptr,
    N,
    K,
    EM,
    num_valid_tokens,
    stride_am,
    stride_ak,
    stride_be,
    stride_bk,
    stride_bn,
    stride_cm,
    stride_cn,
    stride_asm,
    stride_ask,
    stride_bse,
    stride_bsk,
    stride_bsn,
    group_n: tl.constexpr,
    group_k: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
    MUL_ROUTED_WEIGHT: tl.constexpr,
    top_k: tl.constexpr,
    compute_type: tl.constexpr,
    use_fp8_w8a8: tl.constexpr,
    use_int8_w8a8: tl.constexpr,
    use_int8_w8a16: tl.constexpr,
    per_channel_quant: tl.constexpr,
    even_Ks: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    
    num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    half_num_pid_n = num_pid_n // 2

    pid_m = pid // num_pid_n
    pid_n = pid % num_pid_n
    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
    if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
        return
    offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_token = tl.load(sorted_token_ids_ptr + offs_token_id)
    offs_token = offs_token.to(tl.int64)
    token_mask = offs_token < num_valid_tokens
    half_num_block_n = (N // BLOCK_SIZE_N) // 2

    if pid_n < half_num_block_n:
        tmp_ptr = w1_gate_ptr_array
    else:
        tmp_ptr = w1_up_ptr_array


    HALF_N = N // 2
    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % HALF_N

    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = a_ptr + (
        offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak
    )

    off_experts = tl.load(expert_ids_ptr + pid_m)

    tmp_ptr2 = tl.load(
        tmp_ptr + off_experts,
    )

    cur_expert_ptr = tmp_ptr2.to(tl.pointer_type(tl.float16))

    b_ptrs = (
        cur_expert_ptr
        + (offs_k[:, None] * HALF_N + offs_bn[None, :])
    )

    tl.multiple_of(a_ptrs, [BLOCK_SIZE_M, BLOCK_SIZE_K])
    tl.multiple_of(b_ptrs, [BLOCK_SIZE_K, BLOCK_SIZE_N])
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        if even_Ks:
            a = tl.load(
                a_ptrs,
                mask=token_mask[:, None],
                other=0.0,
            )
            b = tl.load(b_ptrs)
        else:
            a = tl.load(
                a_ptrs,
                mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K),
                other=0.0,
            )
            b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)


        accumulator += tl.dot(a, b)
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk

    if MUL_ROUTED_WEIGHT:
        moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
        accumulator = accumulator * moe_weight[:, None]

    accumulator = accumulator.to(compute_type)
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
    c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
    
    tl.store(c_ptrs, accumulator, mask=c_mask)

def invoke_fused_moe_kernel_gate_up(
    A: torch.Tensor,
    B: torch.Tensor,
    C: torch.Tensor,
    A_scale: Optional[torch.Tensor],
    B_scale: Optional[torch.Tensor],
    B_zp: Optional[torch.Tensor],
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_padded: torch.Tensor,
    mul_routed_weight: bool,
    top_k: int,
    config: Dict[str, Any],
    compute_type: tl.dtype,
    use_fp8_w8a8: bool,
    use_int8_w8a8: bool,
    use_int8_w8a16: bool,
    use_int4_w4a16: bool,
    per_channel_quant: bool,
    w1_gate_ptr_array: torch.Tensor,
    w1_up_ptr_array: torch.Tensor,
    block_shape: Optional[List[int]] = None,
    no_combine: bool = False,
) -> None:
    topk_weights = topk_weights.view(-1, topk_weights.shape[-1])
    topk_weights = topk_weights.contiguous()

    assert topk_weights.stride(1) == 1
    assert sorted_token_ids.stride(0) == 1

    padded_size = 0
    A_sacle = None
    B_scale = None

    grid = lambda META: (
        triton.cdiv(sorted_token_ids.shape[0], META["BLOCK_SIZE_M"])
        * triton.cdiv(B.shape[1], META["BLOCK_SIZE_N"]),
    )

    K = B.shape[2] - padded_size
    if K % config["BLOCK_SIZE_K"] == 0:
        even_Ks = True
    else:
        even_Ks = False
    stride_bk = B.shape[1] // 2
    stride_bn = 1
    fused_moe_kernel_gate_up[grid](
        A,
        B,
        w1_gate_ptr_array,
        w1_up_ptr_array,
        C,
        A_scale,
        B_scale,
        topk_weights,
        sorted_token_ids,
        expert_ids,
        num_tokens_post_padded,
        B.shape[1],
        B.shape[2] - padded_size,
        sorted_token_ids.shape[0],
        topk_ids.numel(),
        A.stride(0),
        A.stride(1),
        B.stride(0),
        stride_bk,
        stride_bn,
        C.stride(1),
        C.stride(2),
        A_scale.stride(0) if A_scale is not None and A_scale.ndim == 2 else 0,
        A_scale.stride(1) if A_scale is not None and A_scale.ndim == 2 else 0,
        B_scale.stride(0) if B_scale is not None and B_scale.ndim >= 2 else 0,
        B_scale.stride(2) if B_scale is not None and B_scale.ndim == 3 else 0,
        B_scale.stride(1) if B_scale is not None and B_scale.ndim >= 2 else 0,
        0 if block_shape is None else block_shape[0],
        0 if block_shape is None else block_shape[1],
        MUL_ROUTED_WEIGHT=mul_routed_weight,
        top_k=top_k,
        compute_type=compute_type,
        use_fp8_w8a8=use_fp8_w8a8,
        use_int8_w8a8=use_int8_w8a8,
        use_int8_w8a16=use_int8_w8a16,
        per_channel_quant=per_channel_quant,
        even_Ks=even_Ks,
        **config,)

def moe_align_block_size_triton(
    topk_ids: torch.Tensor,
    num_experts: int,
    block_size: int,
    sorted_token_ids: torch.Tensor,
    expert_ids: torch.Tensor,
    num_tokens_post_pad: torch.Tensor,
) -> None:
    numel = topk_ids.numel()
    grid = (num_experts,)
    tokens_cnts = torch.zeros(
        (num_experts + 1, num_experts), dtype=torch.int32, device=topk_ids.device
    )
    cumsum = torch.zeros((num_experts + 1,), dtype=torch.int32, device=topk_ids.device)
    tokens_per_thread = ceil_div(numel, num_experts)

    moe_align_block_size_stage1[grid](
        topk_ids,
        tokens_cnts,
        num_experts,
        numel,
        tokens_per_thread,
    )
    moe_align_block_size_stage2[grid](
        tokens_cnts,
        num_experts,
    )
    moe_align_block_size_stage3[(1,)](
        num_tokens_post_pad,
        tokens_cnts,
        cumsum,
        num_experts,
        block_size,
    )
    moe_align_block_size_stage4[grid](
        topk_ids,
        sorted_token_ids,
        expert_ids,
        tokens_cnts,
        cumsum,
        num_experts,
        block_size,
        numel,
        tokens_per_thread,
    )

def moe_align_block_size(
    topk_ids: torch.Tensor, block_size: int, num_experts: int
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
    sorted_ids = torch.empty(
        (max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device
    )
    sorted_ids.fill_(topk_ids.numel())
    max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
    expert_ids = torch.empty(
        (max_num_m_blocks,), dtype=torch.int32, device=topk_ids.device
    )
    num_tokens_post_pad = torch.empty((1), dtype=torch.int32, device=topk_ids.device)


    moe_align_block_size_triton(
        topk_ids,
        num_experts,
        block_size,
        sorted_ids,
        expert_ids,
        num_tokens_post_pad,
    )
 
    sorted_ids_check = torch.empty_like(sorted_ids)
    expert_ids_check = torch.empty_like(expert_ids)
    num_tokens_post_pad_check = torch.empty_like(num_tokens_post_pad)
    
    token_cnts_buffer = torch.empty(
        (num_experts + 1) * num_experts,
        dtype=torch.int32,
        device=topk_ids.device,
    )
    cumsum_buffer = torch.empty(
        num_experts + 1, dtype=torch.int32, device=topk_ids.device
    )
    return sorted_ids, expert_ids, num_tokens_post_pad

def fused_experts_impl(
    hidden_states: torch.Tensor,
    w1: torch.Tensor,
    w2: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    w1_gate_ptr_array: torch.Tensor,
    w1_up_ptr_array: torch.Tensor,
    w2_ptr_array: torch.Tensor,
    inplace: bool = False,
    activation: str = "silu",
    apply_router_weight_on_input: bool = False,
    use_fp8_w8a8: bool = False,
    use_int8_w8a8: bool = False,
    use_int8_w8a16: bool = False,
    use_int4_w4a16: bool = False,
    per_channel_quant: bool = False,
    w1_scale: Optional[torch.Tensor] = None,
    w2_scale: Optional[torch.Tensor] = None,
    w1_zp: Optional[torch.Tensor] = None,
    w2_zp: Optional[torch.Tensor] = None,
    a1_scale: Optional[torch.Tensor] = None,
    a2_scale: Optional[torch.Tensor] = None,
    block_shape: Optional[List[int]] = None,
    no_combine: bool = False,
):
    padded_size = 0
    assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
    assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
    assert w1.is_contiguous(), "Expert weights1 must be contiguous"
    assert w2.is_contiguous(), "Expert weights2 must be contiguous"
    assert hidden_states.dtype in [torch.float32, torch.float16, torch.bfloat16]

    hidden_dim = hidden_states.shape[-1]
    hidden_states = hidden_states.view(-1, hidden_dim)
    
    num_tokens, _ = hidden_states.shape
    E, N, _ = w1.shape
    CHUNK_SIZE = 64 * 1024
    M = min(num_tokens, CHUNK_SIZE)

    config = {
        "BLOCK_SIZE_M": 64,
        "BLOCK_SIZE_N": 64,
        "BLOCK_SIZE_K": 256,
        "GROUP_SIZE_M": 1,
        "num_warps": 4,
        "num_stages": 2,
    }

    cache = torch.empty(
        M * topk_ids.shape[1] * max(N, w2.shape[1]),
        device=hidden_states.device,
        dtype=hidden_states.dtype,
    )
    

    intermediate_cache1 = cache[: M * topk_ids.shape[1] * N].view(
        (M, topk_ids.shape[1], N))

    cache1_num = M * topk_ids.shape[1] * (N // 2)
    cache10 = torch.empty(
        cache1_num,
        device=hidden_states.device,
        dtype=hidden_states.dtype,
    ).view(M, topk_ids.shape[1], N // 2)
    cache11 = torch.empty(
        cache1_num,
        device=hidden_states.device,
        dtype=hidden_states.dtype,
    ).view(M, topk_ids.shape[1], N // 2)
    
    intermediate_cache3 = cache[: M * topk_ids.shape[1] * w2.shape[1]].view(
        (M, topk_ids.shape[1], w2.shape[1]))

    compute_type = tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16
    out_hidden_states = torch.empty_like(hidden_states)

    for chunk in range((num_tokens // CHUNK_SIZE) + 1):
        begin_chunk_idx, end_chunk_idx = (
            chunk * CHUNK_SIZE,
            min((chunk + 1) * CHUNK_SIZE, num_tokens),
        )
        curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
        tokens_in_chunk, _ = curr_hidden_states.shape

        curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
        curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]

        sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
            curr_topk_ids, config["BLOCK_SIZE_M"], E)
        
        invoke_fused_moe_kernel_gate_up(
            curr_hidden_states,
            w1,
            intermediate_cache1,
            a1_scale,
            w1_scale,
            w1_zp,
            curr_topk_weights,
            curr_topk_ids,
            sorted_token_ids,
            expert_ids,
            num_tokens_post_padded,
            apply_router_weight_on_input,
            topk_ids.shape[1],
            config,
            compute_type=compute_type,
            use_fp8_w8a8=use_fp8_w8a8,
            use_int8_w8a8=use_int8_w8a8,
            use_int8_w8a16=use_int8_w8a16,
            use_int4_w4a16=use_int4_w4a16,
            per_channel_quant=per_channel_quant,
            block_shape=block_shape,
            w1_gate_ptr_array=w1_gate_ptr_array,
            w1_up_ptr_array=w1_up_ptr_array)
        
        intermediate_cache2 = torch.nn.functional.silu(intermediate_cache1[..., : N // 2]) * intermediate_cache1[..., N // 2 :]

        intermediate_cache2 = intermediate_cache2.view(-1, N // 2)
        module.hip_moe(intermediate_cache2, w2_ptr_array, intermediate_cache3, curr_topk_weights,
                       sorted_token_ids, expert_ids, num_tokens_post_padded,
                       1, curr_topk_ids.numel(), True)

        out_hidden_states = intermediate_cache3.sum(dim = 1)

    return out_hidden_states

class PengcuoExpert(nn.Module):
    def __init__(self, d_hidden, d_expert, n_routed_experts):
        super().__init__()
        self.act_fn = nn.SiLU()
        self.d_hidden = d_hidden
        self.d_expert = d_expert
        
        E = n_routed_experts
        N = d_expert
        D = d_hidden

        self.w1 = torch.zeros((E, N * 2, D), dtype=torch.float16, device='cuda')
        self.w2 = torch.zeros((E, D, N), dtype=torch.float16, device='cuda')
        
        self.w1_gate_ptr_array = torch.zeros(E, dtype=torch.int64, device="cuda")
        self.w1_up_ptr_array = torch.zeros(E, dtype=torch.int64, device="cuda")
        self.w2_ptr_array = torch.zeros(E, dtype=torch.int64, device="cuda")

    def forward(self, x: torch.Tensor, topk_weights, topk_ids) -> torch.Tensor:
        orig_shape = x.shape
        x = x.view(-1, x.shape[-1])
        topk_weights = topk_weights.view(-1, topk_weights.shape[-1])
        topk_ids = topk_ids.view(-1, topk_ids.shape[-1])

        out = fused_experts_impl(x, self.w1, self.w2, topk_weights, topk_ids, self.w1_gate_ptr_array,
                                 self.w1_up_ptr_array, self.w2_ptr_array).view(*orig_shape)
        return out

class PengcuoMoE(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.config = config

        self.experts = PengcuoExpert(config["d_hidden"], config["d_expert"], config["n_routed_experts"])
        self.gating_network = PengcuoMoEGate(config)

        shared_expert_dim = config["d_expert"] * config["n_shared_experts"]
        self.shared_expert = PengcuoSharedExpert(config=config, d_expert=shared_expert_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        shared_output = self.shared_expert(x)
        expert_indices, expert_scores = self.gating_network(x)
        batch_size, seq_len, hidden_dim = x.shape
        orig_shape = x.shape
        routed_output = self.experts(x, expert_scores, expert_indices)
        output = shared_output + routed_output
        return output

def custom_kernel(data: input_t) -> output_t:
    input_tensor, weights, config = data
    num_experts = config["n_routed_experts"]
    moe = PengcuoMoE(config)
    moe.gating_network.update_w(weights['router.weight'].T)

    # start = torch.cuda.Event(enable_timing=True)
    # end = torch.cuda.Event(enable_timing=True)

    # start.record()
    
    N = config["d_expert"]
    w1 = moe.experts.w1
    w2 = moe.experts.w2
    for i in range(num_experts):
        gate_proj_weight = weights[f'experts.{i}.0.weight']
        up_proj_weight = weights[f'experts.{i}.1.weight']
        down_proj_weight = weights[f'experts.{i}.2.weight']
        
        # E, 2 * N, D
        # w1[i][:N].copy_(gate_proj_weight.T)
        # w1[i][N:].copy_(up_proj_weight.T)
        
        moe.experts.w1_gate_ptr_array[i] = gate_proj_weight.data_ptr()
        moe.experts.w1_up_ptr_array[i] = up_proj_weight.data_ptr()
        moe.experts.w2_ptr_array[i] = down_proj_weight.data_ptr()

    # end.record()
    # torch.cuda.synchronize()
    # print(f"   copy weight duration : {start.elapsed_time(end)} {input_tensor.shape}")

    moe.shared_expert.update_w(weights['shared_experts.0.weight'],
                        weights['shared_experts.1.weight'],
                        weights['shared_experts.2.weight'])
    
    output = moe(input_tensor)

    return output
python submission.py
[92mSuccessfully preprocessed all matching files.[0m

/home/runner/.cache/torch_extensions/py310_cpu/hip_moe/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/hip_moe/hip.hip [ok]
/home/runner/.cache/torch_extensions/py310_cpu/hip_moe/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/hip_moe/main.cpp [skipped, no changes]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 1
[1/3] clang++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=hip_moe -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1016\" -isystem /usr/local/lib/python3.10/dist-packages/torch/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/torch/csrc/api/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/THH -isystem /opt/rocm-6.3.1/include -isystem /usr/include/python3.10 -fPIC -std=c++17 -c /home/runner/.cache/torch_extensions/py310_cpu/hip_moe/main.cpp -o main.o -D__HIP_PLATFORM_AMD__=1 -DUSE_ROCM=1 -DHIPBLAS_V2 -fPIC
[2/3] /opt/rocm-6.3.1/bin/hipcc  -DWITH_HIP -DTORCH_EXTENSION_NAME=hip_moe -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1016\" -isystem /usr/local/lib/python3.10/dist-packages/torch/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/torch/csrc/api/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/THH -isystem /opt/rocm-6.3.1/include -isystem /usr/include/python3.10 -fPIC -std=c++17 -D__HIP_PLATFORM_AMD__=1 -DUSE_ROCM=1 -DHIPBLAS_V2 -fPIC -DCUDA_HAS_FP16=1 -D__HIP_NO_HALF_OPERATORS__=1 -D__HIP_NO_HALF_CONVERSIONS__=1 -DHIP_ENABLE_WARP_SYNC_BUILTINS=1 --offload-arch=gfx942 -fno-gpu-rdc --offload-arch=gfx942 -std=c++20 -c /home/runner/.cache/torch_extensions/py310_cpu/hip_moe/hip.hip -o hip.cuda.o 
[3/3] clang++ main.o hip.cuda.o -shared -L/usr/local/lib/python3.10/dist-packages/torch/lib -lc10 -lc10_hip -ltorch_cpu -ltorch_hip -ltorch -ltorch_python -L/opt/rocm-6.3.1/lib -lamdhip64 -o hip_moe.so

python eval.py benchmark /tmp/tmpf1fna56e
[92mSuccessfully preprocessed all matching files.[0m

/home/runner/.cache/torch_extensions/py310_cpu/hip_moe/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/hip_moe/hip.hip [skipped, already hipified]
/home/runner/.cache/torch_extensions/py310_cpu/hip_moe/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/hip_moe/main.cpp [skipped, no changes]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 1
[1/1] clang++ hip.cuda.o main.o -shared -L/usr/local/lib/python3.10/dist-packages/torch/lib -lc10 -lc10_hip -ltorch_cpu -ltorch_hip -ltorch -ltorch_python -L/opt/rocm-6.3.1/lib -lamdhip64 -o hip_moe.so
