submission 710356

"""
MoE_V59: V55 (best=141.8µs) + AMD_DIRECT_DISPATCH=1.

V57 used AMD_DIRECT_DISPATCH=1 + run_1stage=True for d=2048. It failed
because run_1stage caused a timeout, NOT because of AMD_DIRECT_DISPATCH.

AMD_DIRECT_DISPATCH bypasses the HIP runtime dispatch queue, reducing
kernel launch latency. Combined with HIP_FORCE_DEV_KERNARG, this should
further cut dispatch overhead for the 7-shape benchmark.

No config changes from V55 — purely testing the env var impact.
"""
import os
os.environ["HIP_FORCE_DEV_KERNARG"] = "1"
os.environ["GPU_MAX_HW_QUEUES"] = "2"
os.environ["AMD_DIRECT_DISPATCH"] = "1"

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"
_S2_64x32 = "moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
_configs_injected = False
_call_count = 0

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"]
    top_k = cfg["total_top_k"]; est_m = (M * top_k) // E

    if E > 64:
        if est_m <= 0:
            os.environ["AITER_KSPLIT"] = "2"; os.environ["AITER_BYPASS_TUNE_CONFIG"] = "1"
        else:
            os.environ["AITER_KSPLIT"] = "0"; os.environ["AITER_BYPASS_TUNE_CONFIG"] = "0"
    else:
        if est_m <= 34:
            os.environ["AITER_KSPLIT"] = "2"; os.environ["AITER_BYPASS_TUNE_CONFIG"] = "1"
        else:
            os.environ["AITER_KSPLIT"] = "0"; os.environ["AITER_BYPASS_TUNE_CONFIG"] = "0"

    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":_S2_64x32, "run_1stage":False}

        fm.cfg_2stages[mk(128,256,257)] = {**cfg_e257, "block_m": 32}
        fm.cfg_2stages[mk(512,256,257)] = {**cfg_e257, "block_m": 32}

        cfg_sm = {"ksplit":0, "kernelName1":"moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16",
                  "kernelName2":_S2_64x32, "run_1stage":False}
        fm.cfg_2stages[mk(512,512,33)] = {**cfg_sm, "block_m": 32}

        fm.cfg_2stages[mk(512,2048,33)] = {
            "block_m": 64, "ksplit": 0,
            "kernelName1": "", "kernelName2": "",
            "run_1stage": False,
        }
        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)