submission 654246

"""
test_v117_selective_splitk: Selective split_k — keep 2 for E=257 (s2), force 1 for E=33 (s5)
Base: test.py (v109)
Direction: CONTINUING from v116 (attempt 2) — selective split_k
Target: s2 (bs=128, E=257) — split_k=2 saves 5μs; s5 (bs=128, E=33) — split_k=1 avoids regression
Change: Make split_k override expert-aware: force split_k=1 only when token>16 AND
        n_experts<=64 (E=33 shapes). Let E=257 shapes use split_k=2 from get_ksplit.
Rationale: v116 showed split_k=2 helps s2 (-5μs, -2.8%) but hurts s5 (+10.6μs, +10.9%).
Scale: MODERATE
"""
from task import input_t, output_t

import os
os.environ["AITER_USE_OPUS_MOE_SORTING"] = "1"
# NOT setting AITER_USE_NT — we patch use_nt() directly for per-shape control

import functools
import torch
import aiter
from aiter import ActivationType, QuantType, dtypes
from aiter.fused_moe import fused_moe, get_2stage_cfgs, cktile_moe_stage1
import aiter.fused_moe as fmoe_module

# Patch use_nt: force NT=True for bs=512 (token > 128), default heuristic otherwise
def _patched_use_nt(token, topk, e):
    if token > 128:
        # bs=512 shapes: force NT on (benefits s3, s6, s7)
        return True
    # bs<=128 shapes: use default heuristic (estimated_m_per_expert < 64)
    return (token * topk // e) < 64

fmoe_module.use_nt = _patched_use_nt

# Selective split_k: force split_k=1 for E=33 bs=128 (dense), let E=257 use split_k=2 (sparse)
_orig_cktile_stage1 = cktile_moe_stage1

def _patched_cktile_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]
    n_experts = w1.shape[0]
    # Force split_k=1 for E=33 shapes with bs>16 (dense, ~34 tokens/expert)
    # Let E=257 shapes keep split_k=2 (sparse, ~0.5 tokens/expert)
    if token_num > 16 and n_experts <= 64:
        actual_split_k = 1
    else:
        actual_split_k = split_k
    return _orig_cktile_stage1(
        hidden_states, w1, w2, sorted_token_ids, sorted_expert_ids,
        num_valid_ids, out, topk, block_m, a1_scale, w1_scale,
        sorted_weights=sorted_weights, n_pad_zeros=n_pad_zeros,
        k_pad_zeros=k_pad_zeros, bias1=bias1,
        activation=activation, split_k=actual_split_k, dtype=dtype,
    )

fmoe_module.cktile_moe_stage1 = _patched_cktile_stage1

@functools.lru_cache(maxsize=2048)
def _patched_get_ksplit(token, topk, expert, inter_dim, model_dim):
    # cktile path for bs<=128
    if token <= 128:
        if model_dim % 2 == 0 and (model_dim // 2) % 256 == 0:
            return 2
    return 0

fmoe_module.get_ksplit = _patched_get_ksplit

# Tuned CK stage1 kernel for s7 (256x128 tile, block_m=128)
_STAGE1_256x128 = 'moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16'

_injected = False

def _inject_configs():
    global _injected
    if _injected:
        return
    _injected = True

    try:
        get_2stage_cfgs(
            16, 7168, 256, 257, 9,
            dtypes.bf16, dtypes.fp4x2, dtypes.fp4x2,
            QuantType.per_1x32, True, ActivationType.Silu, False,
            0, 0, True,
        )
    except Exception:
        pass

    if fmoe_module.cfg_2stages is None:
        return

    common = (
        'ActivationType.Silu', 'torch.bfloat16',
        'torch.float4_e2m1fn_x2', 'torch.float4_e2m1fn_x2',
        'QuantType.per_1x32', True, False,
    )

    # s3: FlyDSL stage2 reduce
    key = (256, 512, 7168, 256, 257, 9) + common
    if key in fmoe_module.cfg_2stages:
        fmoe_module.cfg_2stages[key]['kernelName2'] = 'flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce'

    # s6: FlyDSL stage2 reduce (original config — empty kernelName1, block_m=64)
    key = (256, 512, 7168, 512, 33, 9) + common
    if key not in fmoe_module.cfg_2stages:
        fmoe_module.cfg_2stages[key] = {
            'block_m': 64, 'ksplit': 0, 'kernelName1': '',
            'kernelName2': 'flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce',
            'run_1stage': False,
        }

    # s7: 256x128 stage1 tile (block_m=128) + FlyDSL t32 stage2
    key_s7 = (256, 512, 7168, 2048, 33, 9) + common
    fmoe_module.cfg_2stages[key_s7] = {
        'block_m': 128, 'ksplit': 0,
        'kernelName1': _STAGE1_256x128,
        'kernelName2': 'flydsl_moe2_afp4_wfp4_bf16_t32x256x256_reduce',
        'run_1stage': False,
    }

    # Remove s1 and s2 CSV entries to let ksplit=2 cktile path kick in
    key_s1 = (256, 16, 7168, 256, 257, 9) + common
    if key_s1 in fmoe_module.cfg_2stages:
        del fmoe_module.cfg_2stages[key_s1]

    key_s2 = (256, 128, 7168, 256, 257, 9) + common
    if key_s2 in fmoe_module.cfg_2stages:
        del fmoe_module.cfg_2stages[key_s2]

    get_2stage_cfgs.cache_clear()


def custom_kernel(data: input_t) -> output_t:
    _inject_configs()

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

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

    return output