submission 587157

#!POPCORN leaderboard amd-moe-mxfp4
#!POPCORN gpu MI355X
import os
os.environ["AITER_USE_NT"] = "1"

import torch
import aiter
import aiter.fused_moe as _fm
from aiter.fused_moe import fused_moe
from aiter import ActivationType, QuantType
from task import input_t, output_t

_orig_ck_stage1 = _fm.ck_moe_stage1
_orig_ck_stage2 = aiter.ck_moe_stage2_fwd


def _force_nt_stage1(*args, **kwargs):
    if len(args) > 7:
        token = args[0].shape[0]
        E = args[1].shape[0]
        topk = args[7]
        kwargs['use_non_temporal_load'] = (token * topk // E) < 64
    return _orig_ck_stage1(*args, **kwargs)


def _force_nt_stage2(*args, **kwargs):
    if len(args) > 7:
        token = args[0].shape[0]
        E = args[1].shape[0]
        topk = args[7]
        kwargs['use_non_temporal_load'] = (token * topk // E) < 64
    return _orig_ck_stage2(*args, **kwargs)


_fm.ck_moe_stage1 = _force_nt_stage1
aiter.ck_moe_stage2_fwd = _force_nt_stage2

_KN1 = "moe_ck2stages_gemm1_64x32x32x128_32x32_1x1_16x4x1_16x4x1_1x4x1x16_4x4x1_1x1_intrawave_v3"

# FlyDSL stage2 — wider tile_n=256 (halves N-dimension tiles from 56 to 28)
_FLYDSL_S2_d256 = "flydsl_moe2_afp4_wfp4_bf16_t32x256x256_atomic"
_FLYDSL_S2_d512 = "flydsl_moe2_afp4_wfp4_bf16_t32x256x256_atomic"
_FLYDSL_S2_d2048 = "flydsl_moe2_afp4_wfp4_bf16_t32x128x256_atomic"

_initialized = False


def _ensure_configs():
    global _initialized
    if _initialized:
        return
    _initialized = True

    from aiter.jit.utils.chip_info import get_cu_num
    cu = get_cu_num()

    s = (
        "ActivationType.Silu",
        "torch.bfloat16",
        "torch.float4_e2m1fn_x2",
        "torch.float4_e2m1fn_x2",
        "QuantType.per_1x32",
        True,
        False,
    )

    configs = {}

    cktile_cfg = {
        'block_m': 32, 'ksplit': 2,
        'kernelName1': '', 'kernelName2': '', 'run_1stage': False,
    }
    for token in [8, 16, 32, 64, 128]:
        for mdim, idim, E in [(7168, 256, 257), (7168, 512, 33), (7168, 2048, 33)]:
            key = (cu, token, mdim, idim, E, 9) + s
            configs[key] = dict(cktile_cfg)

    for token in [256, 512, 1024]:
        key = (cu, token, 7168, 256, 257, 9) + s
        configs[key] = {
            'block_m': 32, 'ksplit': 0,
            'kernelName1': _KN1, 'kernelName2': _FLYDSL_S2_d256,
            'run_1stage': False,
        }

    for token in [256, 512, 1024]:
        key = (cu, token, 7168, 512, 33, 9) + s
        configs[key] = {
            'block_m': 32, 'ksplit': 0,
            'kernelName1': '', 'kernelName2': _FLYDSL_S2_d512,
            'run_1stage': False,
        }

    for token in [256, 512, 1024]:
        key = (cu, token, 7168, 2048, 33, 9) + s
        configs[key] = {
            'block_m': 64, 'ksplit': 0,
            'kernelName1': '', 'kernelName2': _FLYDSL_S2_d2048,
            'run_1stage': False,
        }

    _fm.cfg_2stages = configs


def custom_kernel(data: input_t) -> output_t:
    _ensure_configs()

    (
        hidden_states, gate_up_weight, down_weight,
        gate_up_weight_scale, down_weight_scale,
        gate_up_weight_shuffled, down_weight_shuffled,
        gate_up_weight_scale_shuffled, down_weight_scale_shuffled,
        topk_weights, topk_ids, config,
    ) = data

    hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
    intermediate_pad = config["d_expert_pad"] - config["d_expert"]

    output = fused_moe(
        hidden_states,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        topk_weights,
        topk_ids,
        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=hidden_pad,
        intermediate_pad=intermediate_pad,
    )

    return output