#!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 = {
    # (M_tokens, N_out, K_in, K_inter): (BLOCK_M, BLOCK_N_OUT, BLOCK_K_IN, BLOCK_K_INTER, GROUP_M, WARPS, STAGES)
    (2048, 7168, 7168, 2048): (128, 512, 64, 64, 8, 8, 2), 
    (8192, 7168, 7168, 2048): (128, 512, 64, 64, 8, 8, 2), 
}
DEFAULT_FFN_CONFIG_SHARED = (128, 512, 64, 64, 8, 8, 2)

DEFAULT_FFN_CONFIG_UNIFIED = (128, 256, 64, 64, 8, 4, 2) # BLOCK_M, BLOCK_N_OUT, BLOCK_K_IN, BLOCK_K_INTER, GROUP_M_CFG, WARPS, STAGES
HARDCODED_FFN_CONFIGS_UNIFIED = {
    (16384, 7168, 7168, 2048): (128, 256, 64, 64, 8, 4, 2),
}


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

@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_actual = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    
    # Guard against empty groups if grid is too large (shouldn't happen with tight grid)
    if group_size_m_actual <= 0:
        return

    pid_m_local_offset = (pid % num_pid_in_group) % group_size_m_actual
    pid_m = first_pid_m + pid_m_local_offset
    pid_n = (pid % num_pid_in_group) // group_size_m_actual
    
    # Additional guard if pid_m or pid_n are out of bounds due to ragged grid
    if pid_m >= num_pid_m or pid_n >= num_pid_n:
        return

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

@triton.jit
def unified_fused_moe_kernel(
    x_flat_ptr, w_gate_all_ptr, w_up_all_ptr, w_down_all_ptr,
    final_weighted_output_ptr,
    sorted_assignment_indices_ptr, expert_ids_for_blocks_ptr,
    token_ids_for_flat_assignments_ptr, flat_expert_scores_ptr,
    S_total_assignments_dim: tl.constexpr, num_assignments_post_padded_dim: tl.constexpr,
    N_out_dim: tl.constexpr, K_in_dim: tl.constexpr, K_inter_dim: tl.constexpr, E_num_experts_dim: tl.constexpr,
    stride_x_flat_token_orig, stride_x_flat_k_in,
    stride_w_all_e_g, stride_w_all_k_in_g, stride_w_all_k_inter_g,
    stride_w_all_e_d, stride_w_all_k_inter_d, stride_w_all_n_out_d,
    stride_out_assign_padded, stride_out_n_out,
    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,
):
    pid_flat = tl.program_id(axis=0)
    
    num_pid_m_total = tl.cdiv(num_assignments_post_padded_dim, BLOCK_SIZE_M)
    num_pid_n_total = tl.cdiv(N_out_dim, BLOCK_SIZE_N_OUT)
    
    num_pid_flat_in_m_group = GROUP_SIZE_M_CFG * num_pid_n_total
    
    group_id = pid_flat // num_pid_flat_in_m_group
    pid_m_group_start = group_id * GROUP_SIZE_M_CFG
    
    current_group_actual_num_m_pids = min(num_pid_m_total - pid_m_group_start, GROUP_SIZE_M_CFG)
    
    if current_group_actual_num_m_pids <= 0:
        return

    pid_offset_in_group_flat = pid_flat % num_pid_flat_in_m_group
    
    pid_m_local_offset = pid_offset_in_group_flat % current_group_actual_num_m_pids
    pid_m = pid_m_group_start + pid_m_local_offset
    
    pid_n = pid_offset_in_group_flat // current_group_actual_num_m_pids

    if pid_m >= num_pid_m_total or pid_n >= num_pid_n_total: # Ensure pids are within valid range
        return

    offs_m_block_coords = tl.arange(0, BLOCK_SIZE_M)
    offs_m_global_padded_assignment_rows = pid_m * BLOCK_SIZE_M + offs_m_block_coords
    offs_n_out = pid_n * BLOCK_SIZE_N_OUT + tl.arange(0, BLOCK_SIZE_N_OUT)

    current_expert_id = tl.load(expert_ids_for_blocks_ptr + pid_m) # pid_m is block index for expert_ids

    w_gate_expert_ptr = w_gate_all_ptr + current_expert_id * stride_w_all_e_g
    w_up_expert_ptr = w_up_all_ptr + current_expert_id * stride_w_all_e_g
    w_down_expert_ptr = w_down_all_ptr + current_expert_id * stride_w_all_e_d
    
    mask_rows_in_padded_dim_for_load = offs_m_global_padded_assignment_rows < num_assignments_post_padded_dim
    
    assignment_indices_in_block = tl.load(
        sorted_assignment_indices_ptr + offs_m_global_padded_assignment_rows, 
        mask=mask_rows_in_padded_dim_for_load, 
        other=S_total_assignments_dim 
    )
    
    valid_assignment_mask_in_block = (assignment_indices_in_block < S_total_assignments_dim)

    original_token_indices_in_block = tl.load(
        token_ids_for_flat_assignments_ptr + assignment_indices_in_block,
        mask=valid_assignment_mask_in_block, # Use valid_assignment_mask_in_block
        other=0 
    )

    x_start_ptr_for_block = x_flat_ptr + original_token_indices_in_block[:, None] * stride_x_flat_token_orig

    weight_dtype = w_gate_expert_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_expert_n_block = w_down_expert_ptr + offs_n_out[None, :] * stride_w_all_n_out_d

    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_expert_ptr + offs_k_inter[None, :] * stride_w_all_k_inter_g
        w_up_k_inter_ptr = w_up_expert_ptr + offs_k_inter[None, :] * stride_w_all_k_inter_g

        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_start_ptr_for_block + offs_k_in[None, :] * stride_x_flat_k_in
            x_mask_for_load = valid_assignment_mask_in_block[:, None] & (offs_k_in[None, :] < K_in_dim)
            x_block = tl.load(current_x_ptrs, mask=x_mask_for_load, other=0.0)
            
            w_gate_ptrs = w_gate_k_inter_ptr + offs_k_in[:, None] * stride_w_all_k_in_g
            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_w_all_k_in_g
            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_start_ptr_expert_n_block + offs_k_inter[:, None] * stride_w_all_k_inter_d
        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)

    current_scores_in_block = tl.load(
        flat_expert_scores_ptr + assignment_indices_in_block,
        mask=valid_assignment_mask_in_block, 
        other=0.0
    )
    acc_down = acc_down * current_scores_in_block[:, None]

    final_output_ptrs = final_weighted_output_ptr + offs_m_global_padded_assignment_rows[:, None] * stride_out_assign_padded + offs_n_out[None, :] * stride_out_n_out
    
    output_mask_for_store = valid_assignment_mask_in_block[:, None] & (offs_n_out[None, :] < N_out_dim) & (mask_rows_in_padded_dim_for_load[:,None]) # ensure write within padded buffer bounds
    tl.store(final_output_ptrs, acc_down.to(final_weighted_output_ptr.dtype.element_ty), mask=output_mask_for_store)


