submission 645189

import torch
import functools
import os
import sys

os.environ["HIP_FORCE_DEV_KERNARG"] = "1"
os.environ["AITER_USE_OPUS_MOE_SORTING"] = "0"
import aiter.fused_moe as _fm

_orig_get_ksplit = _fm.get_ksplit.__wrapped__

@functools.lru_cache(maxsize=4096)
def _patched_get_ksplit(token, topk, expert, inter_dim, model_dim):
    if token <= 128:
        return 2
    return _orig_get_ksplit(token, topk, expert, inter_dim, model_dim)

_fm.get_ksplit = _patched_get_ksplit

_orig_get_2stage_cfgs = _fm.get_2stage_cfgs.__wrapped__

@functools.lru_cache(maxsize=4096)
def _patched_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,
):
    if token <= 128:
        old_val = os.environ.get("AITER_BYPASS_TUNE_CONFIG", "0")
        os.environ["AITER_BYPASS_TUNE_CONFIG"] = "1"
        result = _orig_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,
        )
        os.environ["AITER_BYPASS_TUNE_CONFIG"] = old_val
        return result

    if token >= 512 and expert <= 33:
        try:
            from aiter.ops.flydsl.utils import is_flydsl_available
            if is_flydsl_available():
                from aiter.fused_moe import (
                    MOEMetadata, ck_moe_stage1, _flydsl_stage2_wrapper,
                )
                if inter_dim <= 512:
                    flydsl_kernel = "flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce"
                    block_m = 64
                else:
                    flydsl_kernel = "flydsl_moe2_afp4_wfp4_bf16_t64x128x256_reduce"
                    block_m = 64
                stage1_func = functools.partial(
                    ck_moe_stage1,
                    kernelName="",
                    activation=activation,
                    quant_type=q_type,
                    dtype=dtype,
                    splitk=0,
                    use_non_temporal_load=False,
                )
                stage2_func = functools.partial(
                    _flydsl_stage2_wrapper,
                    kernelName=flydsl_kernel,
                )
                return MOEMetadata(
                    stage1_func,
                    stage2_func,
                    block_m,
                    0,
                    False,
                )
        except Exception:
            pass

    if token >= 512 and expert >= 257:
        try:
            from aiter.ops.flydsl.utils import is_flydsl_available
            if is_flydsl_available():
                from aiter.fused_moe import (
                    MOEMetadata, ck_moe_stage1, _flydsl_stage2_wrapper,
                )
                tuned_s1 = ("moe_ck2stages_gemm1_64x32x32x128_1x1"
                            "_MulABScaleShuffled_v3_Nswizzle0"
                            "_Quant3_MulRoutedWeight0_silu"
                            "_FP4X2_FP4X2_B16")
                flydsl_s2 = "flydsl_moe2_afp4_wfp4_bf16_t32x256x256_atomic"
                stage1_func = functools.partial(
                    ck_moe_stage1,
                    kernelName=tuned_s1,
                    activation=activation,
                    quant_type=q_type,
                    dtype=dtype,
                    splitk=0,
                    use_non_temporal_load=True,
                )
                stage2_func = functools.partial(
                    _flydsl_stage2_wrapper,
                    kernelName=flydsl_s2,
                )
                return MOEMetadata(
                    stage1_func,
                    stage2_func,
                    32,
                    0,
                    False,
                )
        except Exception:
            pass

    return _orig_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,
    )

_fm.get_2stage_cfgs = _patched_get_2stage_cfgs

from task import input_t, output_t
from utils import make_match_reference
from reference import ref_kernel
from aiter import ActivationType, QuantType
from aiter.fused_moe import fused_moe


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
    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,
        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,
    )


check_implementation = make_match_reference(ref_kernel, rtol=5e-2, atol=5e-2)