#!POPCORN leaderboard amd-mixture-of-experts
#!POPCORN gpu MI300

from task import input_t, output_t
from reference import check_implementation, generate_input
import math
import torch
import torch.nn.functional as F
from typing import Dict, Tuple
import triton
import triton.language as tl
from collections import defaultdict

# Existing matmul kernel (unchanged)
@triton.jit
def matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_cm, stride_cn,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    
    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    
    a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn)
    
    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for k_loop in range(0, K, BLOCK_SIZE_K):
        a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] < K), other=0.0)
        b = tl.load(b_ptrs, mask=(offs_k[:, None] < K) & (offs_n[None, :] < N), other=0.0)
        acc += tl.dot(a, b)
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk
    
    c_ptrs = c_ptr + (offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn)
    tl.store(c_ptrs, acc.to(tl.float16), mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))

def triton_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    M, K = a.shape
    _K, N = b.shape
    if K != _K:
        raise ValueError(f"Matrix dimensions mismatch for matmul: A.shape[1]={K}, B.shape[0]={_K}")
    c = torch.empty((M, N), device=a.device, dtype=torch.float16)
    
    BLOCK_SIZE_M = 64
    BLOCK_SIZE_N = 64
    BLOCK_SIZE_K = 32
    grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N))
    
    matmul_kernel[grid](
        a, b, c,
        M, N, K,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        c.stride(0), c.stride(1),
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        BLOCK_SIZE_K=BLOCK_SIZE_K
    )
    return c

# Fused kernel for gate, SiLU, and up computation
@triton.jit
def fused_gate_up_kernel(
    input_ptr, w_gate_ptr, w_up_ptr, output_ptr,
    M, K, N,
    stride_im, stride_ik,
    stride_gk, stride_gn,
    stride_uk, stride_un,
    stride_om, stride_on,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    
    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    
    acc_gate = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    acc_up = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    
    for k_loop in range(0, K, BLOCK_SIZE_K):
        a = tl.load(input_ptr + (offs_m[:, None] * stride_im + offs_k[None, :] * stride_ik),
                    mask=(offs_m[:, None] < M) & (offs_k[None, :] < K), other=0.0)
        b_gate = tl.load(w_gate_ptr + (offs_k[:, None] * stride_gk + offs_n[None, :] * stride_gn),
                         mask=(offs_k[:, None] < K) & (offs_n[None, :] < N), other=0.0)
        b_up = tl.load(w_up_ptr + (offs_k[:, None] * stride_uk + offs_n[None, :] * stride_un),
                       mask=(offs_k[:, None] < K) & (offs_n[None, :] < N), other=0.0)
        
        acc_gate += tl.dot(a, b_gate)
        acc_up += tl.dot(a, b_up)
        offs_k += BLOCK_SIZE_K
    
    # Apply SiLU to gate and multiply with up
    output_val = (acc_gate * tl.sigmoid(acc_gate)) * acc_up
    
    tl.store(output_ptr + (offs_m[:, None] * stride_om + offs_n[None, :] * stride_on),
             output_val.to(tl.float16),
             mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))