class TritonExpert(nn.Module):
    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"]
        self.d_expert_actual = d_expert_actual
        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_m_dim = triton.cdiv(M_tokens, BLOCK_SIZE_M)
        grid_n_dim = triton.cdiv(N_out_dim, BLOCK_SIZE_N_OUT)
        grid = (grid_m_dim * grid_n_dim,)
        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), 
            self.W_up.stride(0), self.W_up.stride(1),     
            self.W_down.stride(0), self.W_down.stride(1), 
            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.register_buffer('W_g', W_g_weight) 

    @torch.compile
    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        logits = F.linear(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_pytorch(
    topk_ids_flat: torch.Tensor,
    block_size_m: int,
    num_experts: int
) -> Tuple[torch.Tensor, torch.Tensor, int]:
    S_total_assignments = topk_ids_flat.numel()
    original_indices = torch.arange(S_total_assignments, device=topk_ids_flat.device, dtype=torch.long)
    perm_for_sorting = topk_ids_flat.argsort()
    sorted_original_indices = original_indices[perm_for_sorting] 
    # Use original topk_ids_flat for bincount as it's more direct, though bincount(sorted_topk_ids_flat) is same.
    expert_counts = torch.bincount(topk_ids_flat, minlength=num_experts) 
    padded_expert_counts = ((expert_counts + block_size_m - 1) // block_size_m) * block_size_m
    num_assignments_post_padded = padded_expert_counts.sum().item()
    # This is the number of blocks for the M dimension
    num_padded_m_blocks = num_assignments_post_padded // block_size_m

    final_sorted_assignment_indices = torch.full(
        (num_assignments_post_padded,), S_total_assignments, dtype=torch.long, device=topk_ids_flat.device
    )
    expert_ids_for_blocks = torch.empty((num_padded_m_blocks,), dtype=torch.long, device=topk_ids_flat.device)

    current_out_idx = 0
    current_block_idx = 0
    processed_in_sorted_original_indices_ptr = 0
    for e in range(num_experts):
        num_actual_for_expert = expert_counts[e].item()
        num_padded_for_expert = padded_expert_counts[e].item()
        
        expert_assignment_indices_for_e = sorted_original_indices[
            processed_in_sorted_original_indices_ptr : processed_in_sorted_original_indices_ptr + num_actual_for_expert
        ]
        processed_in_sorted_original_indices_ptr += num_actual_for_expert
        
        if num_padded_for_expert > 0:
             final_sorted_assignment_indices[current_out_idx : current_out_idx + num_actual_for_expert] = expert_assignment_indices_for_e
        
        num_blocks_for_this_expert = num_padded_for_expert // block_size_m
        if num_blocks_for_this_expert > 0:
            expert_ids_for_blocks[current_block_idx : current_block_idx + num_blocks_for_this_expert] = e
        
        current_out_idx += num_padded_for_expert
        current_block_idx += num_blocks_for_this_expert
        
    return final_sorted_assignment_indices, expert_ids_for_blocks, num_assignments_post_padded


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"]
        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'])
        
        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_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)
        S_tokens_orig = 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)
        S_total_assignments = flat_expert_indices.shape[0]
        token_ids_for_flat_assignments = torch.arange(S_tokens_orig, device=x.device, dtype=torch.long).repeat_interleave(self.top_k)

        K_in_dim, K_inter_dim, N_out_dim = self.d_hidden, self.d_expert, self.d_hidden
        
        # Using a simple default config for now.
        # More sophisticated config selection could be based on num_assignments_post_padded.
        config_key_unified_approx = (S_total_assignments, N_out_dim, K_in_dim, K_inter_dim) 
        selected_params_unified = HARDCODED_FFN_CONFIGS_UNIFIED.get(config_key_unified_approx, DEFAULT_FFN_CONFIG_UNIFIED)
        
        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_unified

        sorted_assignment_indices, expert_ids_for_blocks, num_assignments_post_padded = moe_align_block_size_pytorch(flat_expert_indices, BLOCK_SIZE_M, self.num_experts)

        if num_assignments_post_padded == 0:
            routed_expert_contribution = torch.zeros_like(x_flat)
        else:
            final_weighted_output_padded_buffer = torch.empty(
                num_assignments_post_padded, N_out_dim, device=x.device, dtype=x.dtype
            )
            
            grid_m_dim = triton.cdiv(num_assignments_post_padded, BLOCK_SIZE_M)
            grid_n_dim = triton.cdiv(N_out_dim, BLOCK_SIZE_N_OUT)
            grid = (grid_m_dim * grid_n_dim,)

            unified_fused_moe_kernel[grid](
                x_flat, 
                self.W_gate_all_routed, self.W_up_all_routed, self.W_down_all_routed,
                final_weighted_output_padded_buffer,
                sorted_assignment_indices, expert_ids_for_blocks,
                token_ids_for_flat_assignments, flat_expert_scores,
                S_total_assignments, num_assignments_post_padded,
                N_out_dim, K_in_dim, K_inter_dim, self.num_experts,
                x_flat.stride(0), x_flat.stride(1),
                self.W_gate_all_routed.stride(0), self.W_gate_all_routed.stride(1), self.W_gate_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),
                final_weighted_output_padded_buffer.stride(0), final_weighted_output_padded_buffer.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_CFG=GROUP_SIZE_M_cfg,
                num_warps=num_warps, num_stages=num_stages
            )
            
            # Aggregation:
            # final_weighted_output_padded_buffer is ordered by expert, then by original assignment order within that expert, then padded.
            # sorted_assignment_indices tells us, for each slot in ^, which *original* assignment index it corresponds to.
            # Padded slots in sorted_assignment_indices have value S_total_assignments.
            
            # 1. Create a mask for valid (non-padded) entries in the buffer
            valid_entry_mask_in_padded_buffer = sorted_assignment_indices < S_total_assignments
            
            # 2. Get the computed outputs for these valid entries
            # These are still in the expert-sorted order.
            valid_weighted_outputs = final_weighted_output_padded_buffer[valid_entry_mask_in_padded_buffer]
            
            # 3. Get the *original assignment indices* that these valid_weighted_outputs correspond to.
            # These are also in expert-sorted order.
            original_assignment_indices_of_valid_outputs = sorted_assignment_indices[valid_entry_mask_in_padded_buffer]

            # 4. For each valid output, find its *original token index* to scatter-add to.
            # token_ids_for_flat_assignments maps: original_assignment_idx -> original_token_idx.
            target_token_indices_for_scatter = torch.take(token_ids_for_flat_assignments, original_assignment_indices_of_valid_outputs)
            
            routed_expert_contribution = torch.zeros_like(x_flat)
            scatter_indices = target_token_indices_for_scatter.unsqueeze(-1).expand_as(valid_weighted_outputs)
            routed_expert_contribution.scatter_add_(
                dim=0,
                index=scatter_indices,
                src=valid_weighted_outputs
            )

        shared_expert_output_flat = torch.empty_like(x_flat)
        self.shared_expert_module(x_flat, shared_expert_output_flat)
        final_output_flat = shared_expert_output_flat + routed_expert_contribution
        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()
    triton_moe_model = TritonMoE(config, weights, seq_len_for_config_key)
    input_tensor = input_tensor.to('cuda')
    triton_moe_model.to('cuda') 
    input_dtype = next(iter(weights.values())).dtype
    input_tensor = input_tensor.to(input_dtype)
    output = triton_moe_model(input_tensor)
    return output


python eval.py benchmark /tmp/tmp1ufzf7eq