submission 688184

"""
Dynamic AITER_KSPLIT: set per-call based on problem shape.
Small M (token*topk <= expert*2) → KSPLIT=2 (cktile, faster)
Large M → KSPLIT=0 (default CK, no split-K)
"""
import os
import torch
from task import input_t, output_t
from aiter import ActivationType, QuantType
from aiter.fused_moe import fused_moe

_ACT = ActivationType.Silu
_QT = QuantType.per_1x32

def custom_kernel(data: input_t) -> output_t:
    (hs, _, _, _, _, w1, w2, w1s, w2s, tw, ti, cfg) = data
    M = hs.shape[0]
    topk = ti.shape[1]
    E = w1.shape[0]

    # Dynamic split-K: beneficial for small batch (low parallelism)
    # Harmful for large batch (already enough parallelism, split-K adds overhead)
    # Split-K benefits small batch + few experts (more CU parallelism)
    # Harmful for large batch (enough parallelism, split-K adds overhead)
    if M <= 128 and E <= 64:
        os.environ['AITER_KSPLIT'] = '2'
    else:
        os.environ['AITER_KSPLIT'] = '0'

    return fused_moe(
        hs, w1, w2, tw, ti,
        activation=_ACT, quant_type=_QT,
        w1_scale=w1s, w2_scale=w2s,
        hidden_pad=cfg["d_hidden_pad"] - cfg["d_hidden"],
        intermediate_pad=cfg["d_expert_pad"] - cfg["d_expert"],
    )