from pip._internal import main as pip_main
pip_main(['install', 'tilelang-rocm'])

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.profiler import profile, ProfilerActivity
from typing import Dict, Tuple, Optional
from task import input_t, output_t
import tilelang
import tilelang.language as T
from tilelang.autotuner import *

import time

# tilelang.disable_cache()


@tilelang.jit
def moe_forward_tilelang_shared(d_hidden, d_expert, n_shared_experts, dtype, num_tokens, block_token=128, block_dhidden=128, block_dexpert=128, threads=256, num_stages=1):

    scale = 1.44269504  # log2(e)

    # Parameters
    dhidden = d_hidden
    dexpert = d_expert * n_shared_experts

    # Tensors: Note that input shape is reshape to (num_tokens, dhidden)
    input_shape = (num_tokens, dhidden)
    shared_W_gate_shape = (dexpert, dhidden)
    shared_W_up_shape = (dexpert, dhidden)
    shared_W_down_shape = (dhidden, dexpert)

    accum_type = "float32" 

    @T.prim_func
    def kernel_shared(
        input: T.Tensor(input_shape, dtype), # type: ignore
        shared_W_gate: T.Tensor(shared_W_gate_shape, dtype), # type: ignore
        shared_W_up: T.Tensor(shared_W_up_shape, dtype), # type: ignore
        shared_W_down: T.Tensor(shared_W_down_shape, dtype), # type: ignore
        up_logits: T.Tensor((num_tokens, dexpert), dtype), # type: ignore
        output: T.Tensor(input_shape, dtype), # type: ignore
    ):
        # Step 1: Compute gate and up logits
        with T.Kernel(T.ceildiv(num_tokens, block_token), T.ceildiv(dexpert, block_dexpert), threads=threads) as (bx, by):
            # Split the block to shared experts and routed experts
            input_shared = T.alloc_fragment((block_token, block_dhidden), dtype=dtype)
            W_gate_shared = T.alloc_shared((block_dexpert, block_dhidden), dtype=dtype)
            W_up_shared = T.alloc_shared((block_dexpert, block_dhidden), dtype=dtype)
            # W_down_shared = T.alloc_shared((block_dexpert, block_dim), dtype=dtype)                
            # Shared experts: no need to check expert_indices

            gate_logits_local = T.alloc_fragment((block_token, block_dexpert), dtype=accum_type)
            up_logits_local = T.alloc_fragment((block_token, block_dexpert), dtype=accum_type)
            # down_logits_local = T.alloc_fragment((block_token, block_dim), dtype=dtype)

            T.use_swizzle(10)
            T.clear(gate_logits_local)
            T.clear(up_logits_local)
            # T.clear(down_logits_local)
            
            # Parallel for gate and up matmul
            for k in T.Pipelined(T.ceildiv(dhidden, block_dhidden), num_stages=num_stages):
                T.copy(input[bx * block_token, k * block_dhidden], input_shared)
                # for i, j in T.Parallel(block_token, block_dhidden):
                #     input_shared[i, j] = input[bx * block_token + i, k * block_dhidden + j]
                T.copy(shared_W_gate[by * block_dexpert, k * block_dhidden], W_gate_shared)
                T.copy(shared_W_up[by * block_dexpert, k * block_dhidden], W_up_shared)
                # for i, j in T.Parallel(block_dhidden, block_dexpert):
                #     W_gate_shared[i, j] = shareds_W_gate[bz, k * block_dim + i, by * block_dexpert + j]
                #     W_up_shared[i, j] = shareds_W_up[bz, k * block_dim + i, by * block_dexpert + j]
                T.gemm(input_shared, W_gate_shared, gate_logits_local, transpose_B=True)
                T.gemm(input_shared, W_up_shared, up_logits_local, transpose_B=True)

            # Fuse with SiLU and element-wise product
            for i, j in T.Parallel(block_token, block_dexpert):
                gate_logits_local[i, j] = gate_logits_local[i, j] * (1.0 / (1.0 + T.exp2(-gate_logits_local[i, j] * scale)))
                up_logits_local[i, j] = up_logits_local[i, j] * gate_logits_local[i, j]

            T.copy(up_logits_local, up_logits[bx * block_token, by * block_dexpert])
        
        # Step 2: Compute down logits
        with T.Kernel(T.ceildiv(num_tokens, block_token), T.ceildiv(dhidden, block_dhidden), threads=threads) as (bx, by):
            up_logits_shared = T.alloc_fragment((block_token, block_dexpert), dtype=dtype)
            W_down_shared = T.alloc_shared((block_dhidden, block_dexpert), dtype=dtype)
            output_local = T.alloc_fragment((block_token, block_dhidden), dtype=accum_type)
            
            T.use_swizzle(10)
            T.clear(output_local)

            for k in T.Pipelined(T.ceildiv(dexpert, block_dexpert), num_stages=num_stages):
                T.copy(up_logits[bx * block_token, k * block_dexpert], up_logits_shared)
                T.copy(shared_W_down[by * block_dhidden, k * block_dexpert], W_down_shared)
                T.gemm(up_logits_shared, W_down_shared, output_local, transpose_B=True)

            T.copy(output_local, output[bx * block_token, by * block_dhidden])
    
    return kernel_shared


