#!POPCORN leaderboard amd-mixture-of-experts

# This is a submission template for popcorn leaderboard 'amd-mixture-of-experts'.
# Your task is as follows:
# > For a more complete description, see: https://tinyurl.com/amd-comp-moe
# > Implement a DeepSeek-style Mixture of Experts (MoE) layer for efficient transformer models
# > on a single MI300X device.
# > 
# > MoE is a technique that allows scaling model capacity without proportionally increasing computational costs
# > by using a routing mechanism to selectively activate only a subset of parameters for each token.
# > 
# > Your task:
# > - Implement token routing using a simple softmax-based learned router
# > - Route tokens to the top-k experts based on router probabilities
# > - Process tokens through their assigned experts
# > - Combine expert outputs weighted by router probabilities
# > - Calculate appropriate auxiliary losses for training stability
# > 
# > Input:
# > - `data`: Tuple of (input: torch.Tensor, weights: Dict[str, torch.Tensor], config: Dict)
# >   - input: Input tensor of shape [bs, seq_len, d_hidden]
# >   - weights: Dictionary containing model weights
# >   - config: Dictionary containing model configuration parameters
# > 
# > Output:
# > - Tuple containing:
# >   - output: Processed tensor [bs, seq_len, d_model]
# >   - aux_data: Dictionary with auxiliary data like router probabilities and losses
# The deadline for this leaderboard is 2025-06-02 00:00:00+00:00

# You can automatically route this file to specific GPUs by adding a line
# `#!POPCORN gpus <GPUs>` to the header of this file.
# Happy hacking!

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple, Optional
from task import input_t, output_t

