#!POPCORN leaderboard amd-mixture-of-experts
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple

import triton
import triton.language as tl

from task import input_t, output_t # Assuming task.py defines these types

HARDCODED_FFN_CONFIGS_SHARED = {
    (2048, 7168, 7168, 2048): (128, 128, 64, 64, 4, 8, 2),
    (8192, 7168, 7168, 2048): (256, 128, 64, 64, 4, 8, 2),
}
DEFAULT_FFN_CONFIG_SHARED = (128, 64, 64, 64, 4, 4, 2)

HARDCODED_FFN_CONFIGS_ROUTED_BATCHED = {
    (7168, 2048, 7168, 2048): (128, 128, 128, 128, 4, 8, 1),
    (7168, 2048, 7168, 8192): (128, 128, 128, 128, 4, 8, 1),
}
DEFAULT_FFN_CONFIG_ROUTED_BATCHED = (128, 64, 64, 64, 4, 4, 2)


@triton.jit
def silu(x): return (x * tl.sigmoid(x.to(tl.float32))).to(tl.float16) # Assuming float16 from context

# Kernel for single FFN (e.g., for Shared Expert) - Remains the same
@triton.jit
def fused_ffn_kernel_v2(
    # Pointers to Matrices
    x_ptr, w_gate_ptr, w_up_ptr, w_down_ptr, output_ptr,
    # Matrix dimensions
    M, N_out, K_in, K_inter,
    # Strides
    stride_x_m, stride_x_k,
    stride_w_gate_k, stride_w_gate_inter,
    stride_w_up_k, stride_w_up_inter,
    stride_w_down_inter, stride_w_down_n,
    stride_out_m, stride_out_n,
    # Meta-parameters
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N_OUT: tl.constexpr,
    BLOCK_SIZE_K_IN: tl.constexpr, BLOCK_SIZE_K_INTER: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N_out, BLOCK_SIZE_N_OUT)
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + (pid % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    offs_m_block_local = tl.arange(0, BLOCK_SIZE_M)
    offs_m_global = pid_m * BLOCK_SIZE_M + offs_m_block_local
    
    offs_n_out = pid_n * BLOCK_SIZE_N_OUT + tl.arange(0, BLOCK_SIZE_N_OUT)
    x_start_ptr = x_ptr + offs_m_global[:, None] * stride_x_m # Corrected to use offs_m_global
    
    weight_dtype = w_gate_ptr.dtype.element_ty
    accum_dtype = tl.float32
    acc_down = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N_OUT), dtype=accum_dtype)
    w_down_start_ptr = w_down_ptr + offs_n_out[None, :] * stride_w_down_n

    for k_inter_start in range(0, tl.cdiv(K_inter, BLOCK_SIZE_K_INTER)):
        offs_k_inter = k_inter_start * BLOCK_SIZE_K_INTER + tl.arange(0, BLOCK_SIZE_K_INTER)
        acc_gate_fp32 = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_K_INTER), dtype=accum_dtype)
        acc_up_fp32 = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_K_INTER), dtype=accum_dtype)
        w_gate_k_inter_ptr = w_gate_ptr + offs_k_inter[None, :] * stride_w_gate_inter
        w_up_k_inter_ptr = w_up_ptr + offs_k_inter[None, :] * stride_w_up_inter

        for k_in_start in range(0, tl.cdiv(K_in, BLOCK_SIZE_K_IN)):
            offs_k_in = k_in_start * BLOCK_SIZE_K_IN + tl.arange(0, BLOCK_SIZE_K_IN)
            
            current_x_ptrs = x_start_ptr + offs_k_in[None, :] * stride_x_k
            # Mask for input x: valid rows in the input M and valid K_in dimension
            x_mask = (offs_m_global[:, None] < M) & (offs_k_in[None, :] < K_in)
            x_block = tl.load(current_x_ptrs, mask=x_mask, other=0.0)
            
            w_gate_ptrs = w_gate_k_inter_ptr + offs_k_in[:, None] * stride_w_gate_k
            w_gate_mask = (offs_k_in[:, None] < K_in) & (offs_k_inter[None, :] < K_inter)
            w_gate_block = tl.load(w_gate_ptrs, mask=w_gate_mask, other=0.0)
            acc_gate_fp32 += tl.dot(x_block, w_gate_block, out_dtype=accum_dtype)
            
            w_up_ptrs = w_up_k_inter_ptr + offs_k_in[:, None] * stride_w_up_k
            w_up_mask = (offs_k_in[:, None] < K_in) & (offs_k_inter[None, :] < K_inter)
            w_up_block = tl.load(w_up_ptrs, mask=w_up_mask, other=0.0)
            acc_up_fp32 += tl.dot(x_block, w_up_block, out_dtype=accum_dtype)

        acc_gate_fp16 = acc_gate_fp32.to(weight_dtype)
        acc_up_fp16 = acc_up_fp32.to(weight_dtype)
        gate_activated_fp16 = silu(acc_gate_fp16)
        intermediate_block_fp16 = gate_activated_fp16 * acc_up_fp16
        
        w_down_ptrs = w_down_start_ptr + offs_k_inter[:, None] * stride_w_down_inter
        w_down_mask = (offs_k_inter[:, None] < K_inter) & (offs_n_out[None, :] < N_out)
        w_down_block = tl.load(w_down_ptrs, mask=w_down_mask, other=0.0)
        acc_down += tl.dot(intermediate_block_fp16, w_down_block, out_dtype=accum_dtype)

    final_output_ptrs = output_ptr + offs_m_global[:, None] * stride_out_m + offs_n_out[None, :] * stride_out_n
    output_mask = (offs_m_global[:, None] < M) & (offs_n_out[None, :] < N_out)
    tl.store(final_output_ptrs, acc_down.to(output_ptr.dtype.element_ty), mask=output_mask)

