#!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-05-27 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
import time 
from concurrent.futures import ThreadPoolExecutor

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:
        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
    
    
# Global cache
_model_cache = {}

def get_moe_model_v1(config, weights):
    key = id(weights)  # Or use a hash of the weight dict keys or actual content

    if key not in _model_cache:
        moe = MoE(config)
        moe.gating_network.W_g.weight = nn.Parameter(weights['router.weight'])

        for i in range(config["n_routed_experts"]):
            gate_proj_weight = weights[f'experts.{i}.0.weight']
            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())

        # Load shared expert (single instance, no need to parallelize)
        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())

        _model_cache[key] = moe.eval()

    return _model_cache[key]

def get_moe_model(config, weights):
    
    print(weights['router.weight'].device)

    key = id(weights)
    if key not in _model_cache:
        moe = MoE(config)

        state_dict = {}

        # Router weight
        state_dict['gating_network.W_g.weight'] = weights['router.weight']
        print (state_dict['gating_network.W_g.weight'].device)

        # Experts
        for i in range(config["n_routed_experts"]):
            state_dict[f'experts.{i}.W_gate.weight'] = weights[f'experts.{i}.0.weight'].t()
            state_dict[f'experts.{i}.W_up.weight'] = weights[f'experts.{i}.1.weight'].t()
            state_dict[f'experts.{i}.W_down.weight'] = weights[f'experts.{i}.2.weight'].t()

        # Shared expert
        state_dict['shared_expert.W_gate.weight'] = weights['shared_experts.0.weight'].t()
        state_dict['shared_expert.W_up.weight'] = weights['shared_experts.1.weight'].t()
        state_dict['shared_expert.W_down.weight'] = weights['shared_experts.2.weight'].t()

        moe.load_state_dict(state_dict)
        device = weights['router.weight'].device
        dtype = weights['router.weight'].dtype
        moe.to(device=device, dtype=dtype)
        _model_cache[key] = moe.eval()

    return _model_cache[key]



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
    """
    time_1 = time.time()
    input_tensor, weights, config = data
    # num_experts = config["n_routed_experts"]
    # moe = MoE(config)
    moe = get_moe_model(config, weights)
    time_2 = time.time()

    
    
    # Run the model
    output = moe(input_tensor)
    
    time_6 = time.time()
    
    print (f'Time to load config & setup model: {time_2 - time_1}')
    print (f'Time to inference: {time_6 - time_2}')
    

    return output

python eval.py benchmark /tmp/tmpk45_8835
cuda:0
cuda:0
Time to load config & setup model: 8.168728351593018
Time to inference: 1.117619276046753
cuda:0
Time to load config & setup model: 1.52587890625e-05
Time to inference: 0.0061380863189697266
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.005774497985839844
cuda:0
Time to load config & setup model: 5.9604644775390625e-06
Time to inference: 0.005661487579345703
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.0056591033935546875
cuda:0
Time to load config & setup model: 6.198883056640625e-06
Time to inference: 0.005654335021972656
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.005624532699584961
cuda:0
Time to load config & setup model: 5.9604644775390625e-06
Time to inference: 0.005681514739990234
cuda:0
Time to load config & setup model: 5.4836273193359375e-06
Time to inference: 0.005658626556396484
cuda:0
Time to load config & setup model: 5.9604644775390625e-06
Time to inference: 0.005698204040527344
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.0056514739990234375
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.005677461624145508
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.005684614181518555
cuda:0
Time to load config & setup model: 5.4836273193359375e-06
Time to inference: 0.005633115768432617
cuda:0
Time to load config & setup model: 5.4836273193359375e-06
Time to inference: 0.0056858062744140625
cuda:0
Time to load config & setup model: 1.33514404296875e-05
Time to inference: 0.0056304931640625
cuda:0
Time to load config & setup model: 5.245208740234375e-06
Time to inference: 0.0056836605072021484
cuda:0
Time to load config & setup model: 6.4373016357421875e-06
Time to inference: 0.005674600601196289
cuda:0
cuda:0
Time to load config & setup model: 8.105043649673462
Time to inference: 0.006719350814819336
cuda:0
Time to load config & setup model: 1.5020370483398438e-05
Time to inference: 0.005970478057861328
cuda:0
Time to load config & setup model: 6.67572021484375e-06
Time to inference: 0.005833148956298828
cuda:0
Time to load config & setup model: 6.198883056640625e-06
Time to inference: 0.005707263946533203
cuda:0
Time to load config & setup model: 6.4373016357421875e-06
Time to inference: 0.005684375762939453
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.005744457244873047
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.005673408508300781
cuda:0
Time to load config & setup model: 5.9604644775390625e-06
Time to inference: 0.005679130554199219
cuda:0
Time to load config & setup model: 5.4836273193359375e-06
Time to inference: 0.005697488784790039
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.005727052688598633
cuda:0
Time to load config & setup model: 8.106231689453125e-06
Time to inference: 0.0056569576263427734
cuda:0
Time to load config & setup model: 5.4836273193359375e-06
Time to inference: 0.0056989192962646484
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.005674600601196289
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.0056645870208740234
cuda:0
Time to load config & setup model: 5.4836273193359375e-06
Time to inference: 0.0056858062744140625
cuda:0
Time to load config & setup model: 5.4836273193359375e-06
Time to inference: 0.005722999572753906
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.005713224411010742
cuda:0
Time to load config & setup model: 5.245208740234375e-06
Time to inference: 0.005669116973876953
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.0056514739990234375
cuda:0
Time to load config & setup model: 5.4836273193359375e-06
Time to inference: 0.0056819915771484375
cuda:0
Time to load config & setup model: 1.7404556274414062e-05
Time to inference: 0.005722761154174805
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.005712270736694336
cuda:0
Time to load config & setup model: 5.245208740234375e-06
Time to inference: 0.005669832229614258
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.005753040313720703
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.005649089813232422
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.005678653717041016
cuda:0
Time to load config & setup model: 2.0742416381835938e-05
Time to inference: 0.005653858184814453
cuda:0
Time to load config & setup model: 5.245208740234375e-06
Time to inference: 0.005647420883178711
cuda:0
Time to load config & setup model: 5.245208740234375e-06
Time to inference: 0.00570225715637207
cuda:0
Time to load config & setup model: 5.4836273193359375e-06
Time to inference: 0.005673408508300781
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.005647182464599609
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.005644559860229492
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.005672931671142578
cuda:0
Time to load config & setup model: 5.245208740234375e-06
Time to inference: 0.005700111389160156
cuda:0
Time to load config & setup model: 5.4836273193359375e-06
Time to inference: 0.0056552886962890625
cuda:0
Time to load config & setup model: 5.4836273193359375e-06
Time to inference: 0.00562596321105957
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.005704164505004883
cuda:0
Time to load config & setup model: 6.198883056640625e-06
Time to inference: 0.005687236785888672
cuda:0
Time to load config & setup model: 5.245208740234375e-06
Time to inference: 0.005692958831787109
cuda:0
Time to load config & setup model: 5.9604644775390625e-06
Time to inference: 0.0057141780853271484
cuda:0
Time to load config & setup model: 7.3909759521484375e-06
Time to inference: 0.005711793899536133
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.005686044692993164
cuda:0
Time to load config & setup model: 5.9604644775390625e-06
Time to inference: 0.005732536315917969
cuda:0
Time to load config & setup model: 5.7220458984375e-06
Time to inference: 0.0056536197662353516
cuda:0
Time to load config & setup model: 5.245208740234375e-06
Time to inference: 0.0056879520416259766
cuda:0
cuda:0
Time to load config & setup model: 8.030505895614624
Time to inference: 0.008195638656616211
cuda:0
Time to load config & setup model: 2.0742416381835938e-05
Time to inference: 0.00712132453918457
cuda:0
Time to load config & setup model: 8.106231689453125e-06
Time to inference: 0.006824016571044922
cuda:0
Time to load config & setup model: 8.106231689453125e-06
Time to inference: 0.0068204402923583984
cuda:0
Time to load config & setup model: 8.106231689453125e-06
Time to inference: 0.006751298904418945
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.006749629974365234
cuda:0
Time to load config & setup model: 7.3909759521484375e-06
Time to inference: 0.006742238998413086
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.006772756576538086
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.006757497787475586
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.0068242549896240234
cuda:0
Time to load config & setup model: 7.867813110351562e-06
Time to inference: 0.0067675113677978516
cuda:0
Time to load config & setup model: 7.3909759521484375e-06
Time to inference: 0.006831645965576172
cuda:0
Time to load config & setup model: 1.430511474609375e-05
Time to inference: 0.006784200668334961
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.006785869598388672
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.006734371185302734
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.00675511360168457
cuda:0
Time to load config & setup model: 7.3909759521484375e-06
Time to inference: 0.006817817687988281
cuda:0
Time to load config & setup model: 7.62939453125e-06
Time to inference: 0.00675201416015625
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.006735563278198242
cuda:0
Time to load config & setup model: 7.62939453125e-06
Time to inference: 0.00669097900390625
cuda:0
Time to load config & setup model: 7.62939453125e-06
Time to inference: 0.006804943084716797
cuda:0
Time to load config & setup model: 7.3909759521484375e-06
Time to inference: 0.006751298904418945
cuda:0
Time to load config & setup model: 6.67572021484375e-06
Time to inference: 0.0067861080169677734
cuda:0
Time to load config & setup model: 7.3909759521484375e-06
Time to inference: 0.006758928298950195
cuda:0
Time to load config & setup model: 7.3909759521484375e-06
Time to inference: 0.006726503372192383
cuda:0
Time to load config & setup model: 6.9141387939453125e-06
Time to inference: 0.006702899932861328
cuda:0
Time to load config & setup model: 6.9141387939453125e-06
Time to inference: 0.006766080856323242
cuda:0
Time to load config & setup model: 6.9141387939453125e-06
Time to inference: 0.006830930709838867
cuda:0
Time to load config & setup model: 7.3909759521484375e-06
Time to inference: 0.006730556488037109
cuda:0
Time to load config & setup model: 6.9141387939453125e-06
Time to inference: 0.006674766540527344
cuda:0
Time to load config & setup model: 6.9141387939453125e-06
Time to inference: 0.006738901138305664
cuda:0
Time to load config & setup model: 1.811981201171875e-05
Time to inference: 0.00663447380065918
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.006711721420288086
cuda:0
Time to load config & setup model: 6.9141387939453125e-06
Time to inference: 0.006742000579833984
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.006715059280395508
cuda:0
Time to load config & setup model: 6.9141387939453125e-06
Time to inference: 0.00668787956237793
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.0067138671875
cuda:0
Time to load config & setup model: 2.47955322265625e-05
Time to inference: 0.006750583648681641
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.006703853607177734
cuda:0
Time to load config & setup model: 7.3909759521484375e-06
Time to inference: 0.006702423095703125
cuda:0
Time to load config & setup model: 6.9141387939453125e-06
Time to inference: 0.0066983699798583984
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.006734132766723633
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.006778240203857422
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.0067310333251953125
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.0067212581634521484
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.0067577362060546875
cuda:0
Time to load config & setup model: 6.67572021484375e-06
Time to inference: 0.00676274299621582
cuda:0
Time to load config & setup model: 6.9141387939453125e-06
Time to inference: 0.006743669509887695
cuda:0
Time to load config & setup model: 7.62939453125e-06
Time to inference: 0.006761074066162109
cuda:0
Time to load config & setup model: 2.384185791015625e-05
Time to inference: 0.006732463836669922
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.0067670345306396484
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.00673222541809082
cuda:0
Time to load config & setup model: 6.9141387939453125e-06
Time to inference: 0.0067064762115478516
cuda:0
Time to load config & setup model: 7.3909759521484375e-06
Time to inference: 0.006735324859619141
cuda:0
Time to load config & setup model: 6.9141387939453125e-06
Time to inference: 0.006665229797363281
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.006718635559082031
cuda:0
Time to load config & setup model: 6.4373016357421875e-06
Time to inference: 0.006695270538330078
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.006779193878173828
cuda:0
Time to load config & setup model: 7.62939453125e-06
Time to inference: 0.0067212581634521484
cuda:0
Time to load config & setup model: 7.3909759521484375e-06
Time to inference: 0.006748676300048828
cuda:0
Time to load config & setup model: 6.9141387939453125e-06
Time to inference: 0.0067484378814697266
cuda:0
Time to load config & setup model: 7.152557373046875e-06
Time to inference: 0.006749153137207031
cuda:0
Time to load config & setup model: 6.9141387939453125e-06
Time to inference: 0.006763458251953125
cuda:0
Time to load config & setup model: 7.3909759521484375e-06
Time to inference: 0.006791353225708008
cuda:0
Time to load config & setup model: 6.9141387939453125e-06
Time to inference: 0.006690025329589844
