submission 670916

import functools
import os

import aiter
import aiter.fused_moe as aiter_fused_moe_module
import torch
from aiter import ActivationType, QuantType, dtypes as aiter_dtypes
from aiter.fused_moe import fused_moe

from task import input_t, output_t


# Current best confirmed profile:
# - E257/d256: mixed CK/FlyDSL rows
# - E33/d512: CK at bs=16, targeted override only at bs=128, CK+FlyDSL at bs=512
# - E33/d2048: CK at bs=32, CK+FlyDSL at bs=512
_TP4_128_VARIANT = "cktile_both"
_TP4_128_KSPLIT = 2

_KERNEL_64X32 = (
    "moe_ck2stages_gemm1_64x32x32x128_1x1_"
    "MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
)
_KERNEL_256X32 = (
    "moe_ck2stages_gemm1_256x32x128x128_1x4_"
    "MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
)
_KERNEL_256X64 = (
    "moe_ck2stages_gemm1_256x64x128x128_1x4_"
    "MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
)

_STAGE2_64X32_V1 = (
    "moe_ck2stages_gemm2_64x32x32x128_1x1_"
    "MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
)
_STAGE2_256X32_V3 = (
    "moe_ck2stages_gemm2_256x32x128x128_1x4_"
    "MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
)

_FLYDSL_32X128_REDUCE = "flydsl_moe2_afp4_wfp4_bf16_t32x128x256_reduce"
_FLYDSL_32X256_ATOMIC = "flydsl_moe2_afp4_wfp4_bf16_t32x256x256_atomic"
_FLYDSL_64X256_ATOMIC = "flydsl_moe2_afp4_wfp4_bf16_t64x256x256_atomic"

_E257_D256_ROWS = {
    16: (_KERNEL_64X32, _STAGE2_64X32_V1, 32),
    128: (_KERNEL_256X32, _STAGE2_64X32_V1, 32),
    512: (_KERNEL_64X32, _FLYDSL_32X256_ATOMIC, 32),
}
_E33_D512_ROWS = {
    16: (_KERNEL_64X32, _STAGE2_64X32_V1, 32),
    512: (_KERNEL_256X32, _FLYDSL_32X128_REDUCE, 32),
}
_E33_D2048_ROWS = {
    32: (_KERNEL_64X32, _STAGE2_256X32_V3, 32),
    512: (_KERNEL_256X64, _FLYDSL_64X256_ATOMIC, 64),
}


if hasattr(aiter_fused_moe_module, "_codex_orig_get_2stage_cfgs"):
    _ORIGINAL_GET_2STAGE_CFGS = aiter_fused_moe_module._codex_orig_get_2stage_cfgs
else:
    _ORIGINAL_GET_2STAGE_CFGS = aiter_fused_moe_module.get_2stage_cfgs
    aiter_fused_moe_module._codex_orig_get_2stage_cfgs = _ORIGINAL_GET_2STAGE_CFGS

_OVERRIDES_INSTALLED = False


def _is_e257_d256_shape(config: dict) -> bool:
    return (
        config["d_hidden"] == 7168
        and config["d_expert"] == 256
        and config["n_routed_experts"] == 256
        and config["n_shared_experts"] == 1
        and config["n_experts_per_token"] == 8
    )


def _is_tp4_d512_shape(config: dict) -> bool:
    return (
        config["d_hidden"] == 7168
        and config["d_expert"] == 512
        and config["n_routed_experts"] == 32
        and config["n_shared_experts"] == 1
        and config["n_experts_per_token"] == 8
    )


def _set_env(name: str, value: str | None) -> None:
    if value is None:
        os.environ.pop(name, None)
    else:
        os.environ[name] = value


def _configure_runtime(config: dict) -> None:
    use_nt = "1" if (_is_e257_d256_shape(config) or _is_tp4_d512_shape(config)) else None
    _set_env("AITER_USE_NT", use_nt)
    _set_env("AITER_KSPLIT", None)
    aiter_fused_moe_module.use_nt.cache_clear()
    aiter_fused_moe_module.get_ksplit.cache_clear()
    _ORIGINAL_GET_2STAGE_CFGS.cache_clear()
    cache_clear = getattr(aiter_fused_moe_module.get_2stage_cfgs, "cache_clear", None)
    if cache_clear is not None:
        cache_clear()


def _ck_metadata(
    kernel_name1: str,
    kernel_name2: str,
    block_m: int,
    *,
    use_nt: bool,
) -> aiter_fused_moe_module.MOEMetadata:
    stage1 = functools.partial(
        aiter_fused_moe_module.ck_moe_stage1,
        kernelName=kernel_name1,
        activation=ActivationType.Silu,
        quant_type=QuantType.per_1x32,
        dtype=torch.bfloat16,
        splitk=0,
        use_non_temporal_load=use_nt,
    )
    if kernel_name2.startswith("flydsl_"):
        stage2 = functools.partial(
            aiter_fused_moe_module._flydsl_stage2_wrapper,
            kernelName=kernel_name2,
        )
    else:
        stage2 = functools.partial(
            aiter.ck_moe_stage2_fwd,
            kernelName=kernel_name2,
            activation=ActivationType.Silu,
            quant_type=QuantType.per_1x32,
            use_non_temporal_load=use_nt,
        )
    return aiter_fused_moe_module.MOEMetadata(stage1, stage2, block_m, 0, False)