# New Batched FFN Kernel for Routed Experts
@triton.jit
def batched_fused_ffn_kernel_expert_dim(
    # Pointers to Batched/Stacked Matrices
    batched_x_ptr,        # Shape: [E, max_tokens_per_expert, K_in_dim]
    all_w_gate_ptr,       # Shape: [E, K_in_dim, K_inter_dim]
    all_w_up_ptr,         # Shape: [E, K_in_dim, K_inter_dim]
    all_w_down_ptr,       # Shape: [E, K_inter_dim, N_out_dim]
    batched_output_ptr,   # Shape: [E, max_tokens_per_expert, N_out_dim]
    expert_counts_ptr,    # Shape: [E], actual num tokens for each expert (int32)
    # Matrix dimensions (scalar values for THIS expert's FFN structure)
    max_tokens_per_expert_padded_dim, # Padded M dim for array bounds for batched_x_ptr
    N_out_dim, K_in_dim, K_inter_dim,
    # Strides for batched_x_ptr [E, M_padded, K_in]
    stride_b_x_e, stride_b_x_m, stride_b_x_k,
    # Strides for all_w_gate_ptr [E, K_in, K_inter]
    stride_all_w_gate_e, stride_all_w_gate_k, stride_all_w_gate_inter,
    # Strides for all_w_up_ptr [E, K_in, K_inter]
    stride_all_w_up_e, stride_all_w_up_k, stride_all_w_up_inter,
    # Strides for all_w_down_ptr [E, K_inter, N_out]
    stride_all_w_down_e, stride_all_w_down_inter, stride_all_w_down_n,
    # Strides for batched_output_ptr [E, M_padded, N_out]
    stride_b_out_e, stride_b_out_m, stride_b_out_n,
    # Meta-parameters from config
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N_OUT: tl.constexpr,
    BLOCK_SIZE_K_IN: tl.constexpr, BLOCK_SIZE_K_INTER: tl.constexpr,
    GROUP_SIZE_M_CFG: tl.constexpr, # Renamed from GROUP_SIZE_M to avoid conflict
):
    expert_id = tl.program_id(axis=0)
    
    M_curr_expert_actual = tl.load(expert_counts_ptr + expert_id) # Assumes int32 data
    if M_curr_expert_actual == 0:
        return

    pid_token_block_flat = tl.program_id(axis=1)

    num_pid_m_local_padded = tl.cdiv(max_tokens_per_expert_padded_dim, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N_out_dim, BLOCK_SIZE_N_OUT)
    
    num_pid_in_group = GROUP_SIZE_M_CFG * num_pid_n
    group_id = pid_token_block_flat // num_pid_in_group
    first_pid_m_local = group_id * GROUP_SIZE_M_CFG
    
    # Calculate group_size_m_local carefully based on padded M dimension
    # group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    group_size_m_local = min(num_pid_m_local_padded - first_pid_m_local, GROUP_SIZE_M_CFG)

    # If group_size_m_local is non-positive (e.g. if num_pid_m_local_padded < first_pid_m_local),
    # this can lead to issues. Add a check.
    if group_size_m_local <= 0:
        return

    pid_m_local = first_pid_m_local + (pid_token_block_flat % group_size_m_local)
    pid_n = (pid_token_block_flat % num_pid_in_group) // group_size_m_local
    
    if pid_m_local * BLOCK_SIZE_M >= M_curr_expert_actual: # Check against actual tokens
        return

    # Select data for the current expert_id
    x_ptr_curr_expert = batched_x_ptr + expert_id * stride_b_x_e
    w_gate_ptr_curr_expert = all_w_gate_ptr + expert_id * stride_all_w_gate_e
    w_up_ptr_curr_expert = all_w_up_ptr + expert_id * stride_all_w_up_e
    w_down_ptr_curr_expert = all_w_down_ptr + expert_id * stride_all_w_down_e
    output_ptr_curr_expert = batched_output_ptr + expert_id * stride_b_out_e

    offs_m_block_local_coords = tl.arange(0, BLOCK_SIZE_M) # Coordinates within a processing block [0, ..., BLOCK_SIZE_M-1]
    # Offsets relative to the start of this expert's segment in batched_x_ptr
    offs_m_in_expert_segment = pid_m_local * BLOCK_SIZE_M + offs_m_block_local_coords
    
    offs_n_out = pid_n * BLOCK_SIZE_N_OUT + tl.arange(0, BLOCK_SIZE_N_OUT)
    # Start pointer for x for the current M-block of this expert
    x_block_start_ptr = x_ptr_curr_expert + offs_m_in_expert_segment[:, None] * stride_b_x_m
    
    weight_dtype = w_gate_ptr_curr_expert.dtype.element_ty
    accum_dtype = tl.float32
    acc_down = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N_OUT), dtype=accum_dtype)
    
    # Start pointer for w_down for the current N-block of this expert
    w_down_block_start_ptr = w_down_ptr_curr_expert + offs_n_out[None, :] * stride_all_w_down_n

    for k_inter_start in range(0, tl.cdiv(K_inter_dim, BLOCK_SIZE_K_INTER)):
        offs_k_inter = k_inter_start * BLOCK_SIZE_K_INTER + tl.arange(0, BLOCK_SIZE_K_INTER)
        acc_gate_fp32 = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_K_INTER), dtype=accum_dtype)
        acc_up_fp32 = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_K_INTER), dtype=accum_dtype)
        
        w_gate_k_inter_ptr = w_gate_ptr_curr_expert + offs_k_inter[None, :] * stride_all_w_gate_inter
        w_up_k_inter_ptr = w_up_ptr_curr_expert + offs_k_inter[None, :] * stride_all_w_up_inter

        for k_in_start in range(0, tl.cdiv(K_in_dim, BLOCK_SIZE_K_IN)):
            offs_k_in = k_in_start * BLOCK_SIZE_K_IN + tl.arange(0, BLOCK_SIZE_K_IN)
            
            current_x_ptrs = x_block_start_ptr + offs_k_in[None, :] * stride_b_x_k
            x_mask = (offs_m_in_expert_segment[:, None] < M_curr_expert_actual) & (offs_k_in[None, :] < K_in_dim)
            x_block = tl.load(current_x_ptrs, mask=x_mask, other=0.0)
            
            w_gate_ptrs = w_gate_k_inter_ptr + offs_k_in[:, None] * stride_all_w_gate_k
            w_gate_mask = (offs_k_in[:, None] < K_in_dim) & (offs_k_inter[None, :] < K_inter_dim)
            w_gate_block = tl.load(w_gate_ptrs, mask=w_gate_mask, other=0.0)
            acc_gate_fp32 += tl.dot(x_block, w_gate_block, out_dtype=accum_dtype)
            
            w_up_ptrs = w_up_k_inter_ptr + offs_k_in[:, None] * stride_all_w_up_k
            w_up_mask = (offs_k_in[:, None] < K_in_dim) & (offs_k_inter[None, :] < K_inter_dim)
            w_up_block = tl.load(w_up_ptrs, mask=w_up_mask, other=0.0)
            acc_up_fp32 += tl.dot(x_block, w_up_block, out_dtype=accum_dtype)

        acc_gate_fp16 = acc_gate_fp32.to(weight_dtype)
        acc_up_fp16 = acc_up_fp32.to(weight_dtype)
        gate_activated_fp16 = silu(acc_gate_fp16)
        intermediate_block_fp16 = gate_activated_fp16 * acc_up_fp16
        
        w_down_ptrs = w_down_block_start_ptr + offs_k_inter[:, None] * stride_all_w_down_inter
        w_down_mask = (offs_k_inter[:, None] < K_inter_dim) & (offs_n_out[None, :] < N_out_dim)
        w_down_block = tl.load(w_down_ptrs, mask=w_down_mask, other=0.0)
        acc_down += tl.dot(intermediate_block_fp16, w_down_block, out_dtype=accum_dtype)

    final_output_ptrs = output_ptr_curr_expert + offs_m_in_expert_segment[:, None] * stride_b_out_m + offs_n_out[None, :] * stride_b_out_n
    output_mask = (offs_m_in_expert_segment[:, None] < M_curr_expert_actual) & (offs_n_out[None, :] < N_out_dim)
    tl.store(final_output_ptrs, acc_down.to(batched_output_ptr.dtype.element_ty), mask=output_mask)