class Expert(nn.Module):
    def __init__(self, config: Dict, d_expert: Optional[int] = None):
        super().__init__()
        self.config = config
        self.act_fn = nn.SiLU()
        self.d_hidden: int = config["d_hidden"]
        self.d_expert: int = config["d_expert"] if d_expert is None else d_expert

        self.W_gate = nn.Linear(self.d_hidden, self.d_expert, bias=False)
        self.W_up = nn.Linear(self.d_hidden, self.d_expert, bias=False)
        self.W_down = nn.Linear(self.d_expert, self.d_hidden, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if(x.dim() == 2):
            print(self.W_gate(x))
        gate = self.act_fn(self.W_gate(x))
        out = self.W_down(gate * self.W_up(x))
        return out


class MoEGate(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.top_k: int = config["n_experts_per_token"]
        self.num_experts: int = config["n_routed_experts"]
        self.d_hidden: int = config["d_hidden"]

        self.W_g = nn.Linear(self.d_hidden, self.num_experts, bias=False)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        logits = self.W_g(x)
        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 MoE(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.config = config
        self.experts = nn.ModuleList([
            Expert(config)
            for _ in range(config["n_routed_experts"])
        ])
        self.gating_network = MoEGate(config)
        shared_expert_dim = config["d_expert"] * config["n_shared_experts"]
        self.shared_expert = Expert(config=config, d_expert=shared_expert_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        shared_output = self.shared_expert(x)
        expert_indices, expert_scores = self.gating_network(x)
        batch_size, seq_len, hidden_dim = x.shape
        orig_shape = x.shape
        x_flat = x.view(-1, hidden_dim)
        flat_expert_indices = expert_indices.view(-1)
        flat_expert_weights = expert_scores.view(-1, 1)
        routed_output_flat = self.moe_infer(x_flat,
                                            flat_expert_indices,
                                            flat_expert_weights)

        routed_output = routed_output_flat.view(*orig_shape)
        return routed_output + shared_output

    @torch.no_grad()
    def moe_infer(self,
                  x: torch.Tensor,
                  flat_expert_indices: torch.Tensor,
                  flat_expert_weights: torch.Tensor
                 ) -> torch.Tensor:
        expert_cache = torch.zeros_like(x)
        idxs = flat_expert_indices.argsort()
        counts = flat_expert_indices.bincount().cpu().numpy()
        tokens_per_expert = counts.cumsum()
        num_per_tok = self.config["n_experts_per_token"]
        token_idxs = idxs // num_per_tok
        for expert_id, end_idx in enumerate(tokens_per_expert):
            start_idx = 0 if expert_id == 0 else tokens_per_expert[expert_id - 1]
            if start_idx == end_idx:
                continue

            expert = self.experts[expert_id]
            exp_token_idxs = token_idxs[start_idx:end_idx]
            expert_tokens = x[exp_token_idxs]
            expert_out    = expert(expert_tokens)
            expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
            expert_cache.scatter_reduce_(
                0,
                exp_token_idxs.view(-1, 1).repeat(1, x.shape[-1]),
                expert_out,
                reduce='sum'
            )

        return expert_cache


def custom_kernel(data: input_t) -> output_t:
    """
    Submission template for DeepSeek-style Mixture of Experts using PyTorch.
    
    Args:
        data: Tuple of (input: torch.Tensor, weights: Dict[str, torch.Tensor], config: Dict)
            - input: Input tensor of shape [batch_size, seq_len, hidden_size]
            - weights: Dictionary containing model weights
            - config: Dictionary containing model configuration parameters
            
    Returns:
        Tuple containing:
            - output: Processed tensor [batch_size, seq_len, d_model]
            - aux_data: Dictionary with auxiliary data
    """
    input_tensor, weights, config = data
    num_experts = config["n_routed_experts"]
    moe = MoE(config)

    # Fill in the given weights of the model
    moe.gating_network.W_g.weight = nn.Parameter(weights['router.weight'])

    for i in range(num_experts):
        # gate_proj_weight = torch.ones(7168,2048)
        gate_proj_weight = weights[f'experts.{i}.0.weight']
        # print(gate_proj_weight.shape)
        up_proj_weight = weights[f'experts.{i}.1.weight']
        down_proj_weight = weights[f'experts.{i}.2.weight']

        # Transpose weights to match expected shape for nn.Linear
        moe.experts[i].W_gate.weight = nn.Parameter(gate_proj_weight.t())
        moe.experts[i].W_up.weight = nn.Parameter(up_proj_weight.t())
        moe.experts[i].W_down.weight = nn.Parameter(down_proj_weight.t())

    moe.shared_expert.W_gate.weight = nn.Parameter(weights['shared_experts.0.weight'].t())
    moe.shared_expert.W_up.weight = nn.Parameter(weights['shared_experts.1.weight'].t())
    moe.shared_expert.W_down.weight = nn.Parameter(weights['shared_experts.2.weight'].t())

    # Run the model
    output = moe(input_tensor)
    # print(input_tensor.shape)
    # print(output[0][0][0])

    return output

python eval.py test /tmp/tmpz72p56wu
tensor([[-1.3555,  2.1816,  2.0566,  ...,  0.9336, -0.2452, -0.5420],
        [-3.4961,  4.0586, -1.2188,  ...,  0.4331,  1.0459,  2.9727],
        [-1.5039, -0.4509, -0.2791,  ..., -2.1855,  0.7173,  3.0273],
        ...,
        [ 0.6777, -3.0410, -2.5352,  ...,  1.7354, -0.8569,  2.2695],
        [ 2.6582, -1.2471,  0.2390,  ..., -0.9307,  0.7876, -0.3513],
        [ 1.8438,  2.7715,  2.5508,  ...,  2.4590,  5.2422, -0.2998]],
       device='cuda:0', dtype=torch.float16)
tensor([[ 1.4385e+00,  1.9229e+00,  1.1797e+00,  ..., -2.9160e-02,
         -2.1680e-01,  1.6074e+00],
        [ 4.1758e+00, -4.2944e-01,  9.3848e-01,  ..., -1.6931e-01,
          8.1104e-01, -2.2930e+00],
        [ 2.3008e+00, -1.3848e+00,  3.8818e-01,  ...,  3.3154e-01,
         -8.0615e-01, -1.8643e+00],
        ...,
        [ 1.2930e+00,  1.3047e+00, -2.1016e+00,  ..., -2.1426e+00,
          1.3623e+00,  5.9619e-01],
        [ 1.4697e+00,  1.5312e+00,  2.6582e+00,  ...,  5.2832e-01,
          1.8887e+00,  3.7539e+00],
        [ 7.1478e-04, -4.1289e+00, -3.5010e-01,  ...,  1.3916e+00,
         -1.6846e+00,  1.4316e+00]], device='cuda:0', dtype=torch.float16)
tensor([[-0.9268, -2.0352,  1.1553,  ...,  1.0029, -0.0257,  1.9346],
        [ 2.3145,  0.1302, -1.7568,  ..., -0.8359,  1.0361, -1.0273],
        [ 1.3965,  1.1523, -0.4250,  ..., -1.4404, -0.2217, -4.3789],
        ...,
        [ 0.5767, -1.3145, -0.6118,  ...,  3.1484, -0.3679,  2.1543],
        [ 0.4626,  0.0789,  0.2676,  ...,  4.3164,  2.8555, -0.2527],
        [ 2.6367,  1.2080,  3.9551,  ...,  1.7627, -0.2419,  0.0715]],
       device='cuda:0', dtype=torch.float16)
tensor([[ 0.2549, -1.0732,  2.9551,  ...,  1.2939, -3.8770,  0.6582],
        [-3.6895, -1.9531, -0.3374,  ..., -0.3228, -1.8555,  1.2676],
        [-2.2559,  0.8281,  0.3892,  ..., -0.5073, -0.3318,  1.3262],
        ...,
        [ 0.1582, -0.6514,  0.1170,  ...,  0.3655,  0.4600, -0.2092],
        [-1.5000,  1.5293, -0.2913,  ...,  1.1104,  1.7061,  1.2197],
        [ 3.8184, -3.5449, -1.6348,  ...,  1.1699, -0.8872, -0.6460]],
       device='cuda:0', dtype=torch.float16)
tensor([[ 0.3464,  1.7588,  1.7656,  ...,  3.0996, -2.5742, -0.9639],
        [ 3.5527, -4.2695, -2.1289,  ..., -0.5791,  0.0272, -2.7676],
        [-0.4602, -1.9785, -1.5713,  ..., -0.1407, -1.6104, -4.0547],
        ...,
        [-1.3164, -1.6865, -1.0586,  ...,  0.3975, -0.4971,  5.6758],
        [-2.7578,  0.7661, -2.3867,  ...,  0.5649,  1.6299,  1.5469],
        [-0.2585, -1.2500,  0.2583,  ...,  3.1016,  0.6772,  1.7764]],
       device='cuda:0', dtype=torch.float16)
tensor([[ 0.8486,  2.4961, -2.4277,  ...,  1.4551, -2.9395, -0.5698],
        [ 0.9121,  1.5166, -2.4199,  ...,  0.3015, -2.4199,  0.5615],
        [ 2.8828, -2.4355, -0.7593,  ...,  2.1895,  0.7666,  0.2881],
        ...,
        [ 0.5283, -3.1367,  1.2520,  ..., -2.3008,  2.1055,  0.5713],
        [-2.5371, -0.1788, -1.1240,  ...,  0.2878,  3.1934,  2.1191],
        [-0.9277, -2.2871, -2.1562,  ..., -0.0636,  2.3672, -1.5967]],
       device='cuda:0', dtype=torch.float16)
tensor([[ 3.0625, -3.9980,  0.1414,  ..., -2.1895,  2.8105,  4.1133],
        [ 2.3926, -0.3623, -1.0098,  ...,  0.0597,  0.0082,  2.9219],
        [ 1.0400, -1.4277, -3.0020,  ..., -0.6284,  2.3711, -2.4629],
        ...,
        [ 3.6191,  1.8770, -0.4500,  ..., -0.6650, -0.2365, -0.2803],
        [ 2.7969, -1.0205, -0.4143,  ..., -2.9434,  1.2148, -1.3193],
        [-1.0752, -0.2325, -0.6348,  ...,  2.6582,  1.7725, -0.3218]],
       device='cuda:0', dtype=torch.float16)
tensor([[ 0.0528,  2.1211,  0.4653,  ...,  0.1975,  0.1296, -1.3750],
        [ 2.9004, -0.0748,  1.0078,  ..., -0.3962,  0.7842,  0.7031],
        [-1.9062,  2.1523,  1.5674,  ...,  0.1338, -0.8491, -3.4590],
        ...,
        [ 1.9307,  0.1201,  1.2725,  ..., -1.5225,  2.7227,  1.8672],
        [-1.6709,  0.3286, -1.2246,  ...,  0.9902, -0.4277, -0.4539],
        [ 1.0771, -2.1660,  5.8359,  ..., -3.4473,  0.1487,  0.2913]],
       device='cuda:0', dtype=torch.float16)
tensor([[ 1.9092,  1.1201, -1.1543,  ...,  0.7827,  2.1309, -1.2197],
        [-0.3748,  2.8633,  1.5537,  ..., -1.0020,  1.4463, -0.8042],
        [ 1.5459,  1.8184,  0.6514,  ..., -2.2656, -0.9175,  3.2031],
        ...,
        [-3.3223,  0.9458,  0.7568,  ...,  0.8643, -4.3984, -2.8672],
        [ 4.7500, -4.2031,  0.4043,  ...,  0.0140,  0.6152,  1.1436],
        [-1.4492,  3.5625, -0.4695,  ...,  0.9023, -2.2422, -1.4639]],
       device='cuda:0', dtype=torch.float16)
tensor([[-0.4890, -0.3955,  3.9707,  ..., -5.6289,  2.0449,  2.4922],
        [ 1.2480, -0.1854,  2.4082,  ...,  0.5947, -1.0059,  0.4868],
        [ 1.2109, -1.9229, -0.9941,  ...,  2.1348, -0.7944,  0.2367],
        ...,
        [ 2.4238, -4.4609, -1.6455,  ..., -0.6567, -0.5464,  2.1484],
        [ 0.1262,  1.3438, -1.4092,  ...,  0.0971,  1.5566, -0.8042],
        [ 0.7275, -1.6992,  1.2744,  ..., -0.5039, -1.6406, -0.3098]],
       device='cuda:0', dtype=torch.float16)
tensor([[-1.0684,  0.0239,  0.2925,  ...,  2.3887, -1.3398,  1.9795],
        [ 0.0658,  3.1895, -0.5449,  ..., -0.0253, -1.0664, -3.8320],
        [ 0.0925,  2.4922,  1.6416,  ...,  0.1154, -0.3491,  0.7568],
        ...,
        [ 2.3359, -1.5078, -0.4900,  ...,  1.1807,  2.9883, -0.2148],
        [-3.6191,  1.9062, -1.4336,  ..., -1.6494,  1.5068, -2.5020],
        [-0.0821, -1.1621,  0.1124,  ..., -2.9727, -0.9209, -0.0101]],
       device='cuda:0', dtype=torch.float16)
tensor([[ 0.9727, -1.4668,  1.1816,  ..., -0.2632, -1.0000, -1.6104],
        [-2.0762, -0.7144, -2.9785,  ...,  0.0900, -0.2622,  1.6582],
        [ 1.1602,  2.0137,  0.0622,  ..., -1.2422,  1.7373,  1.1357],
        ...,
        [ 0.1941, -1.6543, -1.7754,  ...,  0.5083, -0.8833,  0.9663],
        [ 1.7441, -0.2500,  0.6821,  ..., -3.7070,  2.6973, -0.1223],
        [-2.4199,  2.4336,  0.1764,  ..., -0.2944, -0.6968, -1.1807]],
       device='cuda:0', dtype=torch.float16)
tensor([[ 1.1260, -2.1250,  2.9707,  ..., -1.0078, -0.9014, -1.1406],
        [-2.0977, -0.2461,  3.4805,  ...,  0.0392, -1.3643,  3.5781],
        [ 1.2910, -0.4133,  0.8857,  ..., -1.6982, -0.7036,  0.0693],
        ...,
        [-2.7363, -1.7598,  2.6094,  ..., -0.8735,  0.6060, -0.8960],
        [ 2.5625, -1.6133, -0.6748,  ...,  1.3281,  1.6943,  3.0273],
        [-4.3867,  3.6113,  2.6855,  ...,  1.1562, -1.8770, -1.8125]],
       device='cuda:0', dtype=torch.float16)
tensor([[-3.7832, -4.8086,  0.1234,  ..., -2.3242,  0.7847, -2.0391],
        [ 0.4949,  0.1344, -1.8945,  ..., -1.3809,  0.0397,  4.1797],
        [-0.5684,  1.0771, -3.2812,  ..., -0.7827, -0.5146, -0.0246],
        ...,
        [ 0.2423,  1.4395, -1.7549,  ...,  2.9258,  0.2104,  3.7617],
        [-3.8516, -3.1484,  3.0449,  ...,  0.9971,  0.0839,  2.3418],
        [ 0.1176,  0.2227, -1.4424,  ..., -1.7607, -0.0286,  1.2559]],
       device='cuda:0', dtype=torch.float16)
tensor([[-2.4805,  1.3535, -1.6270,  ..., -3.3047, -2.6445, -0.5337],
        [ 3.2031, -0.1096,  0.1691,  ...,  0.1588, -1.0039, -4.0742],
        [ 3.1152, -0.3745,  1.3164,  ...,  1.3506, -1.4883, -0.6807],
        ...,
        [ 1.0098,  0.2385,  1.6455,  ..., -3.3887, -0.1048, -0.5532],
        [-2.9590,  0.8374,  1.1270,  ...,  2.6113, -0.2549,  0.4155],
        [ 0.5034,  0.0634,  2.4336,  ...,  0.8525, -1.9102, -1.9961]],
       device='cuda:0', dtype=torch.float16)
tensor([[-2.4492,  0.6025, -2.6895,  ...,  1.6299,  2.2109, -1.7109],
        [ 3.9844, -2.6543,  4.0117,  ..., -3.3926,  1.3545, -0.5986],
        [ 1.0615, -1.2734,  1.7061,  ..., -1.6621,  1.8760, -1.1611],
        ...,
        [-0.2566,  0.0063,  1.5371,  ...,  2.1973,  1.7656, -0.5439],
        [-1.4980,  2.1406,  0.7705,  ..., -1.3770, -1.7266, -0.8457],
        [ 0.3616, -2.4668, -3.5059,  ...,  1.1416, -1.9307,  0.4458]],
       device='cuda:0', dtype=torch.float16)
tensor([[ 2.1973, -0.1058, -2.1348,  ...,  2.8906,  1.4639,  1.1621],
        [-0.1761,  2.3809, -0.6025,  ..., -1.7910, -0.5273, -3.7949],
        [-0.6553, -2.1953, -1.3438,  ..., -0.3262,  3.3457, -2.0566],
        ...,
        [ 0.4778, -3.6523, -2.3262,  ..., -0.7393, -3.5957, -0.1877],
        [ 1.6182, -0.7070,  3.1133,  ..., -0.6836, -0.7207, -2.0449],
        [-1.1289, -0.1814,  0.7061,  ..., -0.8271,  1.1152, -1.2588]],
       device='cuda:0', dtype=torch.float16)
tensor([[-0.4373, -3.3887,  0.1464,  ..., -1.8750,  0.3057, -1.4795],
        [-0.1954,  0.3887, -0.6265,  ..., -3.9766,  0.5679, -1.8145],
        [ 1.0918,  0.0614,  1.7129,  ...,  0.3562, -2.6191, -1.6523],
        ...,
        [ 2.5078,  2.6445, -1.3779,  ...,  0.1368,  0.1155, -0.1113],
        [ 1.2412,  0.6978, -1.0400,  ..., -0.8916, -1.4082, -1.1865],
        [-0.6538, -1.5156, -0.9590,  ...,  0.3420, -0.2524, -0.2441]],
       device='cuda:0', dtype=torch.float16)
tensor([[-3.2349e-01,  7.3145e-01,  7.5488e-01,  ..., -2.3359e+00,
          4.1577e-01, -1.8564e+00],
        [-1.6611e+00, -1.4688e+00,  8.5986e-01,  ..., -2.3184e+00,
         -1.2852e+00, -2.0654e-01],
        [ 5.7031e-01, -4.7485e-01,  4.1724e-01,  ..., -4.5391e+00,
         -1.5732e+00, -3.4062e+00],
        ...,
        [ 1.3906e+00,  1.8408e+00,  2.2969e+00,  ...,  4.6992e+00,
         -4.1479e-01,  6.6211e-01],
        [ 1.8428e+00,  3.8306e-01,  9.3506e-01,  ..., -1.1260e+00,
         -1.5537e+00, -1.2715e+00],
        [-8.8574e-01,  2.0625e+00, -1.8921e-01,  ...,  2.3984e+00,
         -9.4509e-04,  5.1514e-01]], device='cuda:0', dtype=torch.float16)
tensor([[ 2.0176, -1.1221,  1.4424,  ..., -0.2159, -3.7129,  4.6562],
        [ 1.8115, -4.7773, -2.7891,  ...,  0.5708,  0.5430,  1.4180],
        [ 0.8613,  2.0957,  0.9365,  ..., -2.1387, -0.1556, -0.7012],
        ...,
        [-2.4473,  1.3428, -5.9062,  ...,  0.3093,  0.7285, -3.0801],
        [ 0.2391, -2.5449, -0.1316,  ..., -0.9600, -0.5161, -2.1680],
        [ 1.5059, -4.0195, -1.2764,  ..., -3.4004,  0.7949, -2.9023]],
       device='cuda:0', dtype=torch.float16)
