submission 636207

"""
MoE MXFP4 v30 — All optimizations combined:
1. OPUS sorting
2. block_m=64 for s7 (d_expert>=2048) — VGPR optimization
3. Injected 256x32x128x128_1x4 stage1 for s5 — CK kernel name optimization
4. FlyDSL stage2 reduce for s3 (E=257, bs=512) — 13us improvement
5. FlyDSL stage2 reduce for s6 (E=33, bs=512, d=512) — 46us improvement!
"""
from task import input_t, output_t

import os
os.environ["AITER_USE_OPUS_MOE_SORTING"] = "1"

import aiter
from aiter import ActivationType, QuantType, dtypes
from aiter.fused_moe import fused_moe, get_2stage_cfgs
import aiter.fused_moe as fmoe_module

_injected = False

def _inject_configs():
    global _injected
    if _injected:
        return
    _injected = True

    try:
        get_2stage_cfgs(
            16, 7168, 256, 257, 9,
            dtypes.bf16, dtypes.fp4x2, dtypes.fp4x2,
            QuantType.per_1x32, True, ActivationType.Silu, False,
            0, 0, True,
        )
    except Exception:
        pass

    if fmoe_module.cfg_2stages is None:
        return

    common = (
        'ActivationType.Silu',
        'torch.bfloat16',
        'torch.float4_e2m1fn_x2',
        'torch.float4_e2m1fn_x2',
        'QuantType.per_1x32',
        True,
        False,
    )

    # s5 (128, 7168, 512, 33, 9): optimized stage1 kernel
    key_s5 = (256, 128, 7168, 512, 33, 9) + common
    if key_s5 not in fmoe_module.cfg_2stages:
        fmoe_module.cfg_2stages[key_s5] = {
            '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': False,
        }

    # s6 (512, 7168, 512, 33, 9): FlyDSL stage2 reduce (-46us!)
    key_s6 = (256, 512, 7168, 512, 33, 9) + common
    if key_s6 not in fmoe_module.cfg_2stages:
        fmoe_module.cfg_2stages[key_s6] = {
            'block_m': 64,
            'ksplit': 0,
            'kernelName1': '',
            'kernelName2': 'flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce',
            'run_1stage': False,
        }

    # s3 (512, 7168, 256, 257, 9): FlyDSL stage2 reduce (-13us)
    key_s3 = (256, 512, 7168, 256, 257, 9) + common
    if key_s3 in fmoe_module.cfg_2stages:
        fmoe_module.cfg_2stages[key_s3]['kernelName2'] = 'flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce'

    get_2stage_cfgs.cache_clear()


def custom_kernel(data: input_t) -> output_t:
    _inject_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"]

    d_expert = config["d_expert"]
    block_m = 64 if d_expert >= 2048 else None

    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,
        block_size_M=block_m,
    )

    return output