#!POPCORN leaderboard amd-mixture-of-experts
from task import input_t, output_t
import torch
from torch.utils.cpp_extension import load_inline
import time
import os
import sys

if "PYTORCH_ROCM_ARCH" not in os.environ:
    os.environ["PYTORCH_ROCM_ARCH"] = "gfx942:xnack-"

kernel_cpp = r"""
#undef __HIP_NO_HALF_OPERATORS__
#undef __HIP_NO_HALF_CONVERSIONS__
#include <hip/hip_runtime.h>
#include <hip/hip_fp16.h>
#include <hip/hip_fp8.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <vector>
#include <string>
#include <iostream>
#include <iomanip>
#include <random>
#include <sstream>
#include <fstream>
#include <chrono>
#include <unistd.h>  

#include <hip/amd_detail/amd_hip_bf16.h>
#include <hip/amd_detail/amd_hip_fp8.h>
using fp16 = __half;

#define WARP_SIZE 64
#define N_WARP_Y_DEF 4
#define N_WARP_X_DEF 4
// Output Tile config
#define TILE_SIZE_Y_DEF 128
#define TILE_SIZE_X_DEF 128
#define TILE_K_DIM_DEF 64


using vec4f = __attribute__((__vector_size__(4 * sizeof(float)))) float;
using vec4fp16 = __attribute__((__vector_size__(4 * sizeof(__fp16)))) __fp16;
using vec2u = __attribute__((__vector_size__(2 * sizeof(uint32_t)))) uint32_t;

#define HIP_CHECK(call)                                                        @
  {                                                                            @
    hipError_t err = call;                                                     @
    if (err != hipSuccess) {                                                   @
      std::cerr << "HIP error: " << hipGetErrorString(err) << " at line "      @
                << __LINE__ << std::endl;                                      @
      exit(err);                                                               @
    }                                                                          @
  }

__device__ void matmul16x16_f16(const fp16 *a, const fp16 *b, vec4f &c, int a_row_offset, int b_col_offset, int k_offset,  int K, int N)
{
int tid = threadIdx.x % WARP_SIZE;
int row_a = tid % 16;
int col_a = (tid / 16);

int col_b = (tid%16)/4;
int row_b = (tid%4) +(tid/16)*4;

//global coordinates
int row_a_g = a_row_offset+row_a;
int col_a_g = k_offset/4+col_a; // convert col_offset to vec4fp16

int row_b_g = k_offset + row_b;
int col_b_g = b_col_offset/4 + col_b;

int col_a_g_swizzled = (col_a_g ^ (row_a_g % 16));
int col_b_g_swizzled = ((col_b_g / 4) ^ (row_b_g % 8)) * 4 + col_b_g % 4;

auto packedA = reinterpret_cast<const vec4fp16 *>(a + row_a_g * K)[col_a_g_swizzled];
auto packedB = reinterpret_cast<const vec4fp16 *>(b + row_b_g * N)[col_b_g_swizzled];

c = __builtin_amdgcn_mfma_f32_16x16x16f16(packedA, packedB, c, 0, 0, 0);
}

__device__ void loadLDS(int warpId, int laneId, int kId, fp16 *s_a, fp16 *s_b,
const fp16 *a, const fp16 *b, int K, int N) {

constexpr int nbReadPerThread = 4;
constexpr int nbWarpPerMatrix = 8; // Todo , compute and add error checking
const fp16 *ptr = a;
fp16 *dst = s_a;
int warpIdLocal = warpId%nbWarpPerMatrix;
int stride = K;
bool transpose = false;
int DIM = TILE_K_DIM_DEF/nbReadPerThread;
uint32_t row_offset = TILE_SIZE_X_DEF * blockIdx.x; 
uint32_t col_offset = kId;

if (warpId>=nbWarpPerMatrix){
ptr=b;
dst=s_b;
stride= N;
transpose= true;
row_offset = kId;
col_offset = TILE_SIZE_Y_DEF * blockIdx.y;
DIM = TILE_SIZE_Y_DEF/nbReadPerThread; //32
}

int nbRowsPerWarp = WARP_SIZE/DIM*nbReadPerThread;//8
vec2u b_value[nbReadPerThread] = {0};

for (int r = 0; r < nbReadPerThread; r++) {
int col = laneId % DIM;
int row = r + nbRowsPerWarp * warpIdLocal + nbReadPerThread * (laneId / DIM);
b_value[r] = (reinterpret_cast<const vec2u *>(ptr + (row_offset + row) * stride + (col_offset + col * nbReadPerThread)))[0];
}

vec2u result_b[nbReadPerThread];
if (transpose){
for (int i = 0; i < nbReadPerThread; ++i)
for (int j = 0; j < nbReadPerThread; ++j) {
reinterpret_cast<uint16_t *>(&result_b[i])[j] =
      reinterpret_cast<uint16_t *>(&b_value[j])[i];
}

for (int r = 0; r < nbReadPerThread; r++) {
  int col = laneId % DIM;
  int row = r + nbRowsPerWarp * warpIdLocal + nbReadPerThread * (laneId / DIM);

int col_swizzled = ((col / 4) ^ (row % 8)) * 4 + col % 4;
  (reinterpret_cast<vec4fp16 *>(dst))[(DIM)*row + col_swizzled] = result_b[r];
}

}else{
for (int i = 0; i < nbReadPerThread; ++i){
  result_b[i]=b_value[i];
}
for (int r = 0; r < nbReadPerThread; r++) {
  int col = laneId % DIM;
  int row = r + nbRowsPerWarp * warpIdLocal + nbReadPerThread * (laneId / DIM);
int col_swizzled = (col ^ (row % 16));
  (reinterpret_cast<vec4fp16 *>(dst))[(DIM)*row + col_swizzled] = result_b[r];
}  
}
}
template<bool applySigmoid, bool useScaling>
__global__ void __launch_bounds__(N_WARP_X_DEF *N_WARP_Y_DEF *WARP_SIZE) 
      custom_kernel(const fp16 *a, 
                    const fp16 **b_array,
                    fp16 *c, int M, int N, int K,fp16* d_scaling, int32_t *d_block_mapping) {

int tid = threadIdx.x;
int laneId = tid % WARP_SIZE;
float result = 0;
const fp16* b = b_array[d_block_mapping[blockIdx.x]];

constexpr int TILE_SIZE_Y = TILE_SIZE_Y_DEF;
constexpr int TILE_SIZE_X = TILE_SIZE_X_DEF;
constexpr int K_DIM = TILE_K_DIM_DEF;
constexpr int NBWARPS_Y = N_WARP_Y_DEF; // 4
constexpr int NBWARPS_X = N_WARP_X_DEF; // 2
// intrinsic config
constexpr int MATMUL_DIM_K = 16;
constexpr int MATMUL_DIM_NM = 16;
constexpr int BLOCK_DIM = NBWARPS_X * NBWARPS_Y * WARP_SIZE; // 512

int warpId = threadIdx.x / WARP_SIZE;
int warpIdx = warpId / NBWARPS_Y; // 0 - > 3
int warpIdy = warpId % NBWARPS_Y;

constexpr int K_DIM_PAD = K_DIM; 

__shared__ fp16 s_a[TILE_SIZE_X][K_DIM_PAD];
__shared__ fp16 s_b[K_DIM_PAD][TILE_SIZE_Y];

constexpr int WARP_TILE_SIZE_Y = TILE_SIZE_Y / NBWARPS_Y / MATMUL_DIM_NM; // 2
constexpr int WARP_TILE_SIZE_X = TILE_SIZE_X / NBWARPS_X / MATMUL_DIM_NM; // 4
vec4f c_s[WARP_TILE_SIZE_X][WARP_TILE_SIZE_Y] = {0};

for (int kId = 0; kId < K; kId += K_DIM) {
loadLDS(warpId, laneId, kId, &s_a[0][0], &s_b[0][0], a, b, K, N); 
__syncthreads();
#pragma unroll
for (int t_wIdx = 0; t_wIdx < WARP_TILE_SIZE_X; t_wIdx++) {
for (int t_wIdy = 0; t_wIdy < WARP_TILE_SIZE_Y; t_wIdy++) {
  constexpr int stepX = TILE_SIZE_X / WARP_TILE_SIZE_X;
  constexpr int stepY = TILE_SIZE_Y / WARP_TILE_SIZE_Y; 
  for (int k = 0; k < K_DIM;k += MATMUL_DIM_K) 
  {
    int a_row = (t_wIdx * stepX) + warpIdx * MATMUL_DIM_NM;
    int b_col = (t_wIdy * stepY) + warpIdy * MATMUL_DIM_NM;
    matmul16x16_f16(&s_a[0][0], &s_b[0][0],c_s[t_wIdx][t_wIdy], a_row, b_col,k,K_DIM_PAD, TILE_SIZE_Y);
  }

}
}

__syncthreads();
}
for (int t_wIdx = 0; t_wIdx < WARP_TILE_SIZE_X; t_wIdx++) {
for (int t_wIdy = 0; t_wIdy < WARP_TILE_SIZE_Y; t_wIdy++) {
// offset pointing to the accumulated 16x16 block
int col_offset =
    blockIdx.y * TILE_SIZE_Y +                 // TILE offset
    (warpIdy * MATMUL_DIM_NM) +                // Warp offset
    t_wIdy * (TILE_SIZE_Y / WARP_TILE_SIZE_Y); // Warp tile offset

int row_offset =
    blockIdx.x * TILE_SIZE_X +                 // TILE offset
    (warpIdx * MATMUL_DIM_NM) +                // Warp offset
    t_wIdx * (TILE_SIZE_X / WARP_TILE_SIZE_X); // Warp tile offset

auto res = c_s[t_wIdx][t_wIdy];
for (int i = 0; i < 4; i++) {
  int col = laneId % 16;
  int row = i + (laneId / 16) * 4;
    auto tmp = (fp16)(res[i]);

    if constexpr(useScaling){
      tmp*=d_scaling[(row_offset + row) * N + (col_offset + col)];
    }
    
    if constexpr(applySigmoid){
      fp16 ratio = fp16(1.0f/(1.0+exp(-res[i])));
      tmp = (fp16)(res[i])*ratio;
    }
    c[(row_offset + row) * N + (col_offset + col)] = tmp;  
}
}
}
}


template<bool applySigmoid, bool useScaling>
__global__ void __launch_bounds__(N_WARP_X_DEF *N_WARP_Y_DEF *WARP_SIZE) 
      custom_kernel_single(const fp16 *a, const fp16 *b, fp16 *c, int M, int N, int K,fp16* d_scaling) {

int tid = threadIdx.x;
int laneId = tid % WARP_SIZE;
float result = 0;

constexpr int TILE_SIZE_Y = TILE_SIZE_Y_DEF;
constexpr int TILE_SIZE_X = TILE_SIZE_X_DEF;
constexpr int K_DIM = TILE_K_DIM_DEF;
constexpr int NBWARPS_Y = N_WARP_Y_DEF; // 4
constexpr int NBWARPS_X = N_WARP_X_DEF; // 2
// intrinsic config
constexpr int MATMUL_DIM_K = 16;
constexpr int MATMUL_DIM_NM = 16;
constexpr int BLOCK_DIM = NBWARPS_X * NBWARPS_Y * WARP_SIZE; // 512

int warpId = threadIdx.x / WARP_SIZE;
int warpIdx = warpId / NBWARPS_Y; // 0 - > 3
int warpIdy = warpId % NBWARPS_Y;

constexpr int K_DIM_PAD = K_DIM;

__shared__ fp16 s_a[TILE_SIZE_X][K_DIM_PAD];
__shared__ fp16 s_b[K_DIM_PAD][TILE_SIZE_Y];

constexpr int WARP_TILE_SIZE_Y = TILE_SIZE_Y / NBWARPS_Y / MATMUL_DIM_NM; // 2
constexpr int WARP_TILE_SIZE_X = TILE_SIZE_X / NBWARPS_X / MATMUL_DIM_NM; // 4
vec4f c_s[WARP_TILE_SIZE_X][WARP_TILE_SIZE_Y] = {0};

for (int kId = 0; kId < K; kId += K_DIM) {
int nbRowsX = BLOCK_DIM / K_DIM; // 4

loadLDS(warpId, laneId, kId, &s_a[0][0], &s_b[0][0], a, b, K, N);

__syncthreads();

#pragma unroll
for (int t_wIdx = 0; t_wIdx < WARP_TILE_SIZE_X; t_wIdx++) {
for (int t_wIdy = 0; t_wIdy < WARP_TILE_SIZE_Y; t_wIdy++) {
  constexpr int stepX = TILE_SIZE_X / WARP_TILE_SIZE_X;
  constexpr int stepY = TILE_SIZE_Y / WARP_TILE_SIZE_Y; // 128/2=64

  for (int k = 0; k < K_DIM;k += MATMUL_DIM_K) 
  {
  
    int a_row = (t_wIdx * stepX) + warpIdx * MATMUL_DIM_NM;
    int b_col = (t_wIdy * stepY) + warpIdy * MATMUL_DIM_NM;

    matmul16x16_f16(&s_a[0][0], &s_b[0][0],c_s[t_wIdx][t_wIdy], a_row, b_col,k,K_DIM_PAD, TILE_SIZE_Y);
  }
  
}
}

__syncthreads();
}
for (int t_wIdx = 0; t_wIdx < WARP_TILE_SIZE_X; t_wIdx++) {
for (int t_wIdy = 0; t_wIdy < WARP_TILE_SIZE_Y; t_wIdy++) {
// offset pointing to the accumulated 16x16 block
int col_offset =
    blockIdx.y * TILE_SIZE_Y +                 // TILE offset
    (warpIdy * MATMUL_DIM_NM) +                // Warp offset
    t_wIdy * (TILE_SIZE_Y / WARP_TILE_SIZE_Y); // Warp tile offset

int row_offset =
    blockIdx.x * TILE_SIZE_X +                 // TILE offset
    (warpIdx * MATMUL_DIM_NM) +                // Warp offset
    t_wIdx * (TILE_SIZE_X / WARP_TILE_SIZE_X); // Warp tile offset

auto res = c_s[t_wIdx][t_wIdy];
// C is stored on 4 VGPRs
for (int i = 0; i < 4; i++) {
  int col = laneId % 16;
  int row = i + (laneId / 16) * 4;

  
    auto tmp = (fp16)(res[i]);

    if constexpr(useScaling){
      tmp*=d_scaling[(row_offset + row) * N + (col_offset + col)];
    }
    
    if constexpr(applySigmoid){
      fp16 ratio = fp16(1.0f/(1.0+exp(-res[i])));
      tmp = (fp16)(res[i])*ratio;
    }
      
    c[(row_offset + row) * N + (col_offset + col)] = tmp;  

  

}
}
}
}

template<bool useSigmoid,bool useScaling>
void matmul_new(const fp16 *d_A, const fp16 **d_B_array, fp16 *d_C, int M, int N, int K,int32_t *d_block_mapping, fp16* d_scaling = nullptr) {

dim3 threadsPerBlock(N_WARP_X_DEF * N_WARP_Y_DEF * WARP_SIZE);
dim3 blocksPerGrid((M + TILE_SIZE_X_DEF - 1) / TILE_SIZE_X_DEF,
                (N + TILE_SIZE_Y_DEF - 1) / TILE_SIZE_Y_DEF);

hipLaunchKernelGGL((custom_kernel<useSigmoid,useScaling>), blocksPerGrid, threadsPerBlock, 0, 0, d_A,d_B_array, d_C, M, N, K,d_scaling,d_block_mapping);
}
template<bool useSigmoid,bool useScaling>
void matmul_new_single(const fp16 *d_A, const fp16 *d_B, fp16 *d_C, int M, int N, int K, fp16* d_scaling = nullptr) {

dim3 threadsPerBlock(N_WARP_X_DEF * N_WARP_Y_DEF * WARP_SIZE);
dim3 blocksPerGrid((M + TILE_SIZE_X_DEF - 1) / TILE_SIZE_X_DEF,
                (N + TILE_SIZE_Y_DEF - 1) / TILE_SIZE_Y_DEF);

hipLaunchKernelGGL((custom_kernel_single<useSigmoid,useScaling>), blocksPerGrid, threadsPerBlock, 0, 0, d_A,d_B, d_C, M, N, K,d_scaling);
}


template <typename T> T *toGPU(const std::vector<T> &h_p) {
auto size = h_p.size() * sizeof(T);
T *d_p;
HIP_CHECK(hipMalloc(&d_p, size));
HIP_CHECK(hipMemcpy(d_p, h_p.data(), size, hipMemcpyHostToDevice));
return d_p;
}

template <typename T> std::vector<T> toCPU(const T *d_p, size_t count) {
std::vector<T> h_p(count);
HIP_CHECK(hipMemcpy(h_p.data(), d_p, count * sizeof(T), hipMemcpyDeviceToHost));
return h_p;
}

__device__ float warp_reduce_sum(float val, int reduction_size = WARP_SIZE) {
#pragma unroll
for (int offset = reduction_size / 2; offset > 0; offset /= 2) {
val += __shfl_down(val, offset,reduction_size);
}
return val;
}

__device__ __forceinline__ float warp_reduce_max(float val,int reduction_size = WARP_SIZE)
{
#pragma unroll
for (int offset = reduction_size / 2; offset > 0; offset /= 2)
{
  val = fmaxf(val, __shfl_down(val, offset, reduction_size));
}
return val;
}
template<int topK, int N, typename T>
__device__ void warpTopK(T val, int idx, T* top_vals, int* top_idxs,int token_id) {
#pragma unroll
for (int i = 0; i < topK; ++i) {
top_vals[i] = -INFINITY;
top_idxs[i] = -1;
}

// Do a local insertion sort within the warp
#pragma unroll
for (int offset = 0; offset < N; ++offset) {
T other_val = __shfl(val, offset);
int other_idx = __shfl(idx, offset);

// Insert in top-K if high enough (break ties with smaller index)
#pragma unroll
for (int i = 0; i < topK; ++i) {
    if (other_val > top_vals[i] || 
        (other_val == top_vals[i] && other_idx < top_idxs[i])) {

        // Shift down
        for (int j = topK - 1; j > i; --j) {
            top_vals[j] = top_vals[j - 1];
            top_idxs[j] = top_idxs[j - 1];
        }

        top_vals[i] = other_val;
        top_idxs[i] = other_idx;
        break;
    }
}
}
}

template<int n_experts,int NB_WARPS>
__global__ void gate_kernel(fp16 *d_scores, int32_t *d_indices, const fp16* W_g, const fp16* x, int d_hidden, int seq_len) {

int token_id = blockIdx.x;
int laneId = threadIdx.x % WARP_SIZE;
int warpId = threadIdx.x / WARP_SIZE;

// One block, one embedding
constexpr int n_iter_warp = n_experts/NB_WARPS; // 2
constexpr int top_k = 4;

int nb_iter = d_hidden/WARP_SIZE;
__shared__ float scores[n_experts]; 
__shared__ uint32_t indices[n_experts];

float sum[n_iter_warp] = {0.0f};
for (int i=0;i<d_hidden;i+=WARP_SIZE){
int col = laneId + i;
int x_row = token_id;  
fp16 x_v = x[x_row*d_hidden+col];

#pragma unroll
for (int j=0;j<n_iter_warp;j++){
int w_row = NB_WARPS*j + warpId;
fp16 x_w = W_g[w_row*d_hidden+col];
sum[j]+=float(x_w)*float(x_v);
}
}

#pragma unroll
for (int j=0;j<n_iter_warp;j++){
sum[j]=warp_reduce_sum(sum[j]);
if (laneId == 0){
scores[NB_WARPS*j + warpId] = fp16(sum[j]);
}
}
__syncthreads();
// Use single warpId to compute softmax
if (warpId == 0 && laneId<n_experts){
auto val = scores[laneId];
// Compute max across warp
auto max_val = warp_reduce_max(val,n_experts);
max_val = __shfl(max_val, 0);  // Broadcast to all threads
// Subtract max and exponentiate
auto exp_val = expf(val - max_val);
// Sum of exponentials
auto sum_val = warp_reduce_sum(exp_val,n_experts);
sum_val = __shfl(sum_val, 0);  // Broadcast
// Normalize
auto softmax_val = exp_val / sum_val;

//check results in fp16 to match pytorch
fp16 top_vals[top_k];
int top_idxs[top_k];

warpTopK<top_k,n_experts,fp16>(fp16(softmax_val), laneId, top_vals, top_idxs,token_id);

// Only lane 0 writes result
if (laneId == 0) {

  for (int i = 0; i < 4; ++i) {
    d_scores[token_id*top_k+i]=fp16(top_vals[i]);
    d_indices[token_id*top_k+i]=top_idxs[i];
  }
}

}

}

template<int n_experts>
__global__ void histogram(const int32_t* d_indices, uint32_t* d_hist, int seq_len) {

constexpr int n_active_experts = 4;
__shared__ uint32_t s_hist[n_experts];

int tid = threadIdx.x;
if (tid < n_experts) {
s_hist[tid] = 0;
}
__syncthreads();
int idx = blockIdx.x * blockDim.x + tid;
if (idx < seq_len*n_active_experts) {
int32_t val = d_indices[idx];
if (val < n_experts) {
    atomicAdd(&s_hist[val], 1);
}
} 

__syncthreads();

if (tid < n_experts) {
atomicAdd(&d_hist[tid], s_hist[tid]);
}
}

int divUp(int x,int n){
return (x + n-1)/n;
}

__global__ void re_order_indices_kernel(fp16* d_output, const fp16 *d_input, const int32_t* d_mapping, int d_hidden,int seq_len,int seq_len_re_order) {

int tid = blockDim.x*blockIdx.x+threadIdx.x;
int dst_row = tid/d_hidden;
int col = tid%d_hidden;

if (dst_row<seq_len_re_order){
int src_row = d_mapping[dst_row];
if (src_row<seq_len){ //make this better by checking hist
d_output[d_hidden*dst_row+col]=d_input[d_hidden*src_row+col];
}else{
d_output[d_hidden*dst_row+col]=fp16(0.0f);
}
}
}

template<int n_experts>
__global__ void id_mapping(int32_t *mapping,int32_t *mapping_rev, const int32_t* d_indices, const uint32_t* d_hist, int seq_len) {

constexpr int n_active_experts=4;
int tid = blockIdx.x*blockDim.x+threadIdx.x;
int total_threads = gridDim.x * blockDim.x;
// Shared memory to compute expert offsets
__shared__ uint32_t expert_offsets[n_experts];

// Thread 0 builds the offsets from the histogram
if (threadIdx.x == 0) {
  uint32_t offset = 0;
  for (int i = 0; i < n_experts; ++i) {
      expert_offsets[i] = offset;
      offset += ((d_hist[i]+(TILE_SIZE_X_DEF-1))/TILE_SIZE_X_DEF)*TILE_SIZE_X_DEF;
  }
}

__syncthreads();

// One thread per row
int row = blockIdx.x * blockDim.x + threadIdx.x;
if (row >= seq_len) return;

//for now single atomic on shared.
for (int row = tid; row < seq_len; row += total_threads) {

for (int i = 0; i < n_active_experts; ++i) {
  int expert = d_indices[row * n_active_experts + i];

  // Atomically fetch and increment to get a unique position for this expert
  int pos = atomicAdd(&expert_offsets[expert], 1);

  // Write the row index at the correct position in mapping
  mapping[pos] = row;
  mapping_rev[row * n_active_experts + i] = pos;
}
}

}

__global__ void reduceK_kernel(fp16 *d_output, const fp16 *d_input_re_order, const int32_t *d_mapping_rev, const fp16*d_scores,const fp16*d_shared_expert, int seq_len, int d_hidden) {

  int tid = blockDim.x*blockIdx.x+threadIdx.x;
  int dst_row = tid/d_hidden;
  int col = tid%d_hidden;
  constexpr int n_active_experts = 4;

  if (tid<seq_len*d_hidden){
    float sum = 0.0f;
    #pragma unroll
    for (int i=0;i<n_active_experts;i++){
      int r = d_mapping_rev[dst_row*n_active_experts+i];
      auto w = d_scores[dst_row*n_active_experts+i];
      sum+= float(w * d_input_re_order[r*d_hidden+col]);
    }
    d_output[dst_row*d_hidden+col] = fp16(sum + float(d_shared_expert[dst_row*d_hidden+col]));
  }
}

void reduceK(fp16 *d_output, const fp16 *d_input_re_order, const int32_t *d_mapping_rev, const fp16*d_scores, const fp16*d_shared_expert, int seq_len, int d_hidden) {
  dim3 threadsPerBlock(WARP_SIZE*4);
  dim3 numBlocks(divUp(seq_len*d_hidden,threadsPerBlock.x));
  hipLaunchKernelGGL(reduceK_kernel, numBlocks, threadsPerBlock, 0, 0, d_output, d_input_re_order, d_mapping_rev,  d_scores,d_shared_expert,seq_len,d_hidden);
}

void re_order_indices(fp16* d_output, const fp16 *d_input, const int32_t* d_mapping, int d_hidden,int seq_len,int seq_len_re_order) {
  dim3 threadsPerBlock(WARP_SIZE*4);
  dim3 numBlocks(divUp(seq_len_re_order*d_hidden,threadsPerBlock.x));
  hipLaunchKernelGGL(re_order_indices_kernel, numBlocks, threadsPerBlock, 0, 0, d_output, d_input, d_mapping,  d_hidden,seq_len,seq_len_re_order);
}



void id_mapping(int32_t *d_mapping, int32_t *d_mapping_rev ,const int32_t* d_indices, const uint32_t* d_hist, int seq_len, int n_experts) {
  dim3 threadsPerBlock(WARP_SIZE*16);
  dim3 numBlocks(1);
  if (n_experts == 32)
  {
    hipLaunchKernelGGL((id_mapping<32>), numBlocks, threadsPerBlock, 0, 0, d_mapping,d_mapping_rev,d_indices,d_hist, seq_len);
  }else if (n_experts == 8){
    hipLaunchKernelGGL((id_mapping<8>), numBlocks, threadsPerBlock, 0, 0, d_mapping,d_mapping_rev,d_indices,d_hist, seq_len);
  }else if (n_experts == 4){
    hipLaunchKernelGGL((id_mapping<4>), numBlocks, threadsPerBlock, 0, 0, d_mapping,d_mapping_rev,d_indices,d_hist, seq_len);
  }
}

void histogram(const int32_t* d_indices, uint32_t* d_hist, int seq_len, int n_active_experts, int n_experts){
  dim3 threadsPerBlock(WARP_SIZE*4);
  dim3 numBlocks(divUp(seq_len*n_active_experts,WARP_SIZE*4));
  if (n_experts == 32){
    hipLaunchKernelGGL((histogram<32>), numBlocks, threadsPerBlock, 0, 0, d_indices,d_hist, seq_len);
  }
  else if (n_experts == 8)
  {
    hipLaunchKernelGGL((histogram<8>), numBlocks, threadsPerBlock, 0, 0, d_indices,d_hist, seq_len);
  }else if (n_experts == 4){
    hipLaunchKernelGGL((histogram<4>), numBlocks, threadsPerBlock, 0, 0, d_indices,d_hist, seq_len);
  }else{
    throw std::runtime_error("Unsupported number of experts. Only 4, 8 or 32 are supported.");
  }
}

void gate(fp16 *d_scores, int32_t *d_indices, const fp16* W_g, const fp16* x, int d_hidden, int seq_len, int n_experts){
  dim3 numBlocks(seq_len);
  if (n_experts == 32){
    constexpr int nbWarps = 16;
    dim3 threadsPerBlock(WARP_SIZE*nbWarps);
    hipLaunchKernelGGL((gate_kernel<32,nbWarps>), numBlocks, threadsPerBlock, 0, 0, d_scores, d_indices,W_g, x, d_hidden, seq_len);
  }else if (n_experts == 8){
    constexpr int nbWarps = 8;
    dim3 threadsPerBlock(WARP_SIZE*nbWarps);
    hipLaunchKernelGGL((gate_kernel<8,nbWarps>), numBlocks, threadsPerBlock, 0, 0, d_scores, d_indices,W_g, x, d_hidden, seq_len);
  }else if (n_experts == 4){
    constexpr int nbWarps = 4;
    dim3 threadsPerBlock(WARP_SIZE*nbWarps);
    hipLaunchKernelGGL((gate_kernel<4,nbWarps>), numBlocks, threadsPerBlock, 0, 0, d_scores, d_indices,W_g, x, d_hidden, seq_len);
  }
  else{
    throw std::runtime_error("Unsupported number of experts. Only 8 or 32 are supported.");
  }
}

void compute_experts(fp16*d_output,const fp16 *d_input_re_order, 
                     const fp16** d_d_W_gates,
                     const fp16** d_d_W_ups,
                     const fp16** d_d_W_downs,
                     const std::vector<int32_t> prefix_sum,
                     const std::vector<int32_t> seq_lengths,
                     int d_hidden, int seq_len, int seq_len_re_order,int d_expert,int n_experts,
                     fp16* d_gate_out,
                     fp16* d_up_out,
                     int32_t*d_block_mapping)
{  
  matmul_new<true,false>(d_input_re_order,
               d_d_W_gates,
               d_gate_out,
               seq_len_re_order,
               d_expert,d_hidden,
               d_block_mapping);

  matmul_new<false,true>( d_input_re_order,
              d_d_W_ups,
              d_up_out,
              seq_len_re_order,
              d_expert,d_hidden,d_block_mapping,
              d_gate_out);

  matmul_new<false,false>(d_up_out,
              d_d_W_downs,
              d_output,
              seq_len_re_order,
              d_hidden,d_expert,d_block_mapping);  

}

void compute_expert(fp16*d_output,const fp16 *d_input, const fp16*d_W_gate,const fp16*d_W_up, const fp16*d_W_down, 
                  int d_hidden, int seq_len, int d_expert,
                  fp16* d_gate_out,
                  fp16* d_up_out){

  matmul_new_single<true,false>(d_input,d_W_gate,d_gate_out,seq_len,d_expert,d_hidden);
  matmul_new_single<false,true>(d_input,d_W_up,d_up_out,seq_len,d_expert,d_hidden, d_gate_out);
  matmul_new_single<false,false>(d_up_out,d_W_down,d_output,seq_len,d_hidden,d_expert);   
}



static const fp16** d_d_w_gate;
static const fp16** d_d_w_up;
static const fp16** d_d_w_down;

static uint32_t *d_hist;
static int32_t *d_mapping;
static int32_t *d_mapping_rev;
static fp16 *d_input_reorder;
static fp16 *d_output_reorder;
static fp16* d_gate_out;
static fp16* d_up_out;
static fp16* d_shared_expert;
static fp16* d_scores;
static int32_t* d_indices;

static int32_t *d_block_mapping;


void allocate(){
  constexpr int max_n_experts = 32;
  constexpr int max_seq_len = 8192;
  constexpr int max_d_expert = 2048;
  constexpr int max_n_active_experts = 4;
  constexpr int padding = max_n_experts*TILE_SIZE_X_DEF;
  constexpr int max_seq_len_re_order = max_seq_len*max_n_active_experts+padding;
  constexpr int max_d_hidden = 7168;
  int max_nb_blocks = divUp(max_seq_len_re_order, TILE_SIZE_X_DEF);

  static bool hasInit = false;
  if (!hasInit){

    //Array of weights
    HIP_CHECK(hipMalloc(&d_d_w_gate, max_n_experts * sizeof(fp16*)));
    HIP_CHECK(hipMalloc(&d_d_w_up, max_n_experts * sizeof(fp16*)));
    HIP_CHECK(hipMalloc(&d_d_w_down, max_n_experts * sizeof(fp16*)));

    HIP_CHECK(hipMalloc(&d_hist, max_n_experts * sizeof(uint32_t)));
    HIP_CHECK(hipMalloc(&d_mapping, max_seq_len_re_order * sizeof(int32_t)));
    HIP_CHECK(hipMalloc(&d_mapping_rev, max_seq_len*max_n_active_experts * sizeof(int32_t)));
    HIP_CHECK(hipMalloc(&d_input_reorder, max_seq_len_re_order*max_d_hidden  * sizeof(fp16)));
    HIP_CHECK(hipMalloc(&d_output_reorder, max_seq_len_re_order*max_d_hidden * sizeof(fp16)));

    HIP_CHECK(hipMalloc(&d_gate_out, max_seq_len_re_order*max_d_expert  * sizeof(fp16)));
    HIP_CHECK(hipMalloc(&d_up_out, max_seq_len_re_order*max_d_expert * sizeof(fp16)));

    HIP_CHECK(hipMalloc(&d_shared_expert, max_seq_len*max_d_hidden * sizeof(fp16)));
    HIP_CHECK(hipMalloc(&d_scores, max_seq_len*max_n_active_experts * sizeof(fp16)));
    HIP_CHECK(hipMalloc(&d_indices, max_seq_len*max_n_active_experts * sizeof(int32_t)));

    HIP_CHECK(hipMalloc(&d_block_mapping, max_nb_blocks * sizeof(int32_t)));
    
    

    hasInit = true;
  }
}

struct Config{
  int d_hidden;
  int d_expert;
  int batch_size;
  int seq_len;
  int n_active_experts;
  int n_experts;
};
static std::vector<uint32_t> h_hist;

void compute_moe(const Config &conf,
                 const fp16* d_input,
                 fp16 *d_output,

                 const fp16* d_W_g,
                 const fp16* d_shared_w_gate,
                 const fp16* d_shared_w_up,
                 const fp16* d_shared_w_down,

                 const std::vector<const fp16*> &d_w_gate,
                 const std::vector<const fp16*> &d_w_up,
                 const std::vector<const fp16*> &d_w_down
                ){

    
  int n_experts = conf.n_experts;
  int d_hidden = conf.d_hidden;
  int d_expert = conf.d_expert;
  int batch_size = conf.batch_size;
  int seq_len = conf.seq_len;
  int n_active_experts = conf.n_active_experts;

  allocate();
  
  // Copy weight ptr to device
  HIP_CHECK(hipMemcpy(d_d_w_gate, d_w_gate.data(), n_experts * sizeof(fp16*), hipMemcpyHostToDevice));
  HIP_CHECK(hipMemcpy(d_d_w_up, d_w_up.data(), n_experts * sizeof(fp16*), hipMemcpyHostToDevice));
  HIP_CHECK(hipMemcpy(d_d_w_down, d_w_down.data(), n_experts * sizeof(fp16*), hipMemcpyHostToDevice));
  
  int padding = n_experts*TILE_SIZE_X_DEF;
  int seq_len_re_order = seq_len*n_active_experts+padding;

  for (int batch=0;batch<batch_size;batch++){
    const fp16* d_input_batch = d_input + batch*seq_len*d_hidden;
    fp16* d_output_batch = d_output + batch*seq_len*d_hidden;

    //reset histogram
    HIP_CHECK(hipMemsetAsync (d_hist, 0, sizeof(int32_t)*n_experts, 0));
    
    //compute shared expert
    compute_expert(d_shared_expert,
      d_input_batch, 
      d_shared_w_gate,
      d_shared_w_up,
      d_shared_w_down,
      d_hidden, 
      seq_len,
      d_expert,
      d_gate_out,
      d_up_out);
    
    gate(d_scores, d_indices, d_W_g,d_input_batch,  d_hidden, seq_len,n_experts);
    histogram(d_indices, d_hist, seq_len, n_active_experts, n_experts);
    auto h_hist = toCPU(d_hist,n_experts);
    
    auto h_hist_up = std::vector<int32_t>(h_hist.size(),0);
    for (int i =0;i<h_hist_up.size();i++){
      h_hist_up[i] = ((h_hist[i]+(TILE_SIZE_X_DEF-1))/TILE_SIZE_X_DEF)*TILE_SIZE_X_DEF;
    }

    std::vector<int32_t> prefix_sum(n_experts,0);
    prefix_sum[0] = 0;
    for (int i = 1; i < n_experts; ++i) {
        prefix_sum[i] = prefix_sum[i - 1] + h_hist_up[i - 1];
    }

    int nb_blocks = divUp(seq_len_re_order, TILE_SIZE_X_DEF);
    std::vector<int32_t> blockMapping(nb_blocks,0);
    for (int i=0;i<nb_blocks;i++){
      int b = 0;
      for (int j=0;j<n_experts;j++){
        if (prefix_sum[j]<=TILE_SIZE_X_DEF*i){
          b = j;
        }
      }
      blockMapping[i]=b;
    }
    HIP_CHECK(hipMemcpy(d_block_mapping, blockMapping.data(), nb_blocks * sizeof(int32_t), hipMemcpyHostToDevice));
    int max_seq = prefix_sum[prefix_sum.size()-1]+h_hist_up[h_hist_up.size()-1];
    id_mapping(d_mapping, d_mapping_rev, d_indices, d_hist, seq_len, n_experts);
    re_order_indices(d_input_reorder, d_input_batch, d_mapping,  d_hidden, seq_len, seq_len_re_order);

    compute_experts(d_output_reorder,
                    d_input_reorder, 
                    d_d_w_gate,
                    d_d_w_up,
                    d_d_w_down,
                    prefix_sum,
                    h_hist_up,
                    d_hidden, seq_len, 
                    max_seq,
                    d_expert, n_experts,
                    d_gate_out,
                    d_up_out,d_block_mapping);
  
    reduceK(d_output_batch, d_output_reorder, d_mapping_rev, d_scores, d_shared_expert, seq_len, d_hidden);
  }

}

void run(
  uintptr_t input_ptr,
  uintptr_t w_router_ptr,
  const std::vector<uintptr_t>& w_gate,
  const std::vector<uintptr_t>& w_up,
  const std::vector<uintptr_t>& w_down,
  uintptr_t w_shared_gate_ptr,
  uintptr_t w_shared_up_ptr,
  uintptr_t w_shared_down_ptr,
  uintptr_t output_ptr,
  int d_hidden,
  int d_expert,
  int batch_size,
  int seq_len,
  int n_active_experts,
  int n_experts
){
  Config conf;
  conf.d_hidden = d_hidden;
  conf.d_expert = d_expert;
  conf.batch_size = batch_size;
  conf.seq_len = seq_len;
  conf.n_active_experts = n_active_experts;
  conf.n_experts = n_experts;

  const auto* d_input = reinterpret_cast<const fp16*>(input_ptr);
  auto *output = reinterpret_cast<fp16*>(output_ptr);

  const auto* d_W_g = reinterpret_cast<const fp16*>(w_router_ptr);
  const auto d_shared_w_gate = reinterpret_cast<const fp16*>(w_shared_gate_ptr);
  const auto d_shared_w_up = reinterpret_cast<const fp16*>(w_shared_up_ptr);
  const auto d_shared_w_down = reinterpret_cast<const fp16*>(w_shared_down_ptr);

  std::vector<const fp16*> d_w_gate;
  std::vector<const fp16*> d_w_up;
  std::vector<const fp16*> d_w_down;

  for (int i=0;i<n_experts;i++){
    d_w_gate.push_back(reinterpret_cast<const fp16*>(w_gate[i]));
    d_w_up.push_back(reinterpret_cast<const fp16*>(w_up[i]));
    d_w_down.push_back(reinterpret_cast<const fp16*>(w_down[i]));
  }

  compute_moe(conf,
    d_input,
    output,
    d_W_g,
    d_shared_w_gate,
    d_shared_w_up,
    d_shared_w_down,
    d_w_gate,
    d_w_up,
    d_w_down
   );
}

PYBIND11_MODULE(moe, m) {
  m.def("moe", &run, "HIP kernel");
}
"""

        
hip_module = load_inline(
    name="moe",
    cpp_sources="",
    cuda_sources=kernel_cpp.replace('@', chr(92)),
    with_cuda=True,
    verbose=False,
    extra_cuda_cflags=["-std=c++17"],
    no_implicit_headers=True,
)

