submission 682216

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

import torch

from task import input_t, output_t

from aiter import ActivationType, QuantType
from aiter.fused_moe import fused_moe
from aiter.fused_moe_bf16_asm import ck_moe_2stages


_TARGET_BATCH_SIZE = 1024
_TARGET_D_EXPERT = 2048
_TARGET_N_ROUTED = 32
_TARGET_TOTAL_TOPK = 9


def _as_contiguous(x):
    if x.is_contiguous():
        return x
    return x.contiguous()


def _should_probe_stage_control(hidden_states, config) -> bool:
    return (
        int(hidden_states.shape[0]) == _TARGET_BATCH_SIZE
        and int(config.get("d_expert", -1)) == _TARGET_D_EXPERT
        and int(config.get("n_routed_experts", -1)) == _TARGET_N_ROUTED
        and int(config.get("total_top_k", -1)) == _TARGET_TOTAL_TOPK
    )


def _run_stage_control(
    hidden_states,
    gate_up_weight_shuffled,
    down_weight_shuffled,
    gate_up_weight_scale_shuffled,
    down_weight_scale_shuffled,
    topk_weights,
    topk_ids,
):
    return ck_moe_2stages(
        hidden_states,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        topk_weights,
        topk_ids,
        quant_type=QuantType.per_1x32,
        fc1_scale=gate_up_weight_scale_shuffled,
        fc2_scale=down_weight_scale_shuffled,
        activation=ActivationType.Silu,
        doweight_stage1=False,
    )


def _run_fused_baseline(
    hidden_states,
    gate_up_weight_shuffled,
    down_weight_shuffled,
    gate_up_weight_scale_shuffled,
    down_weight_scale_shuffled,
    topk_weights,
    topk_ids,
    config,
):
    hidden_pad = int(config["d_hidden_pad"]) - int(config["d_hidden"])
    intermediate_pad = int(config["d_expert_pad"]) - int(config["d_expert"])
    return fused_moe(
        hidden_states,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        topk_weights,
        topk_ids,
        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,
    )


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

    del gate_up_weight
    del down_weight
    del gate_up_weight_scale
    del down_weight_scale

    hidden_states = _as_contiguous(hidden_states)
    topk_weights = _as_contiguous(topk_weights)
    topk_ids = _as_contiguous(topk_ids)

    # Only probe the one large EP-on case that the task docs explicitly single
    # out as the most plausible split-K / large-M opportunity.
    if _should_probe_stage_control(hidden_states, config):
        try:
            output = _run_stage_control(
                hidden_states,
                gate_up_weight_shuffled,
                down_weight_shuffled,
                gate_up_weight_scale_shuffled,
                down_weight_scale_shuffled,
                topk_weights,
                topk_ids,
            )
            if torch.cuda.is_available():
                torch.cuda.synchronize()
            return output
        except Exception:
            pass

    return _run_fused_baseline(
        hidden_states,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        gate_up_weight_scale_shuffled,
        down_weight_scale_shuffled,
        topk_weights,
        topk_ids,
        config,
    )