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

import triton
import triton.language as tl

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

# Global dictionary to store printed FFN config summaries to avoid re-printing
_printed_ffn_config_summaries = {}

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

# Autotuner configs - can be shared or specialized
# For fused_ffn_kernel_v2 (shared expert)
AUTOTUNE_CONFIGS_SINGLE_FFN = [
    triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 512, 'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=8, num_stages=3),
    # 2. Max Pipelining for the champion
    triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 512, 'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=8, num_stages=4),
    # 3. Fewer Warps with champion's stages
    triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 512, 'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=4, num_stages=2),
    # 4. Different Grouping (larger) with champion's M, N, K
    triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 512, 'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 8}, num_warps=8, num_stages=2),

    # --- Push N_OUT even further, or vary M with N_OUT=512 ---
    # 5. N_OUT=512, Larger M
    triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N_OUT': 512, 'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=8, num_stages=2),
    # 6. N_OUT=512, Smaller M (if useful for routed experts with few tokens)
    triton.Config({'BLOCK_SIZE_M': 32,  'BLOCK_SIZE_N_OUT': 512, 'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 2}, num_warps=8, num_stages=2), # Smaller G for smaller M
    # 7. N_OUT even larger (e.g., 1024, if supported and N_out problem dim is huge) - may hit limits
    # triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 1024, 'BLOCK_SIZE_K_IN': 64, 'BLOCK_SIZE_K_INTER': 64, 'GROUP_SIZE_M': 2}, num_warps=8, num_stages=1), # Potentially 1 stage due to resources

    # --- Explore K block sizes with N_OUT=512 and M=64 ---
    # 8. N_OUT=512, M=64, Larger K blocks
    triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 512, 'BLOCK_SIZE_K_IN': 128, 'BLOCK_SIZE_K_INTER': 128, 'GROUP_SIZE_M': 4}, num_warps=8, num_stages=2),
    # 9. N_OUT=512, M=64, Smaller K blocks
    triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 512, 'BLOCK_SIZE_K_IN': 32,  'BLOCK_SIZE_K_INTER': 32,  'GROUP_SIZE_M': 4}, num_warps=8, num_stages=3), # More stages for smaller K

    # --- Try combinations with M=128 (previous good M for some NVIDIA) and N_OUT=256 (previous good N) but vary K wildly ---
    # 10. M=128, N=256, Very Large K
    triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N_OUT': 256, 'BLOCK_SIZE_K_IN': 128, 'BLOCK_SIZE_K_INTER': 128, 'GROUP_SIZE_M': 8}, num_warps=8, num_stages=1), # Fewer stages for large K
    # 11. M=128, N=256, Very Small K
    triton.Config({'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N_OUT': 1024, 'BLOCK_SIZE_K_IN': 32,  'BLOCK_SIZE_K_INTER': 32,  'GROUP_SIZE_M': 8}, num_warps=8, num_stages=2), # More stages for small K

    # --- Revisit M=256 (original good M for large shared FFNs) with new N_OUT insights ---
    # 12. M=256, N_OUT=256 (balanced)
    triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N_OUT': 256, 'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=8, num_stages=2),
    # 13. M=256, N_OUT=512 (winner's N_OUT)
    triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N_OUT': 512, 'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 2}, num_warps=8, num_stages=1), # G and S reduced due to large M, N

    # --- Focus on different warp/stage/grouping for the champion's M,N,K ---
    # (64, 512, 64, 64, G4, W8, S2)
    # 14. Minimal Grouping
    triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 512, 'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 1}, num_warps=8, num_stages=2),
    # 15. Fewer Warps, Max Stages
    triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 512, 'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=4, num_stages=4),
    # 16. Max Warps, Min Stages (potentially good if highly compute-bound and not memory latency bound)
    triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 512, 'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=8, num_stages=1),

    # --- Two "outliers" to check boundaries ---
    # 17. Small M, Small N, but very large K and aggressive grouping/stages (if K is dominant)
    triton.Config({'BLOCK_SIZE_M': 32,  'BLOCK_SIZE_N_OUT': 32,  'BLOCK_SIZE_K_IN': 128, 'BLOCK_SIZE_K_INTER': 128, 'GROUP_SIZE_M': 2}, num_warps=4, num_stages=4),
    # 18. Large M, Large N, but small K and minimal stages (if M/N are dominant and K is small)
    triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N_OUT': 256, 'BLOCK_SIZE_K_IN': 32,  'BLOCK_SIZE_K_INTER': 32,  'GROUP_SIZE_M': 8}, num_warps=8, num_stages=1),

    # --- Explore N_OUT slightly smaller than 512, with other params from winner ---
    # 19. N_OUT = 256 (previous winning N), M=64, K=64, G=4, W8, S2
    triton.Config({'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N_OUT': 256, 'BLOCK_SIZE_K_IN': 64,  'BLOCK_SIZE_K_INTER': 64,  'GROUP_SIZE_M': 4}, num_warps=8, num_stages=2),
]

