submission 722023

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

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

# Keep strict MXFP4 path.
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False

# Prefer opus sorting path for faster dispatch.
os.environ.setdefault("AITER_USE_OPUS_MOE_SORTING", "1")

_PAD_CACHE = {}

def _get_pads(config):
    d_hidden = config["d_hidden"]
    d_hidden_pad = config["d_hidden_pad"]
    d_expert = config["d_expert"]
    d_expert_pad = config["d_expert_pad"]
    
    key = (d_hidden, d_hidden_pad, d_expert, d_expert_pad)
    pads = _PAD_CACHE.get(key)
    if pads is None:
        pads = (
            int(d_hidden_pad - d_hidden),
            int(d_expert_pad - d_expert),
        )
        _PAD_CACHE[key] = pads
    return pads

def custom_kernel(data: input_t) -> output_t:
    hidden_states = data[0]
    w1 = data[5]
    w2 = data[6]
    s1 = data[7]
    s2 = data[8]
    topk_weights = data[9]
    topk_ids = data[10]
    config = data[11]

    hidden_pad, intermediate_pad = _get_pads(config)

    # Calling fused_moe without forcing `block_size_M` or `moe_sorting_dispatch_policy`
    # This avoids the Memory Access Fault, falling back to AITER's internal defaults 
    # and safely handling dynamically tuned parameters from aiter's model_configs csvs.
    return fused_moe(
        hidden_states,
        w1,
        w2,
        topk_weights,
        topk_ids,
        expert_mask=None,
        activation=ActivationType.Silu,
        quant_type=QuantType.per_1x32,
        doweight_stage1=False,
        w1_scale=s1,
        w2_scale=s2,
        a1_scale=None,
        a2_scale=None,
        hidden_pad=hidden_pad,
        intermediate_pad=intermediate_pad,
    )