def _cktile_metadata(token: int, ksplit: int) -> aiter_fused_moe_module.MOEMetadata:
    block_m = 16 if token < 2048 else 32 if token < 16384 else 64
    return aiter_fused_moe_module.MOEMetadata(
        functools.partial(
            aiter_fused_moe_module.cktile_moe_stage1,
            n_pad_zeros=0,
            k_pad_zeros=0,
            activation=ActivationType.Silu,
            split_k=ksplit,
        ),
        functools.partial(
            aiter_fused_moe_module.cktile_moe_stage2,
            n_pad_zeros=0,
            k_pad_zeros=0,
            activation=ActivationType.Silu,
        ),
        block_m,
        ksplit,
        False,
    )


def _tp4_128_stage1_cktile_stage2_ck_metadata(
    token: int,
    ksplit: int,
) -> aiter_fused_moe_module.MOEMetadata:
    block_m = 16 if token < 2048 else 32 if token < 16384 else 64
    return aiter_fused_moe_module.MOEMetadata(
        functools.partial(
            aiter_fused_moe_module.cktile_moe_stage1,
            n_pad_zeros=0,
            k_pad_zeros=0,
            activation=ActivationType.Silu,
            split_k=ksplit,
        ),
        functools.partial(
            aiter.ck_moe_stage2_fwd,
            kernelName=_STAGE2_256X32_V3,
            activation=ActivationType.Silu,
            quant_type=QuantType.per_1x32,
            use_non_temporal_load=True,
        ),
        block_m,
        ksplit,
        False,
    )


def _tp4_128_metadata(token: int) -> aiter_fused_moe_module.MOEMetadata:
    if _TP4_128_VARIANT == "cktile_both":
        return _cktile_metadata(token, _TP4_128_KSPLIT)
    if _TP4_128_VARIANT == "cktile_stage1_ck_stage2":
        return _tp4_128_stage1_cktile_stage2_ck_metadata(token, _TP4_128_KSPLIT)
    raise ValueError(f"unknown _TP4_128_VARIANT={_TP4_128_VARIANT}")


def _row_metadata(
    rows: dict[int, tuple[str, str, int]],
    token: int,
    *,
    use_nt: bool,
) -> aiter_fused_moe_module.MOEMetadata | None:
    row = rows.get(token)
    if row is None:
        return None
    return _ck_metadata(row[0], row[1], row[2], use_nt=use_nt)


def _select_override(
    token: int,
    model_dim: int,
    inter_dim: int,
    expert: int,
    topk: int,
    dtype: torch.dtype,
    q_dtype_a: torch.dtype,
    q_dtype_w: torch.dtype,
    q_type,
    use_g1u1: bool,
    activation,
    doweight_stage1: bool,
) -> aiter_fused_moe_module.MOEMetadata | None:
    if (
        model_dim != 7168
        or topk != 9
        or dtype != torch.bfloat16
        or q_dtype_a != aiter_dtypes.fp4x2
        or q_dtype_w != aiter_dtypes.fp4x2
        or q_type != QuantType.per_1x32
        or not use_g1u1
        or activation != ActivationType.Silu
        or doweight_stage1
    ):
        return None

    if expert == 257 and inter_dim == 256:
        return _row_metadata(_E257_D256_ROWS, token, use_nt=True)

    if expert == 33 and inter_dim == 512:
        if token == 128:
            return _tp4_128_metadata(token)
        return _row_metadata(_E33_D512_ROWS, token, use_nt=True)

    if expert == 33 and inter_dim == 2048:
        return _row_metadata(_E33_D2048_ROWS, token, use_nt=False)

    return None


def _install_overrides() -> None:
    global _OVERRIDES_INSTALLED
    if _OVERRIDES_INSTALLED:
        return

    @functools.lru_cache(maxsize=128)
    def _wrapped_get_2stage_cfgs(
        token,
        model_dim,
        inter_dim,
        expert,
        topk,
        dtype,
        q_dtype_a,
        q_dtype_w,
        q_type,
        use_g1u1,
        activation,
        doweight_stage1,
        hidden_pad,
        intermediate_pad,
        is_shuffled=True,
    ):
        metadata = _select_override(
            token,
            model_dim,
            inter_dim,
            expert,
            topk,
            dtype,
            q_dtype_a,
            q_dtype_w,
            q_type,
            use_g1u1,
            activation,
            doweight_stage1,
        )
        if metadata is not None:
            return metadata
        return _ORIGINAL_GET_2STAGE_CFGS(
            token,
            model_dim,
            inter_dim,
            expert,
            topk,
            dtype,
            q_dtype_a,
            q_dtype_w,
            q_type,
            use_g1u1,
            activation,
            doweight_stage1,
            hidden_pad,
            intermediate_pad,
            is_shuffled,
        )

    aiter_fused_moe_module.get_2stage_cfgs = _wrapped_get_2stage_cfgs
    _OVERRIDES_INSTALLED = True


def _monolithic_fused_moe(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

    del gate_up_weight
    del down_weight
    del gate_up_weight_scale
    del down_weight_scale

    block_m = 32 if _is_tp4_d512_shape(config) else None
    hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
    intermediate_pad = config["d_expert_pad"] - config["d_expert"]

    return fused_moe(
        hidden_states,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        topk_weights,
        topk_ids,
        activation=ActivationType.Silu,
        quant_type=QuantType.per_1x32,
        w1_scale=gate_up_weight_scale_shuffled,
        w2_scale=down_weight_scale_shuffled,
        block_size_M=block_m,
        hidden_pad=hidden_pad,
        intermediate_pad=intermediate_pad,
    )


@torch.inference_mode()
def custom_kernel(data: input_t) -> output_t:
    config = data[-1]
    _install_overrides()
    _configure_runtime(config)
    return _monolithic_fused_moe(data)