def custom_kernel(data: input_t) -> output_t:
    """
    Args:
        data: Tuple of (input: torch.Tensor, weights: Dict[str, torch.Tensor], config: Dict)
            - input: Input tensor of shape [batch_size, seq_len, hidden_dim]
            - weights: Dictionary containing model weights
            - config: Dictionary containing model configuration parameters
            
    Returns:
        Tuple containing:
            - output: Processed tensor [batch_size, seq_len, d_model]
            - aux_data: Dictionary with auxiliary data
    """
       
    input_tensor, weights, config = data
    output = torch.zeros_like(input_tensor)  
    num_experts = config["n_routed_experts"]
    w_gate = []
    w_up = []
    w_down = []
    for i in range(num_experts):
        w_gate.append(weights[f'experts.{i}.0.weight'].data_ptr())
        w_up.append(weights[f'experts.{i}.1.weight'].data_ptr())
        w_down.append(weights[f'experts.{i}.2.weight'].data_ptr())

    hip_module.moe(
        input_tensor.data_ptr(),
        weights['router.weight'].data_ptr(),
        w_gate,
        w_up,
        w_down,
        weights['shared_experts.0.weight'].data_ptr(),
        weights['shared_experts.1.weight'].data_ptr(),
        weights['shared_experts.2.weight'].data_ptr(),
        output.data_ptr(),
        #input_tensor.data_ptr(),
        config["d_hidden"],
        config["d_expert"],
        config["batch_size"],
        config["seq_len"],
        config["n_experts_per_token"],
        config["n_routed_experts"]
    )

    return output

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