submission 639082

"""
MoE MXFP4 v63 — Optimal hybrid combining all discoveries:
- s1, s4 (bs=16): ksplit=2 real split-K (cktile, no quant, split_k=2)
- s2, s5 (bs=128): ksplit=2 but patched split_k=1 (cktile, no quant, no split overhead)
- s3 (bs=512, E=257): CK stage1 + FlyDSL stage2 reduce
- s6 (bs=512, E=33, d=512): CK stage1 + FlyDSL stage2 reduce
- s7 (bs=512, E=33, d=2048): CK with block_m=64
"""
from task import input_t, output_t

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

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

# For bs=128 shapes, force split_k=1 in cktile_moe_stage1
_orig_cktile_stage1 = cktile_moe_stage1
_force_splitk1 = set()  # Set of shape keys that should use split_k=1

def _patched_cktile_stage1(
    hidden_states, w1, w2, sorted_token_ids, sorted_expert_ids,
    num_valid_ids, out, topk, block_m, a1_scale, w1_scale,
    sorted_weights=None, n_pad_zeros=0, k_pad_zeros=0, bias1=None,
    activation=ActivationType.Silu, split_k=1, dtype=torch.bfloat16,
):
    token_num = hidden_states.shape[0]
    # For bs=128 shapes, force split_k=1
    actual_split_k = 1 if token_num > 16 else split_k
    return _orig_cktile_stage1(
        hidden_states, w1, w2, sorted_token_ids, sorted_expert_ids,
        num_valid_ids, out, topk, block_m, a1_scale, w1_scale,
        sorted_weights=sorted_weights, n_pad_zeros=n_pad_zeros,
        k_pad_zeros=k_pad_zeros, bias1=bias1,
        activation=activation, split_k=actual_split_k, dtype=dtype,
    )

fmoe_module.cktile_moe_stage1 = _patched_cktile_stage1

@functools.lru_cache(maxsize=2048)
def _patched_get_ksplit(token, topk, expert, inter_dim, model_dim):
    # cktile path for bs<=128
    if token <= 128:
        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,
    )

    # 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'

    # 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