submission 575278

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

import torch
from typing import Dict
from task import input_t, output_t

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

def custom_kernel(data: input_t) -> output_t:
    """
    Submission template for DeepSeek-R1 MXFP4 MoE kernel.
    Optimized for AMD MI355X using AITER Composable Kernel (CK).
    """
    (
        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

    # 1. 计算内存对齐 Padding
    # CK 算子要求数据 256 字节对齐，这里算出实际维度与对齐维度的差值
    hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
    intermediate_pad = config["d_expert_pad"] - config["d_expert"]

    # 2. 调用高度优化的融合算子
    # 该算子在内部完成：动态激活量化 (bf16 -> mxfp4) -> Stage 1 GEMM (伴随 SwiGLU) -> Stage 2 GEMM -> Top-K 累加
    output = fused_moe(
        hidden_states,
        gate_up_weight_shuffled,      # 必须传入 pre-shuffled (16x16 tile) 的权重
        down_weight_shuffled,         # 必须传入 pre-shuffled 的权重
        topk_weights,
        topk_ids,
        expert_mask=None,
        activation=ActivationType.Silu,
        quant_type=QuantType.per_1x32,  # MXFP4 对应的 1x32 动态分块量化
        doweight_stage1=False,
        w1_scale=gate_up_weight_scale_shuffled,
        w2_scale=down_weight_scale_shuffled,
        a1_scale=None,                  # 激活值缩放由底层 kernel 动态计算
        a2_scale=None,
        hidden_pad=hidden_pad,
        intermediate_pad=intermediate_pad,
    )

    return output