submission 612670

"""v535"""
import functools
import os
import sys

import torch
import triton

from dataclasses import replace

_dsv3_path = "/home/runner/aiter/aiter/configs/model_configs/dsv3_fp4_tuned_fmoe.csv"
_flydsl_s3_stage2 = "flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce"
try:
    with open(_dsv3_path, "r") as f:
        lines = f.readlines()
    header = lines[0].strip()
    modified_lines = [header + "\n"]
    for line in lines[1:]:
        stripped = line.strip()
        if not stripped:
            continue
        fields = stripped.split(",")
        try:
            token_val = int(fields[1])
            expert_val = int(fields[4])
        except (ValueError, IndexError):
            modified_lines.append(line)
            continue
        if expert_val == 257 and token_val <= 128:
            fields[14] = "2"
            modified_lines.append(",".join(fields) + "\n")
        elif expert_val == 257 and token_val == 512:
            flydsl_fields = list(fields)
            flydsl_fields[19] = _flydsl_s3_stage2
            flydsl_fields[20] = "0.1%"
            if len(flydsl_fields) > 25:
                flydsl_fields[25] = ""
            modified_lines.append(",".join(flydsl_fields) + "\n")
            fallback_fields = list(fields)
            if len(fallback_fields) > 25:
                fallback_fields[25] = "flydsl_fallback"
            else:
                fallback_fields.append("flydsl_fallback")
            modified_lines.append(",".join(fallback_fields) + "\n")
        else:
            modified_lines.append(line)
    with open(_dsv3_path, "w") as f:
        f.writelines(modified_lines)
except Exception as e:
    print(f"[v490] dsv3 error: {e}", file=sys.stderr)

_csv_header = (
    "cu_num,token,model_dim,inter_dim,expert,topk,act_type,dtype,q_dtype_a,"
    "q_dtype_w,q_type,use_g1u1,doweight_stage1,block_m,ksplit,us1,kernelName1,"
    "err1,us2,kernelName2,err2,us,run_1stage,tflops,bw,_tag"
)
_common = (
    "ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,"
    "torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0"
)
_k1_512 = (
    "moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_"
    "Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
)
_k2_512 = (
    "moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_"
    "Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
)
_k1_2048 = (
    "moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_"
    "Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
)
_k2_2048 = (
    "moe_ck2stages_gemm2_256x128x128x128_1x4_MulABScaleExpertWeightShuffled_v3_"
    "Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
)
_k2_512_flydsl = "flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce"
_k2_2048_flydsl = "flydsl_moe2_afp4_wfp4_bf16_t32x128x256_atomic"
_csv_rows = [
    f"256,512,7168,512,33,9,{_common},32,0,0,{_k1_512},0.0%,90.0,{_k2_512_flydsl},0.1%,219.79,0,781.78,2884.18,",
    f"256,512,7168,512,33,9,{_common},32,0,0,{_k1_512},0.0%,0,{_k2_512},0.0%,129.79,0,781.78,2884.18,flydsl_fallback",
    f"256,512,7168,2048,33,9,{_common},64,0,0,{_k1_2048},0.0%,180.0,{_k2_2048_flydsl},0.1%,455.08,0,1475.47,5323.27,",
    f"256,512,7168,2048,33,9,{_common},128,0,0,{_k1_2048},0.0%,0,{_k2_2048},0.0%,275.08,0,1475.47,5323.27,flydsl_fallback",
]
try:
    e33_path = "/home/runner/aiter/aiter/configs/model_configs/e33_fp4_tuned_fmoe.csv"
    with open(e33_path, "w") as f:
        f.write(_csv_header + "\n")
        for row in _csv_rows:
            f.write(row + "\n")
except Exception:
    pass

os.environ["AITER_USE_OPUS_MOE_SORTING"] = "1"
os.environ["AITER_USE_NT"] = "1"
os.environ.pop("FLIR_CK_LDS128", None)
os.environ["FLIR_MOE_STAGE1_SCHED"] = "1"
os.environ["FLIR_MOE_STAGE2_SCHED"] = "1"
os.environ["FLIR_MOE_STAGE2_PERSIST_M"] = "1"

