submission 749999

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

import os
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx950'
os.environ.setdefault('AITER_KSPLIT', '0')

"""v57: v50 + block_size_M=32 for bs=128 E=33 d=512 (saves 3µs).
Best combination of all tested optimizations."""
import os
os.environ.setdefault("AITER_KSPLIT", "0")

import torch, torch.nn as nn
import aiter
from aiter import ActivationType, QuantType, dtypes
from aiter.fused_moe import fused_moe, get_cu_num, get_padded_M, get_inter_dim
import aiter.fused_moe as _fm
from task import input_t, output_t


def _optimal_dispatch():
    tune_file = _fm.AITER_CONFIGS.AITER_CONFIG_FMOE_FILE
    if _fm.cfg_2stages is None:
        try:
            import pandas as pd
            df = pd.read_csv(tune_file)
            if "_tag" in df.columns:
                df = df[df["_tag"].fillna("") == ""]
            cols = ["cu_num","token","model_dim","inter_dim","expert","topk",
                    "act_type","dtype","q_dtype_a","q_dtype_w","q_type",
                    "use_g1u1","doweight_stage1"]
            _fm.cfg_2stages = df.set_index(cols).to_dict("index")
        except Exception:
            return
    cu = get_cu_num()
    common = ("ActivationType.Silu", "torch.bfloat16",
              "torch.float4_e2m1fn_x2", "torch.float4_e2m1fn_x2",
              "QuantType.per_1x32", 1, 0)
    cktile = lambda bm=32: {"block_m": bm, "ksplit": 2, "kernelName1": "", "kernelName2": "",
              "us1": 0, "err1": 0, "us2": 0, "err2": 0,
              "us": 0, "run_1stage": 0, "tflops": 0, "bw": 0}
    K1_512 = "moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
    K2_512 = "moe_ck2stages_gemm2_256x32x128x128_1x4_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
    K1_2048 = "moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
    K2_2048 = "moe_ck2stages_gemm2_256x128x128x128_1x4_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
    def ck_entry(bm, k1, k2):
        return {"block_m": bm, "ksplit": 0, "kernelName1": k1, "kernelName2": k2,
                "us1": 0, "err1": 0, "us2": 0, "err2": 0,
                "us": 0, "run_1stage": 0, "tflops": 0, "bw": 0}

    # Small batch (padded_M<=128): cktile ksplit=2
    for token in [16, 32, 64, 128]:
        for E, inter_list in [(257, [256]), (33, [512, 2048])]:
            for inter in inter_list:
                key = (cu, token, 7168, inter, E, 9) + common
                _fm.cfg_2stages[key] = cktile(32)

    # Large batch: CK with tuned names
    for token in [256, 512, 1024]:
        for inter in [512]:
            key = (cu, token, 7168, inter, 33, 9) + common
            if key not in _fm.cfg_2stages:
                _fm.cfg_2stages[key] = ck_entry(32, K1_512, K2_512)
        for inter in [2048]:
            key = (cu, token, 7168, inter, 33, 9) + common
            if key not in _fm.cfg_2stages:
                _fm.cfg_2stages[key] = ck_entry(128, K1_2048, K2_2048)

try:
    _optimal_dispatch()
except Exception:
    pass


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, topk = topk_ids.shape
    padded_M = get_padded_M(M)
    E, _, inter_dim = get_inter_dim(gate_up_weight_shuffled.shape,
                                     down_weight_shuffled.shape)
    hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
    intermediate_pad = config["d_expert_pad"] - config["d_expert"]
    if not getattr(gate_up_weight_shuffled, 'is_shuffled', False):
        gate_up_weight_shuffled.is_shuffled = True
    if not getattr(down_weight_shuffled, 'is_shuffled', False):
        down_weight_shuffled.is_shuffled = True
    # block_size_M=32 for bs=128 E=33 d=512 (measured -3%)
    bsm = 32 if (padded_M == 128 and E < 64 and inter_dim <= 512) else None
    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,
        block_size_M=bsm,
        hidden_pad=hidden_pad, intermediate_pad=intermediate_pad)