submission 638848

"""
MoE MXFP4 v53 — Extend ksplit=2 to both s1 and s4 (all bs=16 shapes).
v52 showed ksplit=2 gives s4: 92→60.9us (-34%). Try for s1 too.
s1 is E=257, d=256 — model_dim=7168 which is divisible by 2*256=512.
"""
from task import input_t, output_t

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

import functools
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

_original_get_ksplit = fmoe_module.get_ksplit.__wrapped__

@functools.lru_cache(maxsize=2048)
def _patched_get_ksplit(token, topk, expert, inter_dim, model_dim):
    # Force ksplit=2 for all bs=16 shapes (token=16)
    if token <= 16:
        if model_dim % 2 == 0 and (model_dim // 2) % 256 == 0:
            return 2
    return 0

fmoe_module.get_ksplit = _patched_get_ksplit

_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,
    )

    ck_stage1_256x128 = 'moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16'
    ck_stage2_small = 'moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16'

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

    # s5: 256x128 stage1
    key = (256, 128, 7168, 512, 33, 9) + common
    if key not in fmoe_module.cfg_2stages:
        fmoe_module.cfg_2stages[key] = {
            'block_m': 32, 'ksplit': 0,
            'kernelName1': ck_stage1_256x128,
            'kernelName2': ck_stage2_small,
            'run_1stage': False,
        }

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

    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