_flydsl_stage1_bug_path = (
    "/home/runner/aiter/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py"
)
try:
    with open(_flydsl_stage1_bug_path, "r") as f:
        _flydsl_src = f.read()
    _old_stage1_sig = (
        "def compute_f8f6f4_tile(\n"
        "                        acc_gate_in,\n"
        "                        acc_up_in,\n"
        "                        b_tile_in_gate,\n"
        "                        b_tile_in_up,\n"
        "                        lds_base,\n"
        "                        *,\n"
        "                        a0_prefetch=None,\n"
        "                        a_scale=None,\n"
        "                        b_scale_gate=None,\n"
        "                        b_scale_up=None,\n"
        "                        prefetch_epilogue: bool = False,\n"
        "                    ):"
    )
    _new_stage1_sig = (
        "def compute_f8f6f4_tile(\n"
        "                        acc_gate_in,\n"
        "                        acc_up_in,\n"
        "                        b_tile_in,\n"
        "                        lds_base,\n"
        "                        *,\n"
        "                        a0_prefetch=None,\n"
        "                        a_scale=None,\n"
        "                        b_scale=None,\n"
        "                        prefetch_epilogue: bool = False,\n"
        "                    ):"
    )
    if _old_stage1_sig in _flydsl_src:
        _flydsl_src = _flydsl_src.replace(_old_stage1_sig, _new_stage1_sig)
    _stage1_sched_disabled = (
        "                        # hot_loop_scheduler()\n"
        "                        gpu.barrier()"
    )
    _stage1_sched_enabled = (
        "                        hot_loop_scheduler()\n"
        "                        gpu.barrier()"
    )
    if _stage1_sched_disabled in _flydsl_src:
        _flydsl_src = _flydsl_src.replace(
            _stage1_sched_disabled,
            _stage1_sched_enabled,
            3,
        )
    with open(_flydsl_stage1_bug_path, "w") as f:
        f.write(_flydsl_src)
except Exception as e:
    print(f"[v530] flydsl stage1 source patch skipped: {e}", file=sys.stderr)

from task import input_t, output_t
from aiter import ActivationType, QuantType, dtypes
import aiter
import aiter.fused_moe as fused_moe_mod
from aiter.ops.triton.quant.fused_mxfp4_quant import (
    fused_dynamic_mxfp4_quant_moe_sort,
    _fused_dynamic_mxfp4_quant_moe_sort_kernel,
)
from aiter.ops.flydsl.moe_kernels import flydsl_moe_stage1, flydsl_moe_stage2


_SORT_BUFS = {}


def _cached_moe_sorting_impl(
    topk_ids, topk_weights, num_experts, model_dim, moebuf_dtype,
    block_size, expert_mask, num_local_tokens, dispatch_policy, use_opus,
):
    device = topk_ids.device
    M, topk = topk_ids.shape
    key = (M, num_experts, block_size, model_dim)

    if key not in _SORT_BUFS:
        max_num_tokens_padded = int(M * topk + num_experts * block_size - topk)
        max_num_m_blocks = int(
            (max_num_tokens_padded + block_size - 1) // block_size
        )
        _SORT_BUFS[key] = (
            torch.empty(max_num_tokens_padded, dtype=dtypes.i32, device=device),
            torch.empty(max_num_tokens_padded, dtype=dtypes.fp32, device=device),
            torch.empty(max_num_m_blocks, dtype=dtypes.i32, device=device),
            torch.empty(2, dtype=dtypes.i32, device=device),
            torch.empty((M, model_dim), dtype=moebuf_dtype, device=device),
        )

    sid, sw, sei, nvi, mb = _SORT_BUFS[key]
    fwd = aiter.moe_sorting_opus_fwd if use_opus else aiter.moe_sorting_fwd
    fwd(
        topk_ids, topk_weights, sid, sw, sei, nvi, mb,
        num_experts, int(block_size), expert_mask, num_local_tokens,
        dispatch_policy,
    )
    return sid, sw, sei, nvi, mb


fused_moe_mod._moe_sorting_impl = _cached_moe_sorting_impl


