submission 600236

"""
Leaderboard-safe MoE optimizer.
- AITER_KSPLIT=7 for E=257, =2 for E=33 small batch
- CK2stages for bs=512 shapes via CSV
- AMD_DIRECT_DISPATCH + HSA_ENABLE_INTERRUPT env vars
- Only update ksplit env var when it changes (skip redundant cache clears)
- Pre-bind function references
"""
import os

os.environ['AMD_DIRECT_DISPATCH'] = '1'
os.environ['HSA_ENABLE_INTERRUPT'] = '0'
os.environ['AITER_KSPLIT'] = '7'

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,_tag"

custom_entries = [
    "256,512,7168,256,257,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,96.0,moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0.0%,69.0,moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16,1.3%,165.0,0,307.0,8632.0,",
    "256,512,7168,512,33,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,90.0,moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0.0%,60.0,moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16,0.0%,150.0,0,50.0,3000.0,",
    "256,512,7168,2048,33,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,128,0,175.0,moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0.0%,100.0,moe_ck2stages_gemm2_256x128x128x128_1x4_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16,7.1%,274.5,0,1478.48,5334.14,",
]

config_path = '/tmp/custom_tuned_fmoe.csv'
with open(config_path, 'w') as f:
    f.write(header + "\n" + "\n".join(custom_entries) + "\n")

aiter_root = '/home/runner/aiter'
default_cfg = f'{aiter_root}/aiter/configs/tuned_fmoe.csv'
os.environ['AITER_CONFIG_FMOE'] = f'{config_path}:{default_cfg}'

from task import input_t, output_t
import torch
from aiter import ActivationType, QuantType
from aiter.fused_moe import fused_moe
try:
    from aiter.fused_moe import get_ksplit
except ImportError:
    get_ksplit = None

_SILU = ActivationType.Silu
_PER_1x32 = QuantType.per_1x32
_fused_moe = fused_moe
_environ = os.environ

_has_cache_clear = get_ksplit is not None and hasattr(get_ksplit, 'cache_clear')
if _has_cache_clear:
    _cache_clear_fn = get_ksplit.cache_clear
    _cache_clear_fn()
else:
    _cache_clear_fn = lambda: None

_last_ksplit = '7'


def custom_kernel(data: input_t) -> output_t:
    global _last_ksplit
    (
        hidden_states, _guw, _dw, _guw_s, _dw_s,
        gate_up_weight_shuffled, down_weight_shuffled,
        gate_up_weight_scale_shuffled, down_weight_scale_shuffled,
        topk_weights, topk_ids, config,
    ) = data

    E = config["n_routed_experts"] + config["n_shared_experts"]
    M = hidden_states.shape[0]
    topk = config["total_top_k"]

    if E > 100:
        ksplit = '7'
    elif M * topk < E * 64:
        ksplit = '2'
    else:
        ksplit = '0'

    if ksplit != _last_ksplit:
        _environ['AITER_KSPLIT'] = ksplit
        _cache_clear_fn()
        _last_ksplit = ksplit

    return _fused_moe(
        hidden_states, gate_up_weight_shuffled, down_weight_shuffled,
        topk_weights, topk_ids,
        expert_mask=None, activation=_SILU,
        quant_type=_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=config["d_hidden_pad"] - config["d_hidden"],
        intermediate_pad=config["d_expert_pad"] - config["d_expert"],
    )