submission 621513

#!POPCORN leaderboard amd-moe-mxfp4
#!POPCORN gpu MI355X
"""
MoE — clean fused_moe with pre-allocated output.
Focus: let fused_moe handle kernel selection, minimize Python overhead.
"""
import torch
import gc
import sys
import os
from task import input_t, output_t

from aiter.fused_moe import fused_moe
from aiter.ops.enum import ActivationType, QuantType

gc.disable()
sys.setswitchinterval(1000.0)

_ACT = ActivationType.Silu
_QT = QuantType.per_1x32
_out = {}


def custom_kernel(data: input_t) -> output_t:
    hidden = data[0]
    w1 = data[5]
    w2 = data[6]
    w1_scale = data[7]
    w2_scale = data[8]
    topk_w = data[9]
    topk_ids = data[10]

    bs = hidden.shape[0]
    d_hidden = hidden.shape[-1]
    key = (bs, d_hidden)

    if key not in _out:
        _out[key] = torch.empty(bs, d_hidden, dtype=torch.bfloat16, device="cuda")

    return fused_moe(
        hidden, w1, w2, topk_w, topk_ids,
        activation=_ACT, quant_type=_QT,
        w1_scale=w1_scale, w2_scale=w2_scale,
    )