_ORIGINAL_GET_2STAGE_CFGS = fused_moe_mod.get_2stage_cfgs

_CKTILE_BUFS = {}


def _cached_cktile_moe_stage1(
    hidden_states, w1, w2,
    sorted_token_ids, sorted_expert_ids, num_valid_ids,
    out, topk, block_m,
    a1_scale, w1_scale, sorted_weights=None,
    n_pad_zeros=0, k_pad_zeros=0, bias1=None,
    activation=ActivationType.Silu, split_k=1, dtype=torch.bfloat16,
):
    token_num = hidden_states.shape[0]
    _, n1, k1 = w1.shape
    _, k2, n2 = w2.shape
    D = n2 if k2 == k1 else n2 * 2
    if w1.dtype is torch.uint32:
        D = D * 8

    buf_key = (token_num, topk, D, w1.shape[1], split_k, hidden_states.device)
    if buf_key not in _CKTILE_BUFS:
        _CKTILE_BUFS[buf_key] = (
            torch.empty((token_num, topk, D), dtype=dtype, device=hidden_states.device),
            torch.zeros(
                (token_num, topk, w1.shape[1]), dtype=hidden_states.dtype,
                device=hidden_states.device,
            ) if split_k > 1 else None,
        )

    out_buf, tmp_buf = _CKTILE_BUFS[buf_key]

    if split_k > 1:
        tmp_buf.zero_()
        aiter.moe_cktile2stages_gemm1(
            hidden_states, w1, tmp_buf,
            sorted_token_ids, sorted_expert_ids, num_valid_ids,
            topk, n_pad_zeros, k_pad_zeros,
            sorted_weights, a1_scale, w1_scale, bias1,
            activation, block_m, split_k,
        )
        aiter.silu_and_mul(out_buf, tmp_buf)
    else:
        aiter.moe_cktile2stages_gemm1(
            hidden_states, w1, out_buf,
            sorted_token_ids, sorted_expert_ids, num_valid_ids,
            topk, n_pad_zeros, k_pad_zeros,
            sorted_weights, a1_scale, w1_scale, bias1,
            activation, block_m, split_k,
        )
    return out_buf


def _make_cktile_metadata(hidden_pad, intermediate_pad, use_g1u1, activation, split_k=2):
    return fused_moe_mod.MOEMetadata(
        functools.partial(
            _cached_cktile_moe_stage1,
            n_pad_zeros=intermediate_pad // 64 * 64 * (2 if use_g1u1 else 1),
            k_pad_zeros=hidden_pad // 128 * 128,
            activation=activation,
            split_k=split_k,
        ),
        functools.partial(
            fused_moe_mod.cktile_moe_stage2,
            n_pad_zeros=hidden_pad // 64 * 64,
            k_pad_zeros=intermediate_pad // 128 * 128,
            activation=activation,
        ),
        16, split_k, False, False, True,
    )


@functools.lru_cache(maxsize=2048)
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,
):
    common = (
        model_dim == 7168 and topk == 9
        and dtype == dtypes.bf16 and q_dtype_a == dtypes.fp4x2
        and q_dtype_w == dtypes.fp4x2 and q_type == QuantType.per_1x32
        and use_g1u1 and not doweight_stage1 and is_shuffled
    )

    if not common:
        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,
        )

    if expert == 257 and inter_dim == 256 and token == 16:
        return _make_cktile_metadata(
            hidden_pad, intermediate_pad, use_g1u1, activation, split_k=2,
        )
    if expert == 257 and inter_dim == 256 and token == 128:
        return _make_cktile_metadata(
            hidden_pad, intermediate_pad, use_g1u1, activation, split_k=4,
        )

    if expert == 33 and inter_dim == 512 and token <= 128:
        return _make_cktile_metadata(
            hidden_pad, intermediate_pad, use_g1u1, activation, split_k=2,
        )

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


fused_moe_mod.get_2stage_cfgs = _patched_get_2stage_cfgs


_DIRECT_BUFS = {}


def _get_split_k(E, M, inter_dim):
    if E == 257 and inter_dim == 256:
        if M == 16:
            return 2
        if M == 128:
            return 4
    if E == 33 and inter_dim == 512:
        if M == 16:
            return 2
        if M == 128:
            return 1
    return 0


