submission 733319

"""ew: Case 5 breakthrough - S1_256x128+FS2_32_128_A+bm128 = 141us vs 166us.
Key finding: FS2_32_128_A (atomic) makes S1_256x128 correct for Case 5.
FS2_64_256_R (reduce) fails correctness for Case 5 with S1_256x128.
Benchmark geomean: ~127.6us. Ranked geomean: ~129.8us (new best).
"""
import os
import torch
from task import input_t, output_t
from aiter import ActivationType, QuantType
from aiter.fused_moe import fused_moe
import aiter.fused_moe as fm

_S1_256x128 = "moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
_S1_256x32 = "moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
_FS2_64_256_R = "flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce"
_FS2_32_128_A = "flydsl_moe2_afp4_wfp4_bf16_t32x128x256_atomic"
_configs_injected = False
_call_count = 0

_KSPLIT_POLICY = {
    (16, 257, 256): ("7", "1"),
    (128, 257, 256): ("0", "0"),
    (512, 257, 256): ("0", "0"),
    (16, 33, 512): ("4", "1"),
    (128, 33, 512): ("0", "0"),
    (512, 33, 512): ("0", "0"),
    (512, 33, 2048): ("0", "0"),
}


@torch.inference_mode()
def custom_kernel(data: input_t) -> output_t:
    (hs, guw, dw, guws, dws, guw_sh, dw_sh, guws_sh, dws_sh, tw, ti, cfg) = data
    hp = cfg["d_hidden_pad"] - cfg["d_hidden"]
    ip = cfg["d_expert_pad"] - cfg["d_expert"]
    M = cfg["bs"]
    E = cfg["n_routed_experts"] + cfg["n_shared_experts"]
    d_expert = cfg["d_expert"]
    top_k = cfg["total_top_k"]
    est_m = (M * top_k) // E

    shape = (M, E, d_expert)
    if shape in _KSPLIT_POLICY:
        ks, bp = _KSPLIT_POLICY[shape]
    else:
        if E > 64:
            ks, bp = ("2", "1") if est_m <= 0 else ("0", "0")
        else:
            ks, bp = ("2", "1") if est_m <= 16 else ("0", "0")
    os.environ["AITER_KSPLIT"] = ks
    os.environ["AITER_BYPASS_TUNE_CONFIG"] = bp

    global _configs_injected, _call_count
    _call_count += 1
    if not _configs_injected and _call_count == 2 and fm.cfg_2stages is not None:
        _configs_injected = True
        _A = "ActivationType.Silu"
        _D = "torch.bfloat16"
        _QA = "torch.float4_e2m1fn_x2"
        _QW = "torch.float4_e2m1fn_x2"
        _QT = "QuantType.per_1x32"

        def mk(t, i, e):
            return (256, t, 7168, i, e, 9, _A, _D, _QA, _QW, _QT, 1, 0)

        cfg_e257 = {
            "ksplit": 0,
            "kernelName1": _S1_256x128,
            "kernelName2": _FS2_64_256_R,
            "run_1stage": False,
            "block_m": 32,
        }
        fm.cfg_2stages[mk(128, 256, 257)] = {**cfg_e257}
        fm.cfg_2stages[mk(512, 256, 257)] = {**cfg_e257}

        fm.cfg_2stages[mk(128, 512, 33)] = {
            "ksplit": 0,
            "kernelName1": _S1_256x128,
            "kernelName2": _FS2_64_256_R,
            "run_1stage": False,
            "block_m": 32,
        }

        fm.cfg_2stages[mk(512, 512, 33)] = {
            "ksplit": 0,
            "kernelName1": _S1_256x128,
            "kernelName2": _FS2_32_128_A,
            "run_1stage": False,
            "block_m": 128,
        }

        fm.cfg_2stages[mk(512, 2048, 33)] = {
            "ksplit": 0,
            "kernelName1": _S1_256x128,
            "kernelName2": _FS2_32_128_A,
            "run_1stage": False,
            "block_m": 128,
        }
        fm.get_2stage_cfgs.cache_clear()

    return fused_moe(
        hs, guw_sh, dw_sh, tw, ti,
        expert_mask=None,
        activation=ActivationType.Silu,
        quant_type=QuantType.per_1x32,
        doweight_stage1=False,
        w1_scale=guws_sh,
        w2_scale=dws_sh,
        a1_scale=None,
        a2_scale=None,
        hidden_pad=hp,
        intermediate_pad=ip,
    )