# For batched_fused_ffn_kernel_expert_dim (routed experts)
# Key will be ['max_tokens_per_expert_padded_dim', 'N_out_dim', 'K_in_dim', 'K_inter_dim']
# We can use similar configs, but GROUP_SIZE_M is named GROUP_SIZE_M_CFG in batched kernel configs
AUTOTUNE_CONFIGS_BATCHED_FFN = [
    triton.Config({'BLOCK_SIZE_M': cfg['BLOCK_SIZE_M'], 
                   'BLOCK_SIZE_N_OUT': cfg['BLOCK_SIZE_N_OUT'], 
                   'BLOCK_SIZE_K_IN': cfg['BLOCK_SIZE_K_IN'], 
                   'BLOCK_SIZE_K_INTER': cfg['BLOCK_SIZE_K_INTER'], 
                   'GROUP_SIZE_M_CFG': cfg['GROUP_SIZE_M']}, # Note the key change here
                  num_warps=config.num_warps, num_stages=config.num_stages)
    for config in AUTOTUNE_CONFIGS_SINGLE_FFN for cfg in [config.kwargs] # Iterate over original configs
]


@triton.autotune(
    configs=AUTOTUNE_CONFIGS_SINGLE_FFN,
    key=['M', 'N_out', 'K_in', 'K_inter']
)
@triton.jit
def fused_ffn_kernel_v2_autotuned( # Renamed for clarity
    # Pointers to Matrices
    x_ptr, w_gate_ptr, w_up_ptr, w_down_ptr, output_ptr,
    # Matrix dimensions for autotune key
    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 from autotuner
    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, # This comes from autotune config
):
    pid = tl.program_id(axis=0)
    num_n_out_blocks_total = tl.cdiv(N_out, BLOCK_SIZE_N_OUT)
    pids_per_m_group_stripe = GROUP_SIZE_M * num_n_out_blocks_total
    m_group_stripe_id = pid // pids_per_m_group_stripe
    first_m_block_in_stripe = m_group_stripe_id * GROUP_SIZE_M
    pid_within_stripe = pid % pids_per_m_group_stripe
    m_offset_within_group = pid_within_stripe % GROUP_SIZE_M 
    current_pid_n_out = pid_within_stripe // GROUP_SIZE_M
    current_pid_m = first_m_block_in_stripe + m_offset_within_group
    
    offs_m_block_local = tl.arange(0, BLOCK_SIZE_M)
    offs_m_global = current_pid_m * BLOCK_SIZE_M + offs_m_block_local
    
    offs_n_out = current_pid_n_out * 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.autotune(
    configs=AUTOTUNE_CONFIGS_BATCHED_FFN,
    key=['max_tokens_per_expert_padded_dim', 'N_out_dim', 'K_in_dim', 'K_inter_dim']
)
@triton.jit
def batched_fused_ffn_kernel_expert_dim_autotuned( # Renamed for clarity
    # Pointers to Batched/Stacked Matrices
    batched_x_ptr, all_w_gate_ptr, all_w_up_ptr, all_w_down_ptr, batched_output_ptr,
    expert_counts_ptr,
    # Matrix dimensions for autotune key
    max_tokens_per_expert_padded_dim, N_out_dim, K_in_dim, K_inter_dim,
    # Strides
    stride_b_x_e, stride_b_x_m, stride_b_x_k,
    stride_all_w_gate_e, stride_all_w_gate_k, stride_all_w_gate_inter,
    stride_all_w_up_e, stride_all_w_up_k, stride_all_w_up_inter,
    stride_all_w_down_e, stride_all_w_down_inter, stride_all_w_down_n,
    stride_b_out_e, stride_b_out_m, stride_b_out_n,
    # Meta-parameters from autotuner
    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, # This comes from autotune config
):
    expert_id = tl.program_id(axis=0)
    M_curr_expert_actual = tl.load(expert_counts_ptr + expert_id)
    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
    
    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

    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)
    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)
    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)
    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 TritonExpertAutotuned(nn.Module): # For Shared Expert
    def __init__(self, W_gate: torch.Tensor, W_up: torch.Tensor, W_down: torch.Tensor, d_expert_actual: int):
        super().__init__()
        self.d_expert_actual = d_expert_actual
        self.W_gate = W_gate
        self.W_up = W_up
        self.W_down = W_down

    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 = x.shape[1] # Output dimension matches input for standard FFN in MoE

        x = x.contiguous()
        
        grid = lambda META: (
            triton.cdiv(M_tokens, META['BLOCK_SIZE_M']) * triton.cdiv(N_out_dim, META['BLOCK_SIZE_N_OUT']),
        )
        
        fused_ffn_kernel_v2_autotuned[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, # Args for autotune key
            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_*, GROUP_SIZE_M are passed by autotuner
        )

