submission 564946

#!POPCORN leaderboard amd-moe-mxfp4
#!POPCORN gpu MI355X
from task import input_t, output_t
from aiter import ActivationType, QuantType
from aiter.fused_moe import fused_moe

# Fully unrolled shape dispatch for ranked benchmark shapes.
# Goal: zero dynamic tuning loops and minimal runtime decision overhead.


def _run(
    hidden_states,
    gate_up_weight_shuffled,
    down_weight_shuffled,
    topk_weights,
    topk_ids,
    gate_up_weight_scale_shuffled,
    down_weight_scale_shuffled,
    hidden_pad,
    intermediate_pad,
    *,
    block_m=None,
    non_temporal_load=None,
    kernel_name_1=None,
    kernel_name_2=None,
):
    kwargs = {
        "expert_mask": None,
        "activation": ActivationType.Silu,
        "quant_type": QuantType.per_1x32,
        "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,
    }
    if block_m is not None:
        kwargs["block_m"] = int(block_m)
    if non_temporal_load is not None:
        kwargs["non_temporal_load"] = bool(non_temporal_load)
    if kernel_name_1 is not None:
        kwargs["kernelName1"] = str(kernel_name_1)
    if kernel_name_2 is not None:
        kwargs["kernelName2"] = str(kernel_name_2)

    try:
        return fused_moe(
            hidden_states,
            gate_up_weight_shuffled,
            down_weight_shuffled,
            topk_weights,
            topk_ids,
            **kwargs,
        )
    except TypeError:
        kwargs.pop("block_m", None)
        kwargs.pop("non_temporal_load", None)
        kwargs.pop("kernelName1", None)
        kwargs.pop("kernelName2", None)
        return fused_moe(
            hidden_states,
            gate_up_weight_shuffled,
            down_weight_shuffled,
            topk_weights,
            topk_ids,
            **kwargs,
        )


def custom_kernel(data: input_t) -> output_t:
    (
        hidden_states,
        _gate_up_weight,
        _down_weight,
        _gate_up_weight_scale,
        _down_weight_scale,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        gate_up_weight_scale_shuffled,
        down_weight_scale_shuffled,
        topk_weights,
        topk_ids,
        config,
    ) = data

    bs = int(config.get("bs", hidden_states.shape[0]))
    d_expert = int(config.get("d_expert", -1))
    n_routed = int(config.get("n_routed_experts", -1))
    n_experts_per_token = int(config.get("n_experts_per_token", topk_ids.shape[-1] - 1))
    hidden_pad = int(config["d_hidden_pad"]) - int(config["d_hidden"])
    intermediate_pad = int(config["d_expert_pad"]) - int(config["d_expert"])

    # Ranked set shape specializations (unrolled).
    if bs == 16 and d_expert == 256 and n_routed == 256 and n_experts_per_token == 8:
        return _run(
            hidden_states,
            gate_up_weight_shuffled,
            down_weight_shuffled,
            topk_weights,
            topk_ids,
            gate_up_weight_scale_shuffled,
            down_weight_scale_shuffled,
            hidden_pad,
            intermediate_pad,
            block_m=32,
            non_temporal_load=True,
            kernel_name_1="moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16",
            kernel_name_2="moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16",
        )

    if bs == 128 and d_expert == 256 and n_routed == 256 and n_experts_per_token == 8:
        return _run(
            hidden_states,
            gate_up_weight_shuffled,
            down_weight_shuffled,
            topk_weights,
            topk_ids,
            gate_up_weight_scale_shuffled,
            down_weight_scale_shuffled,
            hidden_pad,
            intermediate_pad,
            block_m=64,
            non_temporal_load=True,
            kernel_name_1="moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16",
            kernel_name_2="moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16",
        )

    if bs == 512 and d_expert == 256 and n_routed == 256 and n_experts_per_token == 8:
        return _run(
            hidden_states,
            gate_up_weight_shuffled,
            down_weight_shuffled,
            topk_weights,
            topk_ids,
            gate_up_weight_scale_shuffled,
            down_weight_scale_shuffled,
            hidden_pad,
            intermediate_pad,
            block_m=64,
            non_temporal_load=True,
            kernel_name_1="moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16",
            kernel_name_2="moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16",
        )

    if bs == 16 and d_expert == 512 and n_routed == 32 and n_experts_per_token == 8:
        return _run(
            hidden_states,
            gate_up_weight_shuffled,
            down_weight_shuffled,
            topk_weights,
            topk_ids,
            gate_up_weight_scale_shuffled,
            down_weight_scale_shuffled,
            hidden_pad,
            intermediate_pad,
            block_m=32,
            non_temporal_load=True,
        )

    if bs == 128 and d_expert == 512 and n_routed == 32 and n_experts_per_token == 8:
        return _run(
            hidden_states,
            gate_up_weight_shuffled,
            down_weight_shuffled,
            topk_weights,
            topk_ids,
            gate_up_weight_scale_shuffled,
            down_weight_scale_shuffled,
            hidden_pad,
            intermediate_pad,
            block_m=64,
            non_temporal_load=True,
        )

    if bs == 512 and d_expert == 512 and n_routed == 32 and n_experts_per_token == 8:
        return _run(
            hidden_states,
            gate_up_weight_shuffled,
            down_weight_shuffled,
            topk_weights,
            topk_ids,
            gate_up_weight_scale_shuffled,
            down_weight_scale_shuffled,
            hidden_pad,
            intermediate_pad,
            block_m=64,
            non_temporal_load=False,
        )

    if bs == 512 and d_expert == 2048 and n_routed == 32 and n_experts_per_token == 8:
        return _run(
            hidden_states,
            gate_up_weight_shuffled,
            down_weight_shuffled,
            topk_weights,
            topk_ids,
            gate_up_weight_scale_shuffled,
            down_weight_scale_shuffled,
            hidden_pad,
            intermediate_pad,
            block_m=128,
            non_temporal_load=False,
        )

    # Fallback path for unseen shapes.
    return _run(
        hidden_states,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        topk_weights,
        topk_ids,
        gate_up_weight_scale_shuffled,
        down_weight_scale_shuffled,
        hidden_pad,
        intermediate_pad,
    )