submission 693836

"""
v28: Full tuned CSV for both E=33 shapes (d=512 and d=2048).

Combines v26 (d=512 tuned, geomean ~150 µs) with v27's tuning result for d=2048.

E=33 d=512: ck2stages 256x32x128x128_1x4 (gemm1) + 64x32x32x128_1x1 (gemm2), block_m=32
E=33 d=2048: ck2stages 256x128x128x128_1x4 (gemm1) + 256x128x128x128_1x4 (gemm2), block_m=128

v27 tuning showed: d=2048 GEMM time 272 µs (vs heuristic benchmark ~355 µs)
Expected improvement: ~23% on d=2048 shape.
"""
import os
import tempfile

_KN1_D512 = ("moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScale"
             "Shuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0"
             "_silu_FP4X2_FP4X2_B16")
_KN2_D512 = ("moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpert"
             "WeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1"
             "_FP4X2_FP4X2_B16")

_KN1_D2048 = ("moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScale"
              "Shuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0"
              "_silu_FP4X2_FP4X2_B16")
_KN2_D2048 = ("moe_ck2stages_gemm2_256x128x128x128_1x4_MulABScaleExpert"
              "WeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1"
              "_FP4X2_FP4X2_B16")

_HDR = ("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,_tag")

_COMMON = ("ActivationType.Silu,torch.bfloat16,"
           "torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,"
           "QuantType.per_1x32,1,0")

_ROWS = [
    f"256,512,7168,512,33,9,{_COMMON},32,0,0,{_KN1_D512},0,0,{_KN2_D512},0,129,0,787,2905,",
    f"256,512,7168,2048,33,9,{_COMMON},128,0,0,{_KN1_D2048},0,0,{_KN2_D2048},0,272,0,1491,5380,",
]

_csv = tempfile.NamedTemporaryFile(mode='w', suffix='_e33_tuned_fmoe.csv',
                                   delete=False)
_csv.write(_HDR + "\n")
for r in _ROWS:
    _csv.write(r + "\n")
_csv.flush()

_aiter_root = "/home/runner/aiter/aiter/configs"
_default = f"{_aiter_root}/tuned_fmoe.csv"
_dsv3 = f"{_aiter_root}/model_configs/dsv3_fp4_tuned_fmoe.csv"
os.environ["AITER_CONFIG_FMOE"] = f"{_csv.name}:{_default}:{_dsv3}"

os.environ["AITER_USE_OPUS_MOE_SORTING"] = "1"

import aiter.fused_moe as _fm
_fm._USE_OPUS_MOE_SORTING = True

import torch
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:
    (
        hidden_states, _, _, _, _,
        gate_up_weight_shuffled, down_weight_shuffled,
        gate_up_weight_scale_shuffled, down_weight_scale_shuffled,
        topk_weights, topk_ids, config,
    ) = data

    hp = config["d_hidden_pad"] - config["d_hidden"]
    ip = config["d_expert_pad"] - config["d_expert"]
    n_exp = config["n_routed_experts"] + config["n_shared_experts"]
    top_k = config["total_top_k"]
    bs = hidden_states.shape[0]
    est_m = bs * top_k // n_exp

    use_a16w4 = (n_exp > 64 and est_m < 10) or (n_exp <= 64 and est_m < 50)

    if use_a16w4:
        os.environ["AITER_BYPASS_TUNE_CONFIG"] = "1"
        os.environ["AITER_KSPLIT"] = "2"
    else:
        os.environ.pop("AITER_BYPASS_TUNE_CONFIG", None)
        os.environ.pop("AITER_KSPLIT", None)

    gate_up_weight_shuffled.is_shuffled = True
    down_weight_shuffled.is_shuffled = True

    kw = dict(
        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,
        hidden_pad=hp, intermediate_pad=ip,
    )

    return fused_moe(
        hidden_states, gate_up_weight_shuffled, down_weight_shuffled,
        topk_weights, topk_ids, **kw,
    )