class MoEGate(nn.Module): # Same as before
    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

class TritonMoEAutotuned(nn.Module):
    def __init__(self, config: Dict, weights: Dict[str, torch.Tensor]):
        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.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 = TritonExpertAutotuned(
            w_gate_shared, w_up_shared, w_down_shared, shared_expert_inter_dim
        )
        
        # Store dimensions for routed experts for autotuner key construction in print logic
        self.routed_K_in_dim = self.d_hidden
        self.routed_K_inter_dim = self.d_expert
        self.routed_N_out_dim = self.d_hidden
        
        # Store dimensions for shared expert for autotuner key construction in print logic
        self.shared_K_in_dim = self.d_hidden
        self.shared_K_inter_dim = shared_expert_inter_dim
        self.shared_N_out_dim = self.d_hidden

        # Attributes to store last call's dimensions for printing
        self.last_max_tokens_per_expert_routed = 0 
        self.last_M_tokens_shared = 0


    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]

        # --- Routed Experts Computation (Batched Kernel) ---
        expert_counts = torch.bincount(perm_expert_indices.to(torch.int32), minlength=self.num_experts)
        max_tokens_per_expert = expert_counts.max().item()
        self.last_max_tokens_per_expert_routed = max_tokens_per_expert # Save for printing

        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)

        grid_dim0_experts = self.num_experts
        # The META arguments for grid are from the autotuner
        grid = lambda META: (grid_dim0_experts, 
                                triton.cdiv(max_tokens_per_expert, META['BLOCK_SIZE_M']) * triton.cdiv(self.routed_N_out_dim, META['BLOCK_SIZE_N_OUT']))
        
        expert_counts_int32 = expert_counts.to(torch.int32)

        batched_fused_ffn_kernel_expert_dim_autotuned[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,
            # Args for autotune key:
            max_tokens_per_expert, self.routed_N_out_dim, self.routed_K_in_dim, self.routed_K_inter_dim,
            # Strides:
            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_*, GROUP_SIZE_M_CFG are passed by autotuner
        )

        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()
            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
            
        # --- Shared Expert Computation ---
        shared_expert_output_flat = torch.empty_like(x_flat)
        self.shared_expert_module(x_flat, shared_expert_output_flat)
        self.last_M_tokens_shared = x_flat.shape[0] # Save for printing

        # --- Combine Routed and Shared Expert Outputs ---
        final_output_flat = shared_expert_output_flat
        if perm_x_flat.shape[0] > 0: # Ensure there are outputs from routed experts
            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)
        
        output = final_output_flat.view(orig_shape)
        return output