@tilelang.jit
def moe_forward_tilelang_routed(d_hidden, d_expert, n_routed_experts, dtype, group_sum, group_count, block_token=128, block_dhidden=128, block_dexpert=128, threads=256, num_stages=1, k_pack=1, coalesced_width=None):
    
    scale = 1.44269504  # log2(e)

    # Parameters
    dhidden = d_hidden
    dexpert = d_expert
    n_routed_experts = n_routed_experts

    # Group info
    # group_sum = sum(group_sizes_list)
    # group_count = len(group_sizes_list)
    # M = sum([(group_size + block_token - 1) // block_token for group_size in group_sizes_list])
    M = math.ceil(group_sum / block_token) + group_count
    accum_dtype = "float32"

    # Tensors: Note that input shape is reshape to (bs * seq_len * n_experts_per_token, dhidden) for grouped gemm
    input_shape = (group_sum, dhidden)
    intermediate_shape = (group_sum, dexpert)
    routed_expert_gate_shape = (n_routed_experts, dexpert, dhidden)
    routed_expert_up_shape = (n_routed_experts, dexpert, dhidden)
    routed_expert_down_shape = (n_routed_experts, dhidden, dexpert)
    routed_expert_weights_shape = (group_sum)
    group_sizes_shape = (n_routed_experts)

    @T.prim_func
    def kernel(
        input: T.Tensor(input_shape, dtype), # type: ignore
        routed_expert_gate: T.Tensor(routed_expert_gate_shape, dtype), # type: ignore
        routed_expert_up: T.Tensor(routed_expert_up_shape, dtype), # type: ignore
        routed_expert_down: T.Tensor(routed_expert_down_shape, dtype), # type: ignore
        routed_expert_weights: T.Tensor(routed_expert_weights_shape, dtype), # type: ignore
        group_sizes: T.Tensor(group_sizes_shape, "int32"), # type: ignore
        group_offsets: T.Tensor(group_sizes_shape, "int32"), # type: ignore
        group_padded_offsets: T.Tensor(group_sizes_shape, "int32"), # type: ignore
        group_idx_for_bx: T.Tensor((M,), "int32"), # type: ignore
        up_logits: T.Tensor(intermediate_shape, dtype), # type: ignore
        output: T.Tensor(input_shape, dtype), # type: ignore
    ):
        # Step 1: Compute gate and up logits
        with T.Kernel(M, T.ceildiv(dexpert, block_dexpert), threads=threads) as (bx, by):
            input_shared = T.alloc_fragment((block_token, block_dhidden), dtype=dtype)
            routed_expert_gate_shared = T.alloc_shared((block_dexpert, block_dhidden), dtype=dtype)
            routed_expert_up_shared = T.alloc_shared((block_dexpert, block_dhidden), dtype=dtype)
            
            gate_logits_local = T.alloc_fragment((block_token, block_dexpert), dtype=accum_dtype)
            up_logits_local = T.alloc_fragment((block_token, block_dexpert), dtype=accum_dtype)

            cur_group_idx = T.alloc_local([1], "int32")
            cur_group_size = T.alloc_local([1], "int32")
 
            T.use_swizzle(10, enable=True)

            # T.annotate_layout({input_shared: tilelang.layout.make_swizzled_layout(input_shared)})
            # T.annotate_layout({routed_expert_gate_shared: tilelang.layout.make_swizzled_layout(routed_expert_gate_shared)})
            # T.annotate_layout({routed_expert_up_shared: tilelang.layout.make_swizzled_layout(routed_expert_up_shared)})

            m_start_padded = bx * block_token

            # for i in range(group_count):
            #     in_cur_group_idx = (m_start_padded >= group_padded_offsets[i])
            #     cur_group_idx[0] = T.if_then_else(in_cur_group_idx, i, cur_group_idx[0])
            cur_group_idx[0] = group_idx_for_bx[bx]

            cur_group_size[0] = group_sizes[cur_group_idx[0]]
            m_start = m_start_padded - group_padded_offsets[cur_group_idx[0]] + group_offsets[cur_group_idx[0]]
            actual_rows = T.max(0, T.min(block_token, cur_group_size[0] - (m_start_padded - group_padded_offsets[cur_group_idx[0]])))
            
            T.clear(gate_logits_local)
            T.clear(up_logits_local)

            for k in T.Pipelined(T.ceildiv(dhidden, block_dhidden), num_stages=num_stages):
                T.copy(input[m_start:m_start + block_token, k * block_dhidden:(k + 1) * block_dhidden], input_shared, coalesced_width=coalesced_width)
                T.copy(routed_expert_gate[cur_group_idx[0], by * block_dexpert:(by + 1) * block_dexpert, k * block_dhidden:(k + 1) * block_dhidden], routed_expert_gate_shared, coalesced_width=coalesced_width)
                T.gemm(input_shared, routed_expert_gate_shared, gate_logits_local, k_pack=k_pack, transpose_B=True)
                T.copy(routed_expert_up[cur_group_idx[0], by * block_dexpert:(by + 1) * block_dexpert, k * block_dhidden:(k + 1) * block_dhidden], routed_expert_up_shared, coalesced_width=coalesced_width)
                T.gemm(input_shared, routed_expert_up_shared, up_logits_local, k_pack=k_pack, transpose_B=True)

            for i, j in T.Parallel(block_token, block_dexpert):
                gate_logits_local[i, j] = gate_logits_local[i, j] * (1.0 / (1.0 + T.exp2(-gate_logits_local[i, j] * scale)))
                up_logits_local[i, j] = up_logits_local[i, j] * gate_logits_local[i, j]
                
            for i, j in T.Parallel(block_token, block_dexpert):
                with T.If(i < actual_rows), T.Then():
                    up_logits[m_start + i, by * block_dexpert + j] = up_logits_local[i, j]
        
        # Step 2: Compute down logits
        with T.Kernel(M, T.ceildiv(dhidden, block_dhidden), threads=threads) as (bx, by):
            up_logits_shared = T.alloc_fragment((block_token, block_dexpert), dtype=dtype)
            routed_expert_down_shared = T.alloc_shared((block_dexpert, block_dhidden), dtype=dtype)
            output_local = T.alloc_fragment((block_token, block_dhidden), dtype=accum_dtype)

            cur_group_idx = T.alloc_local([1], "int32")
            cur_group_size = T.alloc_local([1], "int32")

            T.use_swizzle(10, enable=True)

            # T.annotate_layout({up_logits_shared: tilelang.layout.make_swizzled_layout(up_logits_shared)})
            # T.annotate_layout({routed_expert_down_shared: tilelang.layout.make_swizzled_layout(routed_expert_down_shared)})

            m_start_padded = bx * block_token

            # for i in range(group_count):
            #     in_cur_group_idx = (m_start_padded >= group_padded_offsets[i])
            #     cur_group_idx[0] = T.if_then_else(in_cur_group_idx, i, cur_group_idx[0])
            cur_group_idx[0] = group_idx_for_bx[bx]

            cur_group_size[0] = group_sizes[cur_group_idx[0]]
            m_start = m_start_padded - group_padded_offsets[cur_group_idx[0]] + group_offsets[cur_group_idx[0]]
            actual_rows = T.max(0, T.min(block_token, cur_group_size[0] - (m_start_padded - group_padded_offsets[cur_group_idx[0]])))

            T.clear(output_local)

            for k in T.Pipelined(T.ceildiv(dexpert, block_dexpert), num_stages=num_stages):
                T.copy(up_logits[m_start:m_start + block_token, k * block_dexpert:(k + 1) * block_dexpert], up_logits_shared, coalesced_width=coalesced_width)
                T.copy(routed_expert_down[cur_group_idx[0], by * block_dhidden:(by + 1) * block_dhidden, k * block_dexpert:(k + 1) * block_dexpert], routed_expert_down_shared, coalesced_width=coalesced_width)                
                T.gemm(up_logits_shared, routed_expert_down_shared, output_local, k_pack=k_pack, transpose_B=True)

            for i, j in T.Parallel(block_token, block_dhidden):
                with T.If(i < actual_rows), T.Then():
                    output[m_start + i, by * block_dhidden + j] = output_local[i, j] * routed_expert_weights[m_start + i]

    return kernel



