submission 513713

import os
# OPUS variant of MoE sorting kernel: 3-6% improvement
os.environ["AITER_USE_OPUS_MOE_SORTING"] = "1"

import torch
import aiter
from aiter import ActivationType, QuantType, dtypes
from aiter.fused_moe import fused_moe, get_2stage_cfgs
from task import input_t, output_t

_current_ksplit = None


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

    M = hidden_states.shape[0]
    E = gate_up_weight_shuffled.shape[0]

    # Adaptive ksplit: cktile kernels with split-K=2 are much faster
    # for small batch sizes with few experts (E=33, M<=128).
    # For large batches, default ksplit=0 is better (less reduction overhead).
    want_ksplit = "2" if (E <= 64 and M <= 128) else None

    if want_ksplit != _current_ksplit:
        if want_ksplit:
            os.environ["AITER_KSPLIT"] = want_ksplit
        else:
            os.environ.pop("AITER_KSPLIT", None)
        get_2stage_cfgs.cache_clear()
        _current_ksplit = want_ksplit

    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