submission 590226

#!POPCORN leaderboard amd-moe-mxfp4
#!POPCORN gpu MI355X

import torch
import aiter
from aiter import ActivationType, QuantType
from utils import make_match_reference

class MoE_V1_Dispatcher:
    def __init__(self):
        self.hidden_pad = None
        self.intermediate_pad = None

    @torch.compiler.disable
    def __call__(self, data):
        # Despacho limpio y directo sin overhead extra
        if self.hidden_pad is None:
            config = data[11]
            self.hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
            self.intermediate_pad = config["d_expert_pad"] - config["d_expert"]

        return aiter.fused_moe.fused_moe(
            data[0],   
            data[5],   
            data[6],   
            data[9],   
            data[10],  
            expert_mask=None,
            activation=ActivationType.Silu,
            quant_type=QuantType.per_1x32, 
            doweight_stage1=False,  # Mantenemos las heurísticas seguras de AMD
            w1_scale=data[7],  
            w2_scale=data[8],  
            a1_scale=None,
            a2_scale=None,
            hidden_pad=self.hidden_pad,
            intermediate_pad=self.intermediate_pad,
        )

custom_kernel = MoE_V1_Dispatcher()
submission = custom_kernel
check_implementation = make_match_reference(custom_kernel, rtol=2e-2, atol=2e-2)