class TritonExpert(nn.Module): # For Shared Expert
    def __init__(self, config: Dict, W_gate: torch.Tensor, W_up: torch.Tensor, W_down: torch.Tensor, d_expert_actual: int, seq_len_for_config_key: int):
        super().__init__()
        self.d_hidden: int = config["d_hidden"] # K_in for FFN
        self.d_expert_actual = d_expert_actual # K_inter for FFN
        # IMPORTANT: Weights must be in [K_in, K_inter] or [K_inter, N_out] format for the kernel.
        # PyTorch nn.Linear weights are [out_features, in_features].
        # If W_gate, W_up, W_down come directly from nn.Linear.weight, they need to be transposed.
        # Example: nn.Linear(d_hidden, d_expert_actual) has weight [d_expert_actual, d_hidden].
        # Kernel needs W_gate as [d_hidden, d_expert_actual]. So, W_gate.T.
        # Current code implies weights are already in kernel-expected format. Let's stick to that.
        self.W_gate = W_gate
        self.W_up = W_up
        self.W_down = W_down
        self.seq_len_for_config_key = seq_len_for_config_key

    def forward(self, x: torch.Tensor, out_tensor_view: torch.Tensor) -> None:
        M_tokens, K_in_dim = x.shape
        K_inter_dim = self.d_expert_actual
        N_out_dim = self.d_hidden 

        x = x.contiguous()
        
        ffn_config_key = (M_tokens, N_out_dim, K_in_dim, K_inter_dim)
        selected_params = HARDCODED_FFN_CONFIGS_SHARED.get(ffn_config_key, DEFAULT_FFN_CONFIG_SHARED)
        
        BLOCK_SIZE_M, BLOCK_SIZE_N_OUT, BLOCK_SIZE_K_IN, BLOCK_SIZE_K_INTER, GROUP_SIZE_M_cfg, num_warps, num_stages = selected_params
        
        grid = (
            triton.cdiv(M_tokens, BLOCK_SIZE_M) * triton.cdiv(N_out_dim, BLOCK_SIZE_N_OUT),
        )

        fused_ffn_kernel_v2[grid](
            x, self.W_gate, self.W_up, self.W_down, 
            out_tensor_view,
            M_tokens, N_out_dim, K_in_dim, K_inter_dim,
            x.stride(0), x.stride(1),
            self.W_gate.stride(0), self.W_gate.stride(1), # Assumes W_gate is [K_in, K_inter]
            self.W_up.stride(0), self.W_up.stride(1),     # Assumes W_up is [K_in, K_inter]
            self.W_down.stride(0), self.W_down.stride(1), # Assumes W_down is [K_inter, N_out]
            out_tensor_view.stride(0), out_tensor_view.stride(1),
            BLOCK_SIZE_M=BLOCK_SIZE_M, BLOCK_SIZE_N_OUT=BLOCK_SIZE_N_OUT,
            BLOCK_SIZE_K_IN=BLOCK_SIZE_K_IN, BLOCK_SIZE_K_INTER=BLOCK_SIZE_K_INTER,
            GROUP_SIZE_M=GROUP_SIZE_M_cfg,
            num_warps=num_warps,
            num_stages=num_stages
        )

