submission 665865

"""
V25: Optimized MoE-MXFP4 with tuned CK configs for E=33, d=512 shapes only.
E=33 d=2048 and E=257 shapes left to DSV3+heuristic.
Key finding: block_m=32 with CK_S1_32/CK_S2_32_v1 is best for E=33, d=512.
"""
#!POPCORN leaderboard amd-moe-mxfp4
#!POPCORN gpu MI355X
import torch
import os
import aiter
from aiter import ActivationType, QuantType, dtypes
from aiter.fused_moe import fused_moe, get_2stage_cfgs
from task import input_t, output_t

CK_S1_32 = "moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
CK_S2_32_v1 = "moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"

_CSV_HEADER = "cu_num,token,model_dim,inter_dim,expert,topk,act_type,dtype,q_dtype_a,q_dtype_w,q_type,use_g1u1,doweight_stage1,block_m,ksplit,us1,kernelName1,err1,us2,kernelName2,err2,us,run_1stage,tflops,bw"

def _row(cu, tok, mdim, idim, E, topk, bm, k1, k2, ks=0):
    return f"{cu},{tok},{mdim},{idim},{E},{topk},ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,True,False,{bm},{ks},0,{k1},0,0,{k2},0,0,False,0,0"

def _build_custom_csv():
    rows = [_CSV_HEADER]
    # Only E=33, d=512 shapes. All others use DSV3/heuristic.
    # V23/V24 showed: block_m=32, CK_S1_32, CK_S2_32_v1 is best for these.
    rows.append(_row(256, 16, 7168, 512, 33, 9, 32, CK_S1_32, CK_S2_32_v1))
    rows.append(_row(256, 128, 7168, 512, 33, 9, 32, CK_S1_32, CK_S2_32_v1))
    rows.append(_row(256, 512, 7168, 512, 33, 9, 32, CK_S1_32, CK_S2_32_v1))
    return "\n".join(rows) + "\n"

_initialized = False

def custom_kernel(data: input_t) -> output_t:
    global _initialized

    (
        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

    d_hidden = config["d_hidden"]
    d_hidden_pad = config["d_hidden_pad"]
    d_expert_pad = config["d_expert_pad"]
    hidden_pad = d_hidden_pad - d_hidden
    intermediate_pad = d_expert_pad - config["d_expert"]

    w1s = gate_up_weight_scale_shuffled.view(dtypes.fp8_e8m0)
    w2s = down_weight_scale_shuffled.view(dtypes.fp8_e8m0)

    if not _initialized:
        csv_path = "/tmp/custom_tuned_fmoe.csv"
        with open(csv_path, 'w') as f:
            f.write(_build_custom_csv())

        aiter_root = os.path.dirname(os.path.abspath(aiter.__file__))
        default_csv = os.path.join(aiter_root, "configs", "tuned_fmoe.csv")
        dsv3_csv = os.path.join(aiter_root, "configs", "model_configs", "dsv3_fp4_tuned_fmoe.csv")

        paths = [csv_path]
        if os.path.exists(dsv3_csv):
            paths.append(dsv3_csv)
        if os.path.exists(default_csv):
            paths.append(default_csv)

        os.environ["AITER_CONFIG_FMOE"] = ":".join(paths)
        get_2stage_cfgs.cache_clear()
        _initialized = True

    return fused_moe(
        hidden_states, gate_up_weight_shuffled, down_weight_shuffled,
        topk_weights, topk_ids,
        activation=ActivationType.Silu, quant_type=QuantType.per_1x32,
        w1_scale=w1s, w2_scale=w2s,
        hidden_pad=hidden_pad, intermediate_pad=intermediate_pad,
    )