submission 517477

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

from typing import Callable, Dict, Optional, Tuple

from task import input_t, output_t

from aiter import ActivationType, QuantType
from aiter.fused_moe import fused_moe as _fused_moe

_ACTIVATION = ActivationType.Silu
_QUANT_TYPE = QuantType.per_1x32

_PAD_CACHE: Dict[Tuple[int, int, int, int], Tuple[int, int]] = {}
_PAD_CACHE_LIMIT = 8

_LAST_SIGNATURE: Optional[Tuple[int, ...]] = None
_LAST_RUNNER: Optional[Callable[[], output_t]] = None


def _get_pads(config: Dict[str, int]) -> Tuple[int, int]:
    key = (
        int(config["d_hidden"]),
        int(config["d_hidden_pad"]),
        int(config["d_expert"]),
        int(config["d_expert_pad"]),
    )
    pads = _PAD_CACHE.get(key)
    if pads is None:
        pads = (key[1] - key[0], key[3] - key[2])
        if len(_PAD_CACHE) >= _PAD_CACHE_LIMIT:
            _PAD_CACHE.clear()
        _PAD_CACHE[key] = pads
    return pads


def _build_runner(data: input_t) -> Callable[[], output_t]:
    hidden_states = data[0]
    gate_up_weight_shuffled = data[5]
    down_weight_shuffled = data[6]
    gate_up_weight_scale_shuffled = data[7]
    down_weight_scale_shuffled = data[8]
    topk_weights = data[9]
    topk_ids = data[10]

    hidden_pad, intermediate_pad = _get_pads(data[11])

    def _run(
        hidden_states=hidden_states,
        gate_up_weight_shuffled=gate_up_weight_shuffled,
        down_weight_shuffled=down_weight_shuffled,
        gate_up_weight_scale_shuffled=gate_up_weight_scale_shuffled,
        down_weight_scale_shuffled=down_weight_scale_shuffled,
        topk_weights=topk_weights,
        topk_ids=topk_ids,
        hidden_pad=hidden_pad,
        intermediate_pad=intermediate_pad,
    ) -> output_t:
        return _fused_moe(
            hidden_states,
            gate_up_weight_shuffled,
            down_weight_shuffled,
            topk_weights,
            topk_ids,
            expert_mask=None,
            activation=_ACTIVATION,
            quant_type=_QUANT_TYPE,
            doweight_stage1=False,
            w1_scale=gate_up_weight_scale_shuffled,
            w2_scale=down_weight_scale_shuffled,
            a1_scale=None,
            a2_scale=None,
            hidden_pad=hidden_pad,
            intermediate_pad=intermediate_pad,
        )

    return _run


def custom_kernel(data: input_t) -> output_t:
    global _LAST_SIGNATURE, _LAST_RUNNER

    signature = (
        id(data),
        id(data[11]),
        int(data[0].data_ptr()),
        int(data[5].data_ptr()),
        int(data[6].data_ptr()),
        int(data[7].data_ptr()),
        int(data[8].data_ptr()),
        int(data[9].data_ptr()),
        int(data[10].data_ptr()),
    )

    runner = _LAST_RUNNER
    if runner is None or signature != _LAST_SIGNATURE:
        runner = _build_runner(data)
        _LAST_RUNNER = runner
        _LAST_SIGNATURE = signature

    return runner()