submission 671029

"""
V56: V52 + CSV priority fix (custom last) + per-shape FlyDSL tuning.
- Custom CSV now LAST in priority chain (highest priority)
- Shape 6 (d=512, K=512): t64x256x256_reduce (original, better for short K)
- Shape 7 (d=2048, K=2048): t64x128x256_reduce (more N-tiles, better for long K)
"""
#!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_S1_256x32 = "moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
CK_S1_64 = "moe_ck2stages_gemm1_256x64x128x128_1x4_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"
FLYDSL_S2_64 = "flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce"
FLYDSL_S2_64_128 = "flydsl_moe2_afp4_wfp4_bf16_t64x128x256_reduce"

_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]
    # E=33, d=512, bs=16/128: CK small tiles (block_m=32)
    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))
    # E=33, d=512, bs=512: block_m=64 + FlyDSL t64x256x256 (optimal for K=512)
    rows.append(_row(256, 512, 7168, 512, 33, 9, 64, CK_S1_64, FLYDSL_S2_64))
    # E=33, d=2048, bs=512: block_m=64 + FlyDSL t64x128x256 (optimal for K=2048)
    rows.append(_row(256, 512, 7168, 2048, 33, 9, 64, CK_S1_64, FLYDSL_S2_64_128))
    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 = []
        if os.path.exists(default_csv):
            paths.append(default_csv)
        if os.path.exists(dsv3_csv):
            paths.append(dsv3_csv)
        paths.append(csv_path)  # Custom last = highest priority (last-wins)

        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,
    )