class MoEGate(nn.Module):
    def __init__(self, config: Dict, W_g_weight: torch.Tensor):
        super().__init__()
        self.top_k: int = config["n_experts_per_token"]
        # self.n_experts: int = config["n_routed_experts"] # Not directly used here
        self.register_buffer('W_g', W_g_weight) # W_g is [n_routed_experts, d_hidden]

    @torch.compile
    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        logits = F.linear(x, self.W_g) # x: [tokens, d_hidden], W_g: [n_routed_experts, d_hidden] -> logits: [tokens, n_routed_experts]
        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 TritonMoE(nn.Module):
    def __init__(self, config: Dict, weights: Dict[str, torch.Tensor], seq_len_for_config_key: int):
        super().__init__()
        self.num_experts = config["n_routed_experts"]
        self.top_k = config["n_experts_per_token"]
        self.d_hidden = config["d_hidden"]
        self.d_expert = config["d_expert"] # K_inter for routed experts
        self.seq_len_for_config_key = seq_len_for_config_key

        self.gating_network = MoEGate(config, weights['router.weight'])
        
        list_w_gate, list_w_up, list_w_down = [], [], []
        for i in range(self.num_experts):
            list_w_gate.append(weights[f'experts.{i}.0.weight'])
            list_w_up.append(weights[f'experts.{i}.1.weight'])
            list_w_down.append(weights[f'experts.{i}.2.weight'])
        
        # Store as non-trainable parameters or buffers if not training
        self.W_gate_all_routed = torch.stack(list_w_gate).contiguous()
        self.W_up_all_routed = torch.stack(list_w_up).contiguous()
        self.W_down_all_routed = torch.stack(list_w_down).contiguous()
        
        # Shared expert
        shared_expert_inter_dim = config["d_expert"] * config["n_shared_experts"]
        w_gate_shared = weights['shared_experts.0.weight']
        w_up_shared = weights['shared_experts.1.weight']
        w_down_shared = weights['shared_experts.2.weight']
        self.shared_expert_module = TritonExpert(config, w_gate_shared, w_up_shared, w_down_shared, 
                                                 shared_expert_inter_dim, seq_len_for_config_key)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_shape = x.shape
        x_flat = x.view(-1, self.d_hidden)
        num_total_tokens_moe_input = x_flat.shape[0]

        expert_indices, expert_scores = self.gating_network(x_flat)

        flat_expert_indices = expert_indices.view(-1)
        flat_expert_scores = expert_scores.view(-1, 1)

        token_ids_for_flat_assignments = torch.arange(num_total_tokens_moe_input, device=x.device).repeat_interleave(self.top_k)
        sort_perm = flat_expert_indices.argsort()

        perm_expert_indices = flat_expert_indices[sort_perm]
        perm_token_ids = token_ids_for_flat_assignments[sort_perm]
        perm_expert_scores = flat_expert_scores[sort_perm]
        perm_x_flat = x_flat[perm_token_ids] # Shape: [total_routed_assignments, d_hidden]

        # --- Routed Experts Computation (Batched Kernel) ---
        if perm_x_flat.shape[0] > 0: # Only proceed if there are tokens for routed experts
            expert_counts = torch.bincount(perm_expert_indices.to(torch.int32), minlength=self.num_experts)
            max_tokens_per_expert = expert_counts.max().item()

            batched_x_input = torch.zeros(self.num_experts, max_tokens_per_expert, self.d_hidden, 
                                          device=x.device, dtype=x.dtype)
            
            expert_boundaries_cpu_list = torch.searchsorted(perm_expert_indices, torch.arange(self.num_experts + 1, device=x.device)).tolist()

            for i in range(self.num_experts):
                start_idx_in_perm = expert_boundaries_cpu_list[i]
                end_idx_in_perm = expert_boundaries_cpu_list[i+1]
                num_tokens_for_this_expert = end_idx_in_perm - start_idx_in_perm
                if num_tokens_for_this_expert > 0:
                    tokens_slice = perm_x_flat.narrow(0, start_idx_in_perm, num_tokens_for_this_expert)
                    batched_x_input[i, :num_tokens_for_this_expert, :] = tokens_slice
            
            batched_output_routed = torch.zeros(self.num_experts, max_tokens_per_expert, self.d_hidden,
                                                device=x.device, dtype=x.dtype)

            K_in_dim, K_inter_dim, N_out_dim = self.d_hidden, self.d_expert, self.d_hidden

            ffn_config_key_routed = (K_in_dim, K_inter_dim, N_out_dim, self.seq_len_for_config_key)
            selected_params_routed = HARDCODED_FFN_CONFIGS_ROUTED_BATCHED.get(ffn_config_key_routed, DEFAULT_FFN_CONFIG_ROUTED_BATCHED)
            
            BLOCK_SIZE_M, BLOCK_SIZE_N_OUT, BLOCK_SIZE_K_IN, BLOCK_SIZE_K_INTER, GROUP_SIZE_M_cfg, num_warps, num_stages = selected_params_routed

            grid_dim0_experts = self.num_experts
            grid_dim1_token_blocks_flat = triton.cdiv(max_tokens_per_expert, BLOCK_SIZE_M) * triton.cdiv(N_out_dim, BLOCK_SIZE_N_OUT)
            grid = (grid_dim0_experts, grid_dim1_token_blocks_flat)
            
            # Ensure expert_counts is int32 for the kernel
            expert_counts_int32 = expert_counts.to(torch.int32)

            batched_fused_ffn_kernel_expert_dim[grid](
                batched_x_input,
                self.W_gate_all_routed, self.W_up_all_routed, self.W_down_all_routed,
                batched_output_routed,
                expert_counts_int32,
                max_tokens_per_expert, N_out_dim, K_in_dim, K_inter_dim,
                batched_x_input.stride(0), batched_x_input.stride(1), batched_x_input.stride(2),
                self.W_gate_all_routed.stride(0), self.W_gate_all_routed.stride(1), self.W_gate_all_routed.stride(2),
                self.W_up_all_routed.stride(0), self.W_up_all_routed.stride(1), self.W_up_all_routed.stride(2),
                self.W_down_all_routed.stride(0), self.W_down_all_routed.stride(1), self.W_down_all_routed.stride(2),
                batched_output_routed.stride(0), batched_output_routed.stride(1), batched_output_routed.stride(2),
                BLOCK_SIZE_M=BLOCK_SIZE_M, BLOCK_SIZE_N_OUT=BLOCK_SIZE_N_OUT,
                BLOCK_SIZE_K_IN=BLOCK_SIZE_K_IN, BLOCK_SIZE_K_INTER=BLOCK_SIZE_K_INTER,
                GROUP_SIZE_M_CFG=GROUP_SIZE_M_cfg, # Pass renamed GROUP_SIZE_M from config
                num_warps=num_warps, num_stages=num_stages
            )

            all_expert_outputs_perm_allocated = torch.empty_like(perm_x_flat)
            for i in range(self.num_experts):
                start_idx_in_perm = expert_boundaries_cpu_list[i]
                num_t = expert_counts[i].item() # Same as expert_boundaries_cpu_list[i+1] - start_idx_in_perm
                if num_t > 0:
                    source_data = batched_output_routed[i, :num_t, :]
                    all_expert_outputs_perm_allocated.narrow(0, start_idx_in_perm, num_t).copy_(source_data)
            
            weighted_expert_outputs_perm = all_expert_outputs_perm_allocated * perm_expert_scores
        else: # No tokens routed to any of the self.num_experts
            weighted_expert_outputs_perm = torch.empty((0, self.d_hidden), device=x.device, dtype=x.dtype)
            # perm_token_ids will also be empty, so scatter_reduce will be a no-op for routed experts.

        # --- Shared Expert Computation ---
        shared_expert_output_flat = torch.empty_like(x_flat) # Shape [total_original_tokens, d_hidden]
        self.shared_expert_module(x_flat, shared_expert_output_flat)

        # --- Combine Routed and Shared Expert Outputs ---
        final_output_flat = shared_expert_output_flat # Initialize with shared expert output
        
        if perm_x_flat.shape[0] > 0: # Check if there are any routed expert outputs to scatter
            scatter_indices = perm_token_ids.unsqueeze(-1).expand_as(weighted_expert_outputs_perm)
            final_output_flat.scatter_reduce_(
                dim=0,
                index=scatter_indices,
                src=weighted_expert_outputs_perm,
                reduce="sum",
                include_self=True # Adds to existing values from shared_expert_output_flat
            )
        
        output = final_output_flat.view(orig_shape)
        return output

def custom_kernel(data: input_t) -> output_t:
    input_tensor, weights, config = data
    seq_len_for_config_key = input_tensor.shape[1]
    
    for k, v in weights.items():
        weights[k] = v.contiguous() # This prepares them before passing to TritonMoE

    triton_moe_model = TritonMoE(config, weights, seq_len_for_config_key)
    
    input_tensor = input_tensor.to('cuda')
    triton_moe_model.to('cuda') # Moves stacked params and shared expert params
    
    input_dtype = next(iter(weights.values())).dtype # Infer dtype from weights
    input_tensor = input_tensor.to(input_dtype)
    
    output = triton_moe_model(input_tensor)
    return output


python eval.py leaderboard /tmp/tmpti8j132l