submission 710606

"""
MoE MXFP4 tuned v11

Narrow cfg_2stages injection candidate:
- Only specialize the single most promising benchmark shape:
  bs=128, n_routed=32 (E=33 total), d_expert=512.
- Leave every other shape on baseline fused_moe() dispatch.
- Uses the proven post-initialization cfg_2stages injection mechanism from v8,
  but narrows the blast radius to one benchmark point.

Why this version:
- TileLang/custom assembled path has not beaten AITER baseline.
- Direct low-level CK assembly was unstable or slower.
- v8 showed the clearest small upside exactly on bs=128, E=33, d=512.
"""

import aiter.fused_moe as _fmoe
from aiter.fused_moe import fused_moe, get_2stage_cfgs
from aiter import ActivationType, QuantType
from task import input_t, output_t

_KN1_M32 = "moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
_KN2_M32 = "moe_ck2stages_gemm2_256x32x128x128_1x4_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"

_ACT = "ActivationType.Silu"
_DT = "torch.bfloat16"
_FP4 = "torch.float4_e2m1fn_x2"
_QT = "QuantType.per_1x32"

_TARGET_BS = 128
_TARGET_MODEL_DIM = 7168
_TARGET_INTER_DIM = 512
_TARGET_EXPERTS = 33
_TARGET_TOPK = 9


def _mk_key(cu_num: int) -> tuple:
    # Key format follows aiter.fused_moe._INDEX_COLS:
    # (cu_num, token, model_dim, inter_dim, expert, topk,
    #  act_type_str, dtype_str, q_dtype_a_str, q_dtype_w_str,
    #  q_type_str, use_g1u1, doweight_stage1)
    return (
        cu_num,
        _TARGET_BS,
        _TARGET_MODEL_DIM,
        _TARGET_INTER_DIM,
        _TARGET_EXPERTS,
        _TARGET_TOPK,
        _ACT,
        _DT,
        _FP4,
        _FP4,
        _QT,
        1,
        0,
    )


def _make_cfg() -> dict:
    return {
        "block_m": 32,
        "ksplit": 0,
        "kernelName1": _KN1_M32,
        "kernelName2": _KN2_M32,
        "run_1stage": 0,
        "use_non_temporal_load": False,
    }


_EXTRA = {}
for _cu in (256, 304):
    _EXTRA[_mk_key(_cu)] = _make_cfg()

_injected = False


def _baseline(
    hidden_states,
    gate_up_weight_shuffled,
    down_weight_shuffled,
    topk_weights,
    topk_ids,
    gate_up_weight_scale_shuffled,
    down_weight_scale_shuffled,
    hidden_pad,
    intermediate_pad,
):
    return 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,
    )


def _force_inject(
    hidden_states,
    gate_up_weight_shuffled,
    down_weight_shuffled,
    topk_weights,
    topk_ids,
    gate_up_weight_scale_shuffled,
    down_weight_scale_shuffled,
    hidden_pad,
    intermediate_pad,
):
    global _injected
    if _injected:
        return

    # Warm one real fused_moe call so cfg_2stages is initialized inside aiter.
    _baseline(
        hidden_states,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        topk_weights,
        topk_ids,
        gate_up_weight_scale_shuffled,
        down_weight_scale_shuffled,
        hidden_pad,
        intermediate_pad,
    )

    if _fmoe.cfg_2stages is not None:
        _fmoe.cfg_2stages.update(_EXTRA)
        get_2stage_cfgs.cache_clear()

    _injected = True


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

    is_target = (
        hidden_states.shape[0] == _TARGET_BS
        and config["d_hidden"] == _TARGET_MODEL_DIM
        and config["d_expert"] == _TARGET_INTER_DIM
        and (config["n_routed_experts"] + config.get("n_shared_experts", 0)) == _TARGET_EXPERTS
        and topk_ids.shape[1] == _TARGET_TOPK
    )

    if is_target and not _injected:
        _force_inject(
            hidden_states,
            gate_up_weight_shuffled,
            down_weight_shuffled,
            topk_weights,
            topk_ids,
            gate_up_weight_scale_shuffled,
            down_weight_scale_shuffled,
            hidden_pad,
            intermediate_pad,
        )

    return _baseline(
        hidden_states,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        topk_weights,
        topk_ids,
        gate_up_weight_scale_shuffled,
        down_weight_scale_shuffled,
        hidden_pad,
        intermediate_pad,
    )