#!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

    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
    
    group_size_m_local = min(num_pid_m_local_padded - first_pid_m_local, GROUP_SIZE_M_CFG)

    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

    # 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
        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]

        expert_counts = torch.bincount(perm_expert_indices.to(torch.int32), minlength=self.num_experts)
        max_tokens_per_expert = expert_counts.max().item() #max count of instances going to an expert
        S_total_assignments = perm_x_flat.shape[0]

        batched_x_input = torch.zeros(
            self.num_experts,
            max_tokens_per_expert,
            self.d_hidden,
            device=x.device,
            dtype=x.dtype
        )
        expert_starts_in_perm_flat = torch.searchsorted(
            perm_expert_indices,
            torch.arange(self.num_experts, device=x.device),
            right=False
        )

        token_abs_indices_in_perm_flat = torch.arange(S_total_assignments, device=x.device)
        start_offset_for_token_group = expert_starts_in_perm_flat[perm_expert_indices] # Shape [S_total_assignments]
        token_intra_expert_indices = token_abs_indices_in_perm_flat - start_offset_for_token_group # Shape [S_total_assignments]
        batched_x_input[perm_expert_indices, token_intra_expert_indices, :] = perm_x_flat
        batched_output_routed = torch.zeros(
            self.num_experts, max_tokens_per_expert, self.d_hidden,
            device=x.device, dtype=x.dtype
        )

        #Batched kernel attrs
        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)
        
        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,
            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
        )

        token_indices_for_mask = torch.arange(max_tokens_per_expert, device=x.device)
        valid_token_mask = token_indices_for_mask < expert_counts.unsqueeze(-1)
        all_expert_outputs_perm_allocated = batched_output_routed[valid_token_mask]
        weighted_expert_outputs_perm = all_expert_outputs_perm_allocated * perm_expert_scores

        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)

        scatter_indices = perm_token_ids.unsqueeze(-1).expand_as(weighted_expert_outputs_perm)
        shared_expert_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 = shared_expert_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 test /tmp/tmpi6ly2369
/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(
