submission 702538

"""V935: Remove d2048 CK override. V911 default=336 vs forced=348."""
import os
os.environ["HIP_FORCE_DEV_KERNARG"] = "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 fmoe_mod

_WINS = {
    (256, 16, 7168, 256, 257, 9, 'ActivationType.Silu', 'torch.bfloat16', 'torch.float4_e2m1fn_x2', 'torch.float4_e2m1fn_x2', 'QuantType.per_1x32', 1, 0): {'block_m': 32, 'ksplit': 4, 'kernelName1': None, 'kernelName2': None, 'run_1stage': 0, '_tag': float('nan')},
    (256, 128, 7168, 256, 257, 9, 'ActivationType.Silu', 'torch.bfloat16', 'torch.float4_e2m1fn_x2', 'torch.float4_e2m1fn_x2', 'QuantType.per_1x32', 1, 0): {'block_m': 32, 'ksplit': 4, 'kernelName1': None, 'kernelName2': None, 'run_1stage': 0, '_tag': float('nan')},
    (256, 512, 7168, 256, 257, 9, 'ActivationType.Silu', 'torch.bfloat16', 'torch.float4_e2m1fn_x2', 'torch.float4_e2m1fn_x2', 'QuantType.per_1x32', 1, 0): {'block_m': 32, 'ksplit': 0, 'kernelName1': 'moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16', 'kernelName2': 'moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16', 'run_1stage': 0, '_tag': float('nan')},
    (256, 16, 7168, 512, 33, 9, 'ActivationType.Silu', 'torch.bfloat16', 'torch.float4_e2m1fn_x2', 'torch.float4_e2m1fn_x2', 'QuantType.per_1x32', 1, 0): {'block_m': 64, 'ksplit': 2, 'kernelName1': None, 'kernelName2': None, 'run_1stage': 0, '_tag': float('nan')},
    (256, 128, 7168, 512, 33, 9, 'ActivationType.Silu', 'torch.bfloat16', 'torch.float4_e2m1fn_x2', 'torch.float4_e2m1fn_x2', 'QuantType.per_1x32', 1, 0): {'block_m': 64, 'ksplit': 2, 'kernelName1': None, 'kernelName2': None, 'run_1stage': 0, '_tag': float('nan')},
    (256, 512, 7168, 512, 33, 9, 'ActivationType.Silu', 'torch.bfloat16', 'torch.float4_e2m1fn_x2', 'torch.float4_e2m1fn_x2', 'QuantType.per_1x32', 1, 0): {'block_m': 32, 'ksplit': 0, 'kernelName1': 'moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16', 'kernelName2': 'moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16', 'run_1stage': 0, '_tag': float('nan')},
}
_patched = False
def _ensure_patched():
    global _patched
    if _patched:
        return
    _patched = True
    try:
        _ = fmoe_mod.get_2stage_cfgs.__wrapped__(
            16, 7168, 256, 257, 9, torch.bfloat16,
            torch.float4_e2m1fn_x2, torch.float4_e2m1fn_x2,
            QuantType.per_1x32, True, ActivationType.Silu, False, 0, 0, True)
    except: pass
    cfg = fmoe_mod.cfg_2stages
    if cfg is not None:
        for key, val in _WINS.items():
            cfg[key] = val
    try:
        fmoe_mod.get_2stage_cfgs.cache_clear()
        fmoe_mod.get_block_size_M.cache_clear()
        fmoe_mod.get_ksplit.cache_clear()
    except: pass
try:
    from aiter import dtypes
    _ensure_patched()
except: pass

def custom_kernel(data: input_t) -> output_t:
    (hidden_states, guw, dw, guws, dws, guws_sh, dws_sh, guws_sc_sh, dws_sc_sh,
     topk_weights, topk_ids, config) = data
    _ensure_patched()
    return fused_moe(
        hidden_states, guws_sh, dws_sh, topk_weights, topk_ids,
        expert_mask=None, activation=ActivationType.Silu,
        quant_type=QuantType.per_1x32, doweight_stage1=False,
        w1_scale=guws_sc_sh, w2_scale=dws_sc_sh, a1_scale=None, a2_scale=None,
        hidden_pad=config["d_hidden_pad"] - config["d_hidden"],
        intermediate_pad=config["d_expert_pad"] - config["d_expert"],
    )