class Expert(nn.Module):
    def __init__(self, config: Dict, gate: torch.Tensor, up: torch.Tensor, down: torch.Tensor, 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.device = torch.device("cuda")

        # self.W_gate = nn.Linear(self.d_hidden, self.d_expert, bias=False)
        self.W_gate_weight = gate.t().contiguous().to(self.device)
        # self.W_up = nn.Linear(self.d_hidden, self.d_expert, bias=False)
        self.W_up_weight = up.t().contiguous().to(self.device)
        # self.W_down = nn.Linear(self.d_expert, self.d_hidden, bias=False)
        self.W_down_weight = down.t().contiguous().to(self.device)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate = self.act_fn(x @ self.W_gate_weight)
        out = (gate * (x @ self.W_up_weight)) @ self.W_down_weight
        return out


class MoEGate(nn.Module):
    def __init__(self, config: Dict, weights: 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"]

        self.W_g_weight = weights['router.weight'].t()

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        logits = x @ self.W_g_weight
        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


class MoE(nn.Module):
    def __init__(self, config: Dict, shared_kernel: tilelang.JITKernel, routed_kernel: tilelang.JITKernel, weights: Dict, padding_M: int = 128):
        super().__init__()
        self.config = config
        self.shared_kernel = shared_kernel
        self.routed_kernel = routed_kernel
        self.padding_M = padding_M
        self.experts = nn.ModuleList([
            Expert(config, gate=weights[f'experts.{i}.0.weight'], up=weights[f'experts.{i}.1.weight'], down=weights[f'experts.{i}.2.weight'])
            for i in range(config["n_routed_experts"])
        ])
        self.device = torch.device("cuda")
        self.gating_network = MoEGate(config, weights).to(self.device)
        shared_expert_dim = config["d_expert"] * config["n_shared_experts"]
        self.shared_expert = Expert(config=config, gate=weights['shared_experts.0.weight'], up=weights['shared_experts.1.weight'], down=weights['shared_experts.2.weight'], d_expert=shared_expert_dim).to(self.device)
        self.expert_cache = torch.zeros((config["batch_size"] * config["seq_len"], config["d_hidden"]), dtype=torch.float16, device=self.device)
        self.stacked_expert_w_gate = torch.stack([expert.W_gate_weight for expert in self.experts], dim=0)
        self.stacked_expert_w_up = torch.stack([expert.W_up_weight for expert in self.experts], dim=0)
        # self.stacked_expert_w_gate_up_cat = torch.cat([self.stacked_expert_w_gate, self.stacked_expert_w_up], dim=-1)
        self.stacked_expert_w_down = torch.stack([expert.W_down_weight for expert in self.experts], dim=0)
        self.stacked_expert_tokens = torch.empty((config["batch_size"] * config["seq_len"] * config["n_experts_per_token"], self.config["d_hidden"]), dtype=torch.float16, device=self.device)
        self.stacked_expert_weights = torch.empty((config["batch_size"] * config["seq_len"] * config["n_experts_per_token"], 1), dtype=torch.float16, device=self.device)
        self.stacked_expert_tokens_idxs = torch.empty((config["batch_size"] * config["seq_len"] * config["n_experts_per_token"], 1), dtype=torch.int64, device=self.device)

        self.up_logits_shared = torch.empty((config["batch_size"] * config["seq_len"], self.config["d_expert"]), dtype=torch.float16, device=self.device)
        self.expert_output_shared = torch.empty((config["batch_size"] * config["seq_len"], self.config["d_hidden"]), dtype=torch.float16, device=self.device)

        # self.gate_up_cat_logits_routed = torch.empty((config["batch_size"] * config["seq_len"] * config["n_experts_per_token"], self.config["d_expert"] * 2), dtype=torch.float16, device=self.device)
        self.up_logits_routed = torch.empty((config["batch_size"] * config["seq_len"] * config["n_experts_per_token"], self.config["d_expert"]), dtype=torch.float16, device=self.device)
        self.expert_output_routed = torch.empty((config["batch_size"] * config["seq_len"] * config["n_experts_per_token"], self.config["d_hidden"]), dtype=torch.float16, device=self.device)

    @torch.no_grad()
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_shape = x.shape
        batch_size, seq_len, hidden_dim = x.shape
        expert_indices, expert_scores = self.gating_network(x)
        flat_expert_indices = expert_indices.view(-1)
        flat_expert_weights = expert_scores.view(-1, 1)
        x_flat = x.view(-1, hidden_dim)

        # shared_output = torch.empty_like(x_flat, device=self.device)
        # up_logits_shared = torch.empty((batch_size * seq_len, self.config["d_expert"]), dtype=x.dtype, device=self.device)

        # Prepare for grouped GEMM
        idxs = flat_expert_indices.argsort()
        counts = flat_expert_indices.bincount().cpu().numpy()
        # counts = flat_expert_indices.bincount()
        tokens_per_expert = counts.cumsum()
        # tokens_per_expert = torch.cumsum(counts, dim=0)
        num_per_tok = self.config["n_experts_per_token"]
        token_idxs = idxs // num_per_tok

        # Get stacked expert tokens and expert weights
        # stacked_expert_tokens = []
        # stacked_expert_tokens_idxs = []
        # stacked_expert_weights = []

        for expert_id, end_idx in enumerate(tokens_per_expert):
            start_idx = 0 if expert_id == 0 else tokens_per_expert[expert_id - 1]
            if start_idx == end_idx:
                continue

            exp_token_idxs = token_idxs[start_idx:end_idx]
            expert_tokens = x_flat[exp_token_idxs]

            self.stacked_expert_tokens[start_idx:end_idx] = expert_tokens
            self.stacked_expert_tokens_idxs[start_idx:end_idx, 0] = exp_token_idxs
            self.stacked_expert_weights[start_idx:end_idx] = flat_expert_weights[idxs[start_idx:end_idx]]
            # stacked_expert_tokens_idxs.append(exp_token_idxs) # [bs * seq_len * n_experts_per_token]
            # stacked_expert_tokens.append(expert_tokens)
            # stacked_expert_weights.append(flat_expert_weights[idxs[start_idx:end_idx]])
        
        # stacked_expert_tokens = torch.cat(stacked_expert_tokens, dim=0).to(self.device)
        # stacked_expert_tokens_idxs = torch.cat(stacked_expert_tokens_idxs, dim=0).to(self.device)
        # stacked_expert_weights = torch.cat(stacked_expert_weights, dim=0).to(self.device)

        group_sizes = torch.tensor(counts, dtype=torch.int32, device=self.device)
        group_offset = torch.tensor(tokens_per_expert - counts, dtype=torch.int32, device=self.device)

        # group_padded_offsets = torch.empty(len(group_sizes), dtype=torch.int32, device=self.device)
        group_padded_offsets = [0 for _ in range(len(group_sizes))]
        for i in range(1, len(group_sizes)):
            group_padded_offsets[i] = group_padded_offsets[i - 1] + math.ceil((counts[i - 1] + 1) / self.padding_M) * self.padding_M

        block_token = 128
        M = math.ceil(self.config["batch_size"] * self.config["seq_len"] * self.config["n_experts_per_token"] / block_token) + self.config["n_routed_experts"]
        group_idx_for_bx = [0 for _ in range(M)]
        
        for bx in range(M):
            m_start_padded = bx * block_token
            for i in range(self.config["n_routed_experts"]):
                if m_start_padded >= group_padded_offsets[i]:
                    group_idx_for_bx[bx] = i

        group_padded_offsets = torch.tensor(group_padded_offsets, dtype=torch.int32, device=self.device)
        group_idx_for_bx = torch.tensor(group_idx_for_bx, dtype=torch.int32, device=self.device)
        # up_logits_routed = torch.empty((x_flat.shape[0] * num_per_tok, self.config["d_expert"]), dtype=x.dtype, device=self.device)
        # expert_output_routed = torch.empty((x_flat.shape[0] * num_per_tok, self.config["d_hidden"]), dtype=x.dtype, device=self.device)

        shared_stream = torch.cuda.Stream()
        routed_stream = torch.cuda.default_stream()
        torch.cuda.synchronize()

        with torch.cuda.stream(routed_stream):
            # Tilelang version: Grouped GEMM
            self.routed_kernel(self.stacked_expert_tokens, self.stacked_expert_w_gate, self.stacked_expert_w_up, self.stacked_expert_w_down, self.stacked_expert_weights, group_sizes, group_offset, group_padded_offsets, group_idx_for_bx, self.up_logits_routed, self.expert_output_routed)
            # self.routed_kernel(stacked_expert_tokens, self.stacked_expert_w_gate_up_cat, self.stacked_expert_w_down, stacked_expert_weights, group_sizes, group_offset, group_padded_offsets, self.gate_up_cat_logits_routed, self.up_logits_routed, self.expert_output_routed)

            # Scatter reduce
            self.expert_cache = torch.scatter_reduce(
                self.expert_cache,
                0,
                self.stacked_expert_tokens_idxs.view(-1, 1).repeat(1, x_flat.shape[-1]),
                self.expert_output_routed,
                reduce='sum'
            )
            routed_output = self.expert_cache.view(*orig_shape)

        with torch.cuda.stream(shared_stream):
            # Use tilelang to implement the moe inference of shared experts
            self.shared_kernel(x_flat, self.shared_expert.W_gate_weight, self.shared_expert.W_up_weight, self.shared_expert.W_down_weight, self.up_logits_shared, self.expert_output_shared)
            shared_output = self.expert_output_shared.view(*orig_shape)

            # shared_output = self.shared_expert(x)

        torch.cuda.synchronize()

        return routed_output + shared_output

    # @torch.no_grad()
    # def moe_infer(self,
    #               x: torch.Tensor,
    #               flat_expert_indices: torch.Tensor,
    #               flat_expert_weights: torch.Tensor
    #              ) -> torch.Tensor:

    #     return self.expert_cache


def custom_kernel(data: input_t) -> output_t:
    """
    Submission template for DeepSeek-style Mixture of Experts using PyTorch.
    
    Args:
        data: Tuple of (input: torch.Tensor, weights: Dict[str, torch.Tensor], config: Dict)
            - input: Input tensor of shape [batch_size, seq_len, hidden_size]
            - 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

    dtype_str = "float16"

    # start_time = time.perf_counter_ns()
    # def get_configs():
    #     import itertools
    #     BLOCK_M = [32, 64, 128]
    #     BLOCK_N = [32, 64, 128]
    #     BLOCK_K = [32, 64, 128]
    #     num_stages = [0, 1, 2]
    #     threads = [128, 256]
    #     k_pack = [1, 2, 4, 8]
    #     coalesced_width = [1, 2, 4]
    #     # num_stages = [1]
    #     # threads = [128]
    #     # k_pack = [1]
    #     # coalesced_width = [4]
    #     _configs = list(itertools.product(BLOCK_M, BLOCK_N, BLOCK_K, num_stages, threads, k_pack, coalesced_width))
        
    #     return [{
    #         "block_M": c[0],
    #         "block_N": c[1],
    #         "block_K": c[2],
    #         "num_stages": c[3],
    #         "threads": c[4],
    #         "k_pack": c[5],
    #         "coalesced_width": c[6],
    #     } for c in _configs]
    
    # def wrapped_kernel(block_M=None, block_N=None, block_K=None, num_stages=None, threads=None, k_pack=None, coalesced_width=None):
    #     return moe_forward_tilelang_routed(config["d_hidden"], config["d_expert"], config["n_routed_experts"], dtype=dtype_str, group_sum=config["batch_size"] * config["seq_len"] * config["n_experts_per_token"], group_count=config["n_routed_experts"], block_token=block_M, block_dhidden=block_N, block_dexpert=block_K, threads=threads, num_stages=num_stages, k_pack=k_pack, coalesced_width=coalesced_width)
    
    # autotuner = AutoTuner.from_kernel(
    #     kernel=wrapped_kernel,
    #     configs=get_configs()
    # ).set_compile_args(
    #     out_idx=[-1],
    #     target="hip",
    # )
    # tune_result = autotuner.run(warmup=3, rep=10)
    # best_latency = tune_result.latency
    # best_config = tune_result.config
    # print(f"Best latency: {best_latency}")
    # print(f"Best config: {best_config}")
    
    shared_kernel = moe_forward_tilelang_shared(config["d_hidden"], config["d_expert"], config["n_shared_experts"], dtype=dtype_str, num_tokens=config["batch_size"] * config["seq_len"])
    # shared_kernel = tilelang.compile(shared_program, out_idx=[-1], pass_configs={"tl.disable_tma_lower": True, "tl.disable_warp_specialized": True})
    routed_kernel = moe_forward_tilelang_routed(config["d_hidden"], config["d_expert"], config["n_routed_experts"], dtype=dtype_str, group_sum=config["batch_size"] * config["seq_len"] * config["n_experts_per_token"], group_count=config["n_routed_experts"], block_token=128, block_dhidden=128, block_dexpert=128, threads=256, num_stages=1, k_pack=1, coalesced_width=2)
    # routed_kernel = tilelang.compile(routed_program, out_idx=[-1], pass_configs={"tl.disable_tma_lower": True, "tl.disable_warp_specialized": True})

    # end_time = time.perf_counter_ns()
    # compile_time = end_time - start_time
    # print(f"Time taken to compile kernels: {compile_time} ns")

    # start_time = time.perf_counter_ns()
    moe = MoE(config, shared_kernel, routed_kernel, weights, padding_M=128)
    # end_time = time.perf_counter_ns()
    # init_time = end_time - start_time
    # print(f"Time taken to initialize model: {init_time} ns")

    # start_time = time.perf_counter_ns()
    # Run the model
    # with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
    output = moe(input_tensor)
    # end_time = time.perf_counter_ns()
    # inference_time = end_time - start_time
    # print(f"Time taken to run model: {inference_time} ns")

    # print(prof.key_averages().table(sort_by="cpu_time_total"))
    # prof.export_chrome_trace("moe_multistream.json")


    return output
python eval.py leaderboard /tmp/tmpkhuvfll9
WARNING: pip is being invoked by an old script wrapper. This will fail in a future version of pip.
Please see https://github.com/pypa/pip/issues/5599 for advice on fixing the underlying issue.
To avoid this problem you can invoke Python with '-m pip' instead of running pip directly.

Defaulting to user installation because normal site-packages is not writeable
Requirement already satisfied: tilelang-rocm in /home/runner/.local/lib/python3.10/site-packages (0.1.4.post13)
Requirement already satisfied: Cython in /home/runner/.local/lib/python3.10/site-packages (from tilelang-rocm) (3.1.1)
Requirement already satisfied: numpy>=1.23.5 in /usr/local/lib/python3.10/dist-packages (from tilelang-rocm) (1.26.4)
Requirement already satisfied: tqdm>=4.62.3 in /home/runner/.local/lib/python3.10/site-packages (from tilelang-rocm) (4.67.1)
Requirement already satisfied: typing_extensions>=4.10.0 in /usr/local/lib/python3.10/dist-packages (from tilelang-rocm) (4.12.2)
Requirement already satisfied: cloudpickle in /home/runner/.local/lib/python3.10/site-packages (from tilelang-rocm) (3.1.1)
Requirement already satisfied: ml_dtypes in /home/runner/.local/lib/python3.10/site-packages (from tilelang-rocm) (0.5.1)
Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from tilelang-rocm) (7.0.0)
Requirement already satisfied: torch in /usr/local/lib/python3.10/dist-packages (from tilelang-rocm) (2.8.0.dev20250602+rocm6.3)
Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch->tilelang-rocm) (3.16.1)
Requirement already satisfied: sympy>=1.13.3 in /usr/local/lib/python3.10/dist-packages (from torch->tilelang-rocm) (1.13.3)
Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch->tilelang-rocm) (3.4.2)
Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch->tilelang-rocm) (3.1.4)
Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch->tilelang-rocm) (2024.10.0)
Requirement already satisfied: pytorch-triton-rocm==3.3.1+gitc8757738 in /usr/local/lib/python3.10/dist-packages (from torch->tilelang-rocm) (3.3.1+gitc8757738)
Requirement already satisfied: setuptools>=40.8.0 in /usr/lib/python3/dist-packages (from pytorch-triton-rocm==3.3.1+gitc8757738->torch->tilelang-rocm) (59.6.0)
Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from sympy>=1.13.3->torch->tilelang-rocm) (1.3.0)
Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch->tilelang-rocm) (2.1.5)
