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);
"""

CPP_WRAPPER2 = """
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 num_tokens_post_pad,
    torch::Tensor token_cnts_buffer,
    torch::Tensor cumsum_buffer);
"""

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];
            }
        }
    }
}

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);
    
    #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);
}
"""

CUDA_SRC2 = """
#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 CEILDIV(x, y) (((x) + (y)-1) / (y))

#define WARP_SIZE 64
template <typename scalar_t>
__global__ void count_and_sort_expert_tokens_kernel(
    const scalar_t* __restrict__ topk_ids,
    int32_t* __restrict__ sorted_token_ids,
    int32_t* __restrict__ cumsum_buffer,
    size_t numel) {
  const size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
  const size_t stride = blockDim.x * gridDim.x;

  for (size_t i = tid; i < numel; i += stride) {
    int32_t expert_id = topk_ids[i];
    int32_t rank_post_pad = atomicAdd(&cumsum_buffer[expert_id], 1);
    sorted_token_ids[rank_post_pad] = i;
  }
}

template <typename scalar_t>
__global__ void moe_align_block_size_kernel(
    const scalar_t* __restrict__ topk_ids,
    int32_t* __restrict__ sorted_token_ids,
    int32_t* __restrict__ expert_ids,
    int32_t* __restrict__ total_tokens_post_pad,
    int32_t num_experts,
    int32_t padded_num_experts,
    int32_t experts_per_warp,
    int32_t block_size,
    size_t numel,
    int32_t* __restrict__ cumsum) {
  extern __shared__ int32_t shared_counts[];

  const int warp_id = threadIdx.x / WARP_SIZE;
  const int my_expert_start = warp_id * experts_per_warp;

  for (int i = 0; i < experts_per_warp; ++i) {
    if (my_expert_start + i < padded_num_experts) {
      shared_counts[warp_id * experts_per_warp + i] = 0;
    }
  }

  __syncthreads();

  const size_t tid = threadIdx.x;
  const size_t stride = blockDim.x;

  for (size_t i = tid; i < numel; i += stride) {
    int expert_id = topk_ids[i];
    int warp_idx = expert_id / experts_per_warp;
    int expert_offset = expert_id % experts_per_warp;
    atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset], 1);
  }

  __syncthreads();

  if (threadIdx.x == 0) {
    cumsum[0] = 0;
    for (int i = 1; i <= num_experts; ++i) {
      int expert_count = 0;
      int warp_idx = (i - 1) / experts_per_warp;
      int expert_offset = (i - 1) % experts_per_warp;
      expert_count = shared_counts[warp_idx * experts_per_warp + expert_offset];

      cumsum[i] = cumsum[i - 1] + CEILDIV(expert_count, block_size) * block_size;
    }
    *total_tokens_post_pad = cumsum[num_experts];
  }

  __syncthreads();

  if (threadIdx.x < num_experts) {
    for (int i = cumsum[threadIdx.x]; i < cumsum[threadIdx.x + 1]; i += block_size) {
      expert_ids[i / block_size] = threadIdx.x;
    }
  }
}

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 num_tokens_post_pad,
    torch::Tensor token_cnts_buffer,
    torch::Tensor cumsum_buffer) {
    int64_t padded_num_experts = ((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;

    int experts_per_warp = WARP_SIZE;
    int threads = 1024;

    threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
    auto align_kernel = moe_align_block_size_kernel<int32_t>;

    size_t num_warps = CEILDIV(padded_num_experts, experts_per_warp);
    size_t shared_mem_size = num_warps * experts_per_warp * sizeof(int32_t);

    cumsum_buffer.zero_();

    align_kernel<<<1, threads, shared_mem_size>>>(
        topk_ids.data_ptr<int32_t>(),
        sorted_token_ids.data_ptr<int32_t>(),
        experts_ids.data_ptr<int32_t>(),
        num_tokens_post_pad.data_ptr<int32_t>(),
        num_experts,
        padded_num_experts,
        experts_per_warp,
        block_size,
        topk_ids.numel(),
        cumsum_buffer.data_ptr<int32_t>());

    const int block_threads = std::min(256, (int)threads);
    const int num_blocks = (topk_ids.numel() + block_threads - 1) / block_threads;
    const int max_blocks = 65535;
    const int actual_blocks = std::min(num_blocks, max_blocks);

    auto sort_kernel = count_and_sort_expert_tokens_kernel<int32_t>;
    sort_kernel<<<actual_blocks, block_threads, 0>>>(
        topk_ids.data_ptr<int32_t>(),
        sorted_token_ids.data_ptr<int32_t>(),
        cumsum_buffer.data_ptr<int32_t>(),
        topk_ids.numel());
}
"""


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"],
)
module2 = load_inline(
    name='moe_align_block_size',
    cpp_sources=[CPP_WRAPPER2],
    cuda_sources=[CUDA_SRC2],
    functions=['moe_align_block_size'],
    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

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)

    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
    )
    module2.moe_align_block_size(topk_ids, num_experts, block_size,
                                    sorted_ids, expert_ids, num_tokens_post_pad,
                                    token_cnts_buffer, cumsum_buffer)
 
    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)

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

    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]

        # start.record()

        sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
            curr_topk_ids, config["BLOCK_SIZE_M"], E)

        # end.record()
        # torch.cuda.synchronize()
        # print(f"   moe_align_block_size duration : {start.elapsed_time(end)} {curr_topk_ids.shape}")

        # start.record()
        module.hip_moe(curr_hidden_states, w1_gate_ptr_array, cache10, curr_topk_weights,
                       sorted_token_ids,  expert_ids, num_tokens_post_padded,
                       topk_ids.shape[1], curr_topk_ids.numel(), False)
        module.hip_moe(curr_hidden_states, w1_up_ptr_array, cache11, curr_topk_weights,
                       sorted_token_ids,  expert_ids, num_tokens_post_padded,
                       topk_ids.shape[1], curr_topk_ids.numel(), False)

        intermediate_cache2 = torch.nn.functional.silu(cache10) * cache11
        
        # 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)
        
        # end.record()
        # torch.cuda.synchronize()
        # print(f"   hip_moe duration : {start.elapsed_time(end)} {curr_topk_ids.shape}")

        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='meta')
        self.w2 = torch.zeros((E, D, N), dtype=torch.float16, device='meta')
        
        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.to(torch.int32).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)
    
    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']
        
        moe.experts.w1_gate_ptr_array[i] = weights[f'experts.{i}.0.weight'].data_ptr()
        moe.experts.w1_up_ptr_array[i] = weights[f'experts.{i}.1.weight'].data_ptr()
        moe.experts.w2_ptr_array[i] = weights[f'experts.{i}.2.weight'].data_ptr()

    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 eval.py leaderboard /tmp/tmpoek8tg26
[92mSuccessfully preprocessed all matching files.[0m
[92mSuccessfully preprocessed all matching files.[0m

/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]
/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]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 1
ninja: no work to do.
/home/runner/.cache/torch_extensions/py310_cpu/moe_align_block_size/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/moe_align_block_size/main.cpp [skipped, no changes]
/home/runner/.cache/torch_extensions/py310_cpu/moe_align_block_size/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/moe_align_block_size/hip.hip [skipped, already hipified]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 2
ninja: no work to do.
