#!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 test /tmp/tmp_j5fuw7m
/usr/local/lib/python3.10/dist-packages/torch/_inductor/lowering.py:7078: UserWarning: 
Online softmax is disabled on the fly since Inductor decides to
split the reduction. Cut an issue to PyTorch if this is an
important use case and you want to speed it up with online
softmax.

  warnings.warn(
