submission 682356

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

"""MXFP4 MoE with aiter.fused_moe using proper enum types."""

import torch
from task import input_t, output_t
from aiter.fused_moe import fused_moe
from aiter import ActivationType, QuantType


def custom_kernel(data: input_t) -> output_t:
    x = data[0]
    w1_shuffled = data[5]
    w2_shuffled = data[6]
    w1_shuffled_scale = data[7]
    w2_shuffled_scale = data[8]
    routing_weights = data[9]
    routing_indices = data[10]

    out = fused_moe(
        hidden_states=x,
        w1=w1_shuffled,
        w2=w2_shuffled,
        topk_weight=routing_weights,
        topk_ids=routing_indices,
        expert_mask=None,
        activation=ActivationType.Silu,
        quant_type=QuantType.per_1x32,
        doweight_stage1=False,
        w1_scale=w1_shuffled_scale,
        w2_scale=w2_shuffled_scale,
    )
    return out