submission 592373

#!POPCORN leaderboard amd-moe-mxfp4
#!POPCORN gpu MI355X
"""MoE MXFP4 baseline — use aiter fused_moe directly."""
from task import input_t, output_t
import torch
import aiter
from aiter import dtypes
from aiter.fused_moe import fused_moe, ActivationType, QuantType

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

    M = hidden_states.shape[0]
    d_hidden = config['d_hidden']
    d_expert = config['d_expert']
    n_routed = config['n_routed_experts']
    n_shared = config['n_shared_experts']
    total_top_k = config['total_top_k']

    # Use aiter's fused_moe with MXFP4 quantization
    # gate_up_weight_shuffled is [E, 2*d_expert_pad, d_hidden_pad//2] fp4x2 (pre-shuffled)
    # down_weight_shuffled is [E, d_hidden_pad, d_expert_pad//2] fp4x2 (pre-shuffled)
    result = fused_moe(
        hidden_states=hidden_states,
        w1=gate_up_weight_shuffled,
        w2=down_weight_shuffled,
        topk_weight=topk_weights,
        topk_ids=topk_ids,
        w1_scale=gate_up_weight_scale_shuffled,
        w2_scale=down_weight_scale_shuffled,
        activation=ActivationType.Silu,
        quant_type=QuantType.per_1x32,
        dtype=torch.bfloat16,
    )
    return result