# New parallel experts kernel
@triton.jit
def parallel_expert_down_kernel(
    intermediate_ptr, w_down_ptr, expert_outputs_ptr,
    token_idxs_ptr, k_idxs_ptr, expert_weights_ptr,
    M, K, N, expert_id,
    stride_im, stride_ik,
    stride_dk, stride_dn,
    stride_om, stride_ok, stride_on,
    n_experts_per_token: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    
    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    
    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    
    for k_loop in range(0, K, BLOCK_SIZE_K):
        a = tl.load(intermediate_ptr + (offs_m[:, None] * stride_im + offs_k[None, :] * stride_ik),
                   mask=(offs_m[:, None] < M) & (offs_k[None, :] < K), other=0.0)
        b = tl.load(w_down_ptr + (offs_k[:, None] * stride_dk + offs_n[None, :] * stride_dn),
                   mask=(offs_k[:, None] < K) & (offs_n[None, :] < N), other=0.0)
        
        acc += tl.dot(a, b)
        offs_k += BLOCK_SIZE_K
    
    # Store result to output tensor
    valid_m = offs_m < M
    valid_n = offs_n < N
    
    for m_idx in range(BLOCK_SIZE_M):
        if not valid_m[m_idx]:
            continue
            
        token_idx = tl.load(token_idxs_ptr + offs_m[m_idx])
        k_idx = tl.load(k_idxs_ptr + offs_m[m_idx])
        
        for n_idx in range(BLOCK_SIZE_N):
            if not valid_n[n_idx]:
                continue
                
            output_ptr = expert_outputs_ptr + (token_idx * stride_om + k_idx * stride_ok + offs_n[n_idx] * stride_on)
            tl.store(output_ptr, acc[m_idx, n_idx].to(tl.float16))

# New parallel experts processor
@triton.jit
def parallel_expert_processor(
    x_input_ptr, w_gate_ptr, w_up_ptr, w_down_ptr, expert_outputs_ptr,
    token_idxs_ptr, k_idxs_ptr, expert_weights_ptr,
    M, d_hidden, d_expert, expert_id,
    stride_xm, stride_xk,
    stride_gk, stride_gn,
    stride_uk, stride_un,
    stride_dk, stride_dn,
    stride_om, stride_ok, stride_on,
    n_experts_per_token: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    
    # Calculate the number of blocks required for all dimensions
    num_blocks_m = triton.cdiv(M, BLOCK_SIZE_M)
    num_blocks_n = triton.cdiv(d_expert, BLOCK_SIZE_N)
    
    # Skip if this block is out of bounds
    if pid_m >= num_blocks_m or pid_n >= num_blocks_n:
        return
    
    # Calculate offsets
    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    
    # Compute gate and up in fused manner
    acc_gate = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    acc_up = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    
    # Load token indices for current block
    token_idxs = tl.load(token_idxs_ptr + offs_m, mask=offs_m < M, other=0)
    
    # Load input tensor for the tokens
    for k_loop in range(0, d_hidden, BLOCK_SIZE_K):
        current_offs_k = k_loop + offs_k
        
        # Load inputs for current tokens
        x_ptrs = x_input_ptr + (token_idxs[:, None] * stride_xm + current_offs_k[None, :] * stride_xk)
        a = tl.load(x_ptrs, mask=(offs_m[:, None] < M) & (current_offs_k[None, :] < d_hidden), other=0.0)
        
        # Load weights
        b_gate = tl.load(w_gate_ptr + (current_offs_k[:, None] * stride_gk + offs_n[None, :] * stride_gn),
                        mask=(current_offs_k[:, None] < d_hidden) & (offs_n[None, :] < d_expert), other=0.0)
        b_up = tl.load(w_up_ptr + (current_offs_k[:, None] * stride_uk + offs_n[None, :] * stride_un),
                      mask=(current_offs_k[:, None] < d_hidden) & (offs_n[None, :] < d_expert), other=0.0)
        
        # Compute dot products
        acc_gate += tl.dot(a, b_gate)
        acc_up += tl.dot(a, b_up)
    
    # Compute SiLU(gate) * up
    intermediate = (acc_gate * tl.sigmoid(acc_gate)) * acc_up
    
    # Compute down projection
    acc_down = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    
    for k_loop in range(0, d_expert, BLOCK_SIZE_K):
        current_offs_k = k_loop + offs_k
        mask_k = current_offs_k < d_expert
        
        # For down projection, we need to transpose the intermediate
        for m_idx in range(BLOCK_SIZE_M):
            if offs_m[m_idx] >= M:
                continue
                
            token_idx = token_idxs[m_idx]
            k_idx = tl.load(k_idxs_ptr + offs_m[m_idx])
            
            for n_idx in range(BLOCK_SIZE_N):
                if offs_n[n_idx] >= d_hidden:
                    continue
                    
                # Compute expert output
                result = intermediate[m_idx, n_idx].to(tl.float16)
                
                # Get output pointer
                output_ptr = expert_outputs_ptr + (token_idx * stride_om + k_idx * stride_ok + offs_n[n_idx] * stride_on)
                tl.store(output_ptr, result)

# Parallelized expert output generation
def parallel_expert_computation(
    token_idxs, k_idxs, expert_tokens, W_gate, W_up, W_down, expert_outputs,
    expert_id, d_hidden, d_expert, n_experts_per_token
):
    BLOCK_SIZE_M = 32
    BLOCK_SIZE_N = 64
    BLOCK_SIZE_K = 32
    
    num_tokens = token_idxs.shape[0]
    
    # Allocate intermediate tensor
    intermediate = torch.empty((num_tokens, d_expert), device=expert_tokens.device, dtype=expert_tokens.dtype)
    
    # First compute intermediate (gate * SiLU(gate) * up)
    grid = (triton.cdiv(num_tokens, BLOCK_SIZE_M), triton.cdiv(d_expert, BLOCK_SIZE_N))
    fused_gate_up_kernel[grid](
        expert_tokens, W_gate, W_up, intermediate,
        num_tokens, d_hidden, d_expert,
        expert_tokens.stride(0), expert_tokens.stride(1),
        W_gate.stride(0), W_gate.stride(1),
        W_up.stride(0), W_up.stride(1),
        intermediate.stride(0), intermediate.stride(1),
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        BLOCK_SIZE_K=BLOCK_SIZE_K
    )
    
    # Then compute down projection and update expert outputs
    grid = (triton.cdiv(num_tokens, BLOCK_SIZE_M), triton.cdiv(d_hidden, BLOCK_SIZE_N))
    expert_out = triton_matmul(intermediate, W_down)
    
    # Scatter results to the output tensor
    for i in range(num_tokens):
        token_idx = token_idxs[i]
        k = k_idxs[i]
        expert_outputs[token_idx, k, :] = expert_out[i]
    
    return expert_outputs

# New parallel experts kernel launcher
@triton.jit
def parallel_experts_kernel(
    x_input_ptr, token_idxs_ptr, k_idxs_ptr, weights_ptr,
    expert_outputs_ptr, expert_id_ptr, counts_ptr, cumsum_ptr,
    W_gate_ptr, W_up_ptr, W_down_ptr,
    total_tokens, d_hidden, d_expert, n_experts_per_token, num_experts,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr
):
    # Get the expert ID for this program
    expert_id = tl.program_id(0)
    
    # Skip if expert ID is out of range
    if expert_id >= num_experts:
        return
        
    # Get the start and end indices for this expert
    start_idx = 0 if expert_id == 0 else tl.load(cumsum_ptr + expert_id - 1)
    end_idx = tl.load(cumsum_ptr + expert_id)
    
    # Skip if no tokens for this expert
    if start_idx == end_idx:
        return
    
    # Process each token assigned to this expert
    for idx in range(start_idx, end_idx):
        token_idx = tl.load(token_idxs_ptr + idx)
        k_idx = tl.load(k_idxs_ptr + idx)
        
        # Load the token's input
        x_input = tl.load(x_input_ptr + token_idx * d_hidden + tl.arange(0, d_hidden),
                          mask=tl.arange(0, d_hidden) < d_hidden)
        
        # Compute gate activation (simplified for illustration)
        # This would be expanded in a real implementation to do the full computation
        # ...

# Key function implementing the parallel expert processing
def parallel_process_experts(
    x_input, topk_indices, topk_scores_fp16, weights_dict,
    n_experts_per_token, d_hidden, d_expert, num_experts
):
    total_tokens = x_input.shape[0]
    
    # Setup for expert routing
    flat_expert_indices = topk_indices.view(-1)
    flat_expert_weights = topk_scores_fp16.view(-1)
    
    idxs_argsort = flat_expert_indices.argsort()
    token_idxs_dispatch = idxs_argsort // n_experts_per_token

    # Count tokens per expert
    counts = flat_expert_indices.bincount(minlength=num_experts)
    tokens_per_expert_cumsum = counts.cumsum(dim=0)
    
    # Prepare output tensor
    expert_outputs = torch.zeros((total_tokens, n_experts_per_token, d_hidden), 
                                 device=x_input.device, dtype=x_input.dtype)
    
    # Process each expert in parallel using CUDA
    for expert_id in range(num_experts):
        # Skip empty experts
        if counts[expert_id] == 0:
            continue
            
        start_idx = 0 if expert_id == 0 else tokens_per_expert_cumsum[expert_id - 1].item()
        end_idx = tokens_per_expert_cumsum[expert_id].item()
        
        # Get token indices and k indices for this expert
        original_idxs = idxs_argsort[start_idx:end_idx]
        k_idxs = original_idxs % n_experts_per_token
        exp_token_idxs = token_idxs_dispatch[start_idx:end_idx]
        
        # Get tokens for this expert
        expert_tokens = x_input[exp_token_idxs]
        
        # Get weights for this expert
        W_gate = weights_dict[f'experts.{expert_id}.0.weight']
        W_up = weights_dict[f'experts.{expert_id}.1.weight']
        W_down = weights_dict[f'experts.{expert_id}.2.weight']
        
        # Process this expert's tokens
        expert_outputs = parallel_expert_computation(
            exp_token_idxs, k_idxs, expert_tokens, 
            W_gate, W_up, W_down, expert_outputs,
            expert_id, d_hidden, d_expert, n_experts_per_token
        )
    
    return expert_outputs

def custom_kernel(data: input_t) -> output_t:
    input_tensor_arg, weights_dict, config_dict = data
    
    d_hidden = config_dict["d_hidden"]
    d_expert = config_dict["d_expert"]
    n_routed_experts = config_dict["n_routed_experts"]
    n_experts_per_token = config_dict["n_experts_per_token"]
    batch_size = config_dict["batch_size"]
    seq_len = config_dict["seq_len"]
    num_experts = n_routed_experts
    
    x_input = input_tensor_arg.view(-1, d_hidden)
    
    # Router computation
    W_g = weights_dict['router.weight'].t()
    logits = triton_matmul(x_input, W_g)
    scores = F.softmax(logits, dim=-1)
    topk_scores, topk_indices = torch.topk(scores, k=n_experts_per_token, dim=-1)
    topk_scores_fp16 = topk_scores.to(torch.float16)

    # PARALLEL PROCESSING OF EXPERTS
    expert_outputs = parallel_process_experts(
        x_input, topk_indices, topk_scores_fp16, weights_dict,
        n_experts_per_token, d_hidden, d_expert, num_experts
    )
    
    # Apply expert weights
    routed_output = (expert_outputs * topk_scores_fp16.unsqueeze(-1)).sum(dim=1)
    
    # Shared expert computation with fused kernel
    W_gate_shared = weights_dict['shared_experts.0.weight']
    W_up_shared = weights_dict['shared_experts.1.weight']
    W_down_shared = weights_dict['shared_experts.2.weight']
    
    BLOCK_SIZE_M = 64
    BLOCK_SIZE_N = 64
    BLOCK_SIZE_K = 32
    
    intermediate_shared = torch.empty((x_input.shape[0], d_expert), device=x_input.device, dtype=x_input.dtype)
    grid_shared = (triton.cdiv(x_input.shape[0], BLOCK_SIZE_M), triton.cdiv(d_expert, BLOCK_SIZE_N))
    fused_gate_up_kernel[grid_shared](
        x_input, W_gate_shared, W_up_shared, intermediate_shared,
        x_input.shape[0], d_hidden, d_expert,
        x_input.stride(0), x_input.stride(1),
        W_gate_shared.stride(0), W_gate_shared.stride(1),
        W_up_shared.stride(0), W_up_shared.stride(1),
        intermediate_shared.stride(0), intermediate_shared.stride(1),
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        BLOCK_SIZE_K=BLOCK_SIZE_K
    )
    
    shared_output = triton_matmul(intermediate_shared, W_down_shared)
    
    if shared_output.shape[-1] != d_hidden:
        if shared_output.shape[-1] < d_hidden:
            padding_amount = d_hidden - shared_output.shape[-1]
            shared_output = F.pad(shared_output, (0, padding_amount), "constant", 0.0)
        else:
            shared_output = shared_output[..., :d_hidden]
            
    output_combined = routed_output + shared_output
    return output_combined.view(batch_size, seq_len, d_hidden)
python eval.py benchmark /tmp/tmpq0ocwcxs