submission 699911

from task import input_t, output_t

import os

from aiter import ActivationType, QuantType
from aiter.fused_moe import fused_moe


def _normalize_aiter_env() -> None:
    # Clear stale AITER process state so this shell behaves deterministically
    # even if the runner previously executed a different AITER-based variant.
    os.environ["AITER_ONLINE_TUNE"] = "0"
    os.environ["AITER_BYPASS_TUNE_CONFIG"] = "0"
    os.environ["AITER_USE_OPUS_MOE_SORTING"] = "0"
    os.environ.pop("AITER_CONFIG_FMOE", None)
    os.environ.pop("AITER_LOG_TUNED_CONFIG", None)


def _configure_aiter_env(routed_experts: int, token_count: int) -> None:
    _normalize_aiter_env()

    if routed_experts in (32, 33) and token_count <= 128:
        os.environ["AITER_KSPLIT"] = "2"
    else:
        os.environ["AITER_KSPLIT"] = "0"

    if routed_experts in (32, 33) and token_count <= 128:
        os.environ["AITER_USE_NT"] = "1"
    else:
        os.environ["AITER_USE_NT"] = "-1"


def _block_size_m(routed_experts: int, token_count: int) -> int | None:
    # block_size_M=32 helps shape 5 (M=128, 34tok/expert) by 8% vs cktile
    # default 16 but hurts shape 4 (M=16, 4tok/expert) by 15%.
    if routed_experts in (32, 33) and 64 < token_count <= 128:
        return 32
    return None


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

    E = config["n_routed_experts"]
    M = hidden_states.shape[0]
    _configure_aiter_env(E, M)
    bm = _block_size_m(E, M)

    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=config["d_hidden_pad"] - config["d_hidden"],
        intermediate_pad=config["d_expert_pad"] - config["d_expert"],
        block_size_M=bm,
    )