def _direct_cktile_pipeline(
    hidden_states, w1, w2, w1_scale, w2_scale,
    topk_ids, topk_weights, E, M, topk, model_dim,
    inter_dim, hidden_pad, intermediate_pad, split_k,
):
    block_m = 16

    sid, sw, sei, nvi, mb = _cached_moe_sorting_impl(
        topk_ids, topk_weights, E, model_dim, torch.bfloat16,
        block_m, None, None, 0, True,
    )

    _, n1, k1 = w1.shape
    _, k2, n2 = w2.shape
    D = n2 if k2 == k1 else n2 * 2
    if w1.dtype is torch.uint32:
        D = D * 8

    n_pad1 = intermediate_pad // 64 * 64 * 2
    k_pad1 = hidden_pad // 128 * 128
    n_pad2 = hidden_pad // 64 * 64
    k_pad2 = intermediate_pad // 128 * 128

    buf_key = (M, topk, D, w1.shape[1], split_k)
    if buf_key not in _DIRECT_BUFS:
        dev = hidden_states.device
        out = torch.empty((M, topk, D), dtype=torch.bfloat16, device=dev)
        tmp = (
            torch.zeros((M, topk, w1.shape[1]), dtype=torch.bfloat16, device=dev)
            if split_k > 1 else None
        )
        _DIRECT_BUFS[buf_key] = (out, tmp)

    out_buf, tmp_buf = _DIRECT_BUFS[buf_key]

    w1_scale_e8m0 = w1_scale.view(dtypes.fp8_e8m0)
    w2_scale_e8m0 = w2_scale.view(dtypes.fp8_e8m0)

    if split_k > 1:
        tmp_buf.zero_()
        aiter.moe_cktile2stages_gemm1(
            hidden_states, w1, tmp_buf,
            sid, sei, nvi,
            topk, n_pad1, k_pad1,
            None, None, w1_scale_e8m0, None,
            ActivationType.Silu, block_m, split_k,
        )
        aiter.silu_and_mul(out_buf, tmp_buf)
    else:
        aiter.moe_cktile2stages_gemm1(
            hidden_states, w1, out_buf,
            sid, sei, nvi,
            topk, n_pad1, k_pad1,
            None, None, w1_scale_e8m0, None,
            ActivationType.Silu, block_m, 1,
        )

    aiter.moe_cktile2stages_gemm2(
        out_buf, w2, mb,
        sid, sei, nvi,
        topk, n_pad2, k_pad2,
        sw, None, w2_scale_e8m0, None,
        ActivationType.Silu, block_m,
    )

    return mb


_A2_BUFS = {}

_S3_CK_K1 = (
    "moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_"
    "Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
)

_DIRECT_CK_CONFIGS = {
    (257, 256): (_S3_CK_K1, 32, 16, 256, 256, "reduce"),
    (33, 2048): (_k1_2048, 64, 16, 256, 256, "atomic"),
}


_QUANT_BUFS = {}