def custom_kernel(data: input_t) -> output_t:
    input_tensor, weights, config = data
    
    for k, v_orig in weights.items():
        weights[k] = v_orig.contiguous()

    triton_moe_model = TritonMoEAutotuned(config, weights)
    
    input_tensor_cuda = input_tensor.to('cuda')
    triton_moe_model.to('cuda')
    
    input_dtype = next(iter(weights.values())).dtype
    input_tensor_cuda = input_tensor_cuda.to(input_dtype)
    
    output = triton_moe_model(input_tensor_cuda)

    # Print autotune configs
    global _printed_ffn_config_summaries
    
    # Shared Expert Config Print
    shared_kernel = fused_ffn_kernel_v2_autotuned
    actual_shared_key_for_autotune = ( 
        triton_moe_model.last_M_tokens_shared,
        triton_moe_model.shared_N_out_dim,
        triton_moe_model.shared_K_in_dim,
        triton_moe_model.shared_K_inter_dim
    )
    print_suppression_key_shared = ('shared',) + actual_shared_key_for_autotune + (config['seq_len'],)

    if print_suppression_key_shared not in _printed_ffn_config_summaries:
        best_config_to_print_shared = None
        key_of_best_config_printed = "N/A" # The key for which the printed config was best

        if hasattr(shared_kernel, 'cache') and actual_shared_key_for_autotune in shared_kernel.cache:
            cache_entry = shared_kernel.cache[actual_shared_key_for_autotune]
            if cache_entry: # This is a triton.Config object
                best_config_to_print_shared = cache_entry
                key_of_best_config_printed = actual_shared_key_for_autotune
        
        if not best_config_to_print_shared and hasattr(shared_kernel, 'best_config') and shared_kernel.best_config:
            best_config_to_print_shared = shared_kernel.best_config # This is the kernel's global best
            # Now, find which key this global best_config corresponds to
            if hasattr(shared_kernel, 'cache'):
                found_key_for_global_best = False
                for k_iter, entry_iter in shared_kernel.cache.items():
                    # entry_iter is AutotuneCacheEntry, entry_iter.best_config is triton.Config
                    if entry_iter is best_config_to_print_shared: # Object identity check
                        key_of_best_config_printed = f"OverallBest (for key {k_iter})"
                        found_key_for_global_best = True
                        break
                if not found_key_for_global_best:
                    key_of_best_config_printed = "OverallBest (original key not found in cache)"
            else: # Should not happen if best_config exists
                key_of_best_config_printed = "OverallBest (cache not found)"


        if best_config_to_print_shared:
            try:
                cfg_kwargs = best_config_to_print_shared.kwargs
                summary = (f"S.{cfg_kwargs.get('BLOCK_SIZE_M')}-{cfg_kwargs.get('BLOCK_SIZE_N_OUT')}-"
                           f"{cfg_kwargs.get('BLOCK_SIZE_K_IN')}-{cfg_kwargs.get('BLOCK_SIZE_K_INTER')}-"
                           f"G{cfg_kwargs.get('GROUP_SIZE_M')}-W{best_config_to_print_shared.num_warps}-S{best_config_to_print_shared.num_stages}")
                print(f"SharedExpert Autotune (CurrentActualKey={actual_shared_key_for_autotune}, PrintedFor={key_of_best_config_printed}, SL={config['seq_len']}): {summary}", file=sys.stderr)
            except Exception as e:
                print(f"Error summarizing SharedExpert FFN config for {print_suppression_key_shared}: {e}", file=sys.stderr)
        _printed_ffn_config_summaries[print_suppression_key_shared] = True


    # Routed Experts Batched Config Print
    routed_batched_kernel = batched_fused_ffn_kernel_expert_dim_autotuned
    actual_routed_key_for_autotune = (
        triton_moe_model.last_max_tokens_per_expert_routed,
        triton_moe_model.routed_N_out_dim,
        triton_moe_model.routed_K_in_dim,
        triton_moe_model.routed_K_inter_dim
    )
    print_suppression_key_routed = ('routed_batched',) + actual_routed_key_for_autotune + (config['seq_len'],)
    
    if triton_moe_model.last_max_tokens_per_expert_routed > 0: # Only print if routed path was taken
        if print_suppression_key_routed not in _printed_ffn_config_summaries:
            best_config_to_print_routed = None
            key_of_best_config_printed_routed = "N/A"

            if hasattr(routed_batched_kernel, 'cache') and actual_routed_key_for_autotune in routed_batched_kernel.cache:
                cache_entry_routed = routed_batched_kernel.cache[actual_routed_key_for_autotune]
                if cache_entry_routed: # This is a triton.Config
                    best_config_to_print_routed = cache_entry_routed
                    key_of_best_config_printed_routed = actual_routed_key_for_autotune

            if not best_config_to_print_routed and hasattr(routed_batched_kernel, 'best_config') and routed_batched_kernel.best_config:
                best_config_to_print_routed = routed_batched_kernel.best_config # Kernel's global best
                if hasattr(routed_batched_kernel, 'cache'):
                    found_key_for_global_best_routed = False
                    for k_iter,  entry_iter in routed_batched_kernel.cache.items():
                        if entry_iter is best_config_to_print_routed:
                            key_of_best_config_printed_routed = f"OverallBest (for key {k_iter})"
                            found_key_for_global_best_routed = True
                            break
                    if not found_key_for_global_best_routed:
                        key_of_best_config_printed_routed = "OverallBest (original key not found in cache)"
                else:
                     key_of_best_config_printed_routed = "OverallBest (cache not found)"


            if best_config_to_print_routed:
                try:
                    cfg_kwargs = best_config_to_print_routed.kwargs
                    summary = (f"RB.{cfg_kwargs.get('BLOCK_SIZE_M')}-{cfg_kwargs.get('BLOCK_SIZE_N_OUT')}-"
                               f"{cfg_kwargs.get('BLOCK_SIZE_K_IN')}-{cfg_kwargs.get('BLOCK_SIZE_K_INTER')}-"
                               f"G{cfg_kwargs.get('GROUP_SIZE_M_CFG')}-W{best_config_to_print_routed.num_warps}-S{best_config_to_print_routed.num_stages}")
                    print(f"RoutedBatched Autotune (CurrentActualKey={actual_routed_key_for_autotune}, PrintedFor={key_of_best_config_printed_routed}, SL={config['seq_len']}): {summary}", file=sys.stderr)
                except Exception as e:
                    print(f"Error summarizing RoutedBatched FFN config for {print_suppression_key_routed}: {e}", file=sys.stderr)
            _printed_ffn_config_summaries[print_suppression_key_routed] = True
                
    return output

python eval.py benchmark /tmp/tmp4mgnzcbh
SharedExpert Autotune (CurrentActualKey=(2048, 7168, 7168, 2048), PrintedFor=OverallBest (for key (2048, 7168, 7168, 2048, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.float16', 'torch.float16')), SL=2048): S.128-512-64-64-G4-W8-S2
RoutedBatched Autotune (CurrentActualKey=(292, 7168, 7168, 2048), PrintedFor=OverallBest (for key (292, 7168, 7168, 2048, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.float16', 'torch.float16', 'torch.int32')), SL=2048): RB.64-512-64-64-G4-W4-S2
SharedExpert Autotune (CurrentActualKey=(8192, 7168, 7168, 2048), PrintedFor=OverallBest (for key (2048, 7168, 7168, 2048, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.float16', 'torch.float16')), SL=8192): S.128-512-64-64-G4-W8-S2
RoutedBatched Autotune (CurrentActualKey=(1104, 7168, 7168, 2048), PrintedFor=OverallBest (for key (1104, 7168, 7168, 2048, 'torch.float16', 'torch.float16', 'torch.float16', 'torch.float16', 'torch.float16', 'torch.int32')), SL=8192): RB.128-512-64-64-G4-W8-S2