def _cached_quant(x, sorted_ids, num_valid_ids, token_num, topk, block_size=32):
    M, N = x.shape
    MXFP4_QUANT_BLOCK_SIZE = 32
    scaleN = triton.cdiv(N, MXFP4_QUANT_BLOCK_SIZE)
    BLOCK_SIZE_Mx = 128
    BLOCK_SIZE_M, BLOCK_SIZE_N = 32, 8
    BLOCK_SIZE_M_u32, BLOCK_SIZE_N_u32 = 16, 4
    M_o = sorted_ids.shape[0]
    N_o = scaleN

    key = (M, N, M_o)
    if key not in _QUANT_BUFS:
        x_fp4 = torch.empty((M, N // 2), dtype=torch.uint8, device=x.device)
        blockscale = torch.empty(
            (
                triton.cdiv(M_o, BLOCK_SIZE_M),
                triton.cdiv(N_o, BLOCK_SIZE_N),
                BLOCK_SIZE_N_u32,
                BLOCK_SIZE_M_u32,
                4,
            ),
            dtype=torch.uint8,
            device=x.device,
        )
        _QUANT_BUFS[key] = (
            x_fp4,
            blockscale,
            x_fp4.view(dtypes.fp4x2),
            blockscale.view(dtypes.fp8_e8m0).view(-1, N_o),
        )

    x_fp4, blockscale, x_fp4_view, blockscale_view = _QUANT_BUFS[key]

    M_i, N_i = M, scaleN
    num_pid = triton.cdiv(M, BLOCK_SIZE_Mx) * scaleN + triton.cdiv(
        M_o, BLOCK_SIZE_M
    ) * triton.cdiv(N_i, BLOCK_SIZE_N)
    _fused_dynamic_mxfp4_quant_moe_sort_kernel[(num_pid,)](
        x,
        x_fp4,
        sorted_ids,
        num_valid_ids,
        blockscale,
        M,
        N,
        scaleN,
        *x.stride(),
        *x_fp4.stride(),
        *blockscale.stride(),
        token_num=token_num,
        M_i=M_i,
        N_i=N_i,
        MXFP4_QUANT_BLOCK_SIZE=MXFP4_QUANT_BLOCK_SIZE,
        BLOCK_SIZE_Mx=BLOCK_SIZE_Mx,
        BLOCK_SIZE_M=BLOCK_SIZE_M // 2,
        BLOCK_SIZE_N=BLOCK_SIZE_N // 2,
        TOPK=topk,
    )

    return x_fp4_view, blockscale_view


def _direct_ck_flydsl_pipeline(
    hidden_states, w1, w2, w1_scale, w2_scale,
    topk_ids, topk_weights, E, M, topk, model_dim, inter_dim,
):
    ck_k1, block_m, fly_tm, fly_tn, fly_tk, fly_mode = _DIRECT_CK_CONFIGS[(E, inter_dim)]

    sid, sw, sei, nvi, mb = _cached_moe_sorting_impl(
        topk_ids, topk_weights, E, model_dim, torch.bfloat16,
        block_m, None, None, 0, True,
    )

    a1, a1_scale = _cached_quant(
        hidden_states, sorted_ids=sid, num_valid_ids=nvi,
        token_num=M, topk=1, block_size=block_m,
    )

    buf_key = (M, topk, inter_dim)
    if buf_key not in _A2_BUFS:
        _A2_BUFS[buf_key] = torch.empty(
            (M, topk, inter_dim), dtype=torch.bfloat16, device=hidden_states.device,
        )
    a2 = _A2_BUFS[buf_key]

    w1_scale_e8m0 = w1_scale.view(dtypes.fp8_e8m0)
    aiter.ck_moe_stage1_fwd(
        a1, w1, w2, sid, sei, nvi, a2, topk,
        ck_k1, w1_scale_e8m0, a1_scale, block_m,
        None, QuantType.per_1x32, ActivationType.Silu, 0, True,
        torch.bfloat16,
    )

    a2_flat = a2.view(-1, inter_dim)
    a2_quant, a2_scale = _cached_quant(
        a2_flat, sorted_ids=sid, num_valid_ids=nvi,
        token_num=M, topk=topk, block_size=block_m,
    )
    a2_quant = a2_quant.view(M, topk, -1)

    w2_scale_e8m0 = w2_scale.view(dtypes.fp8_e8m0)
    flydsl_moe_stage2(
        inter_states=a2_quant, w2=w2,
        sorted_token_ids=sid, sorted_expert_ids=sei,
        num_valid_ids=nvi, out=mb, topk=topk,
        tile_m=fly_tm, tile_n=fly_tn, tile_k=fly_tk,
        a_dtype="fp4", b_dtype="fp4", out_dtype="bf16",
        mode=fly_mode,
        w2_scale=w2_scale_e8m0, a2_scale=a2_scale,
        sorted_weights=sw,
    )

    return mb


_FLYDSL_S1_CONFIGS = {
    (257, 256): (32, 128, 256, 256, 64, 256, 256),
    (33, 512): (32, 128, 256, 256, 16, 256, 256),
}


def _direct_flydsl_pipeline(
    hidden_states, w1, w2, w1_scale, w2_scale,
    topk_ids, topk_weights, E, M, topk, model_dim, inter_dim,
):
    block_m, fly_tm, fly_tn, fly_tk, s2_tm, s2_tn, s2_tk = _FLYDSL_S1_CONFIGS[(E, inter_dim)]
    fly_mode = "reduce"

    sid, sw, sei, nvi, mb = _cached_moe_sorting_impl(
        topk_ids, topk_weights, E, model_dim, torch.bfloat16,
        block_m, None, None, 0, True,
    )

    a1, a1_scale = _cached_quant(
        hidden_states, sorted_ids=sid, num_valid_ids=nvi,
        token_num=M, topk=1, block_size=block_m,
    )

    buf_key = (M, topk, inter_dim, "fly")
    if buf_key not in _A2_BUFS:
        _A2_BUFS[buf_key] = torch.empty(
            (M, topk, inter_dim), dtype=torch.bfloat16, device=hidden_states.device,
        )
    a2 = _A2_BUFS[buf_key]

    w1_scale_e8m0 = w1_scale.view(dtypes.fp8_e8m0)
    flydsl_moe_stage1(
        a=a1, w1=w1,
        sorted_token_ids=sid, sorted_expert_ids=sei,
        num_valid_ids=nvi, out=a2, topk=topk,
        tile_m=fly_tm, tile_n=fly_tn, tile_k=fly_tk,
        a_dtype="fp4", b_dtype="fp4", out_dtype="bf16",
        w1_scale=w1_scale_e8m0, a1_scale=a1_scale,
        sorted_weights=None,
    )

    a2_flat = a2.view(-1, inter_dim)
    a2_quant, a2_scale = _cached_quant(
        a2_flat, sorted_ids=sid, num_valid_ids=nvi,
        token_num=M, topk=topk, block_size=block_m,
    )
    a2_quant = a2_quant.view(M, topk, -1)

    w2_scale_e8m0 = w2_scale.view(dtypes.fp8_e8m0)
    flydsl_moe_stage2(
        inter_states=a2_quant, w2=w2,
        sorted_token_ids=sid, sorted_expert_ids=sei,
        num_valid_ids=nvi, out=mb, topk=topk,
        tile_m=s2_tm, tile_n=s2_tn, tile_k=s2_tk,
        a_dtype="fp4", b_dtype="fp4", out_dtype="bf16",
        mode=fly_mode,
        w2_scale=w2_scale_e8m0, a2_scale=a2_scale,
        sorted_weights=sw,
    )

    return mb


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"]

    M = hidden_states.shape[0]
    E = gate_up_weight.shape[0]
    inter_dim = config["d_expert"]
    model_dim = hidden_states.shape[1]
    topk = topk_ids.shape[1]

    split_k = _get_split_k(E, M, inter_dim)

    if split_k > 0 and model_dim == 7168 and topk == 9:
        return _direct_cktile_pipeline(
            hidden_states,
            gate_up_weight_shuffled, down_weight_shuffled,
            gate_up_weight_scale_shuffled, down_weight_scale_shuffled,
            topk_ids, topk_weights, E, M, topk, model_dim,
            inter_dim, hidden_pad, intermediate_pad, split_k,
        )

    if M == 512 and model_dim == 7168 and topk == 9 and (E, inter_dim) in _FLYDSL_S1_CONFIGS:
        return _direct_flydsl_pipeline(
            hidden_states,
            gate_up_weight_shuffled, down_weight_shuffled,
            gate_up_weight_scale_shuffled, down_weight_scale_shuffled,
            topk_ids, topk_weights, E, M, topk, model_dim, inter_dim,
        )

    if M == 512 and model_dim == 7168 and topk == 9 and (E, inter_dim) in _DIRECT_CK_CONFIGS:
        return _direct_ck_flydsl_pipeline(
            hidden_states,
            gate_up_weight_shuffled, down_weight_shuffled,
            gate_up_weight_scale_shuffled, down_weight_scale_shuffled,
            topk_ids, topk_weights, E, M, topk, model_dim, inter_dim,
        )

    return fused_moe_mod.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,
    )