submission 731681

"""
v52: ksplit=4 for E=257 shapes (both bs=16 and bs=128) and E=33 bs=16.
     ksplit=2 for E=33 d=512 bs=128 (ksplit=4 regressed +6µs in v51).

v51 experiment revealed:
- ksplit=4 helps E=257 bs=128: 190µs → 173µs (-17µs!)
- ksplit=4 helps E=257 bs=16: 96µs → 91.2µs (-4.8µs)
- ksplit=4 helps E=33 d=512 bs=16: 64µs → 61.0µs (-3µs)
- ksplit=4 hurts E=33 d=512 bs=128: 110µs → 116µs (+6µs) ← revert

Expected geomean: ~150µs (vs v39 ~154µs, v51 ~151.6µs).

Key format: (cu_num=256, token, hidden=7168, inter, n_experts, topk, act, dtype,
             q_dtype_a, q_dtype_w, q_type, use_g1u1=True, doweight_stage1=False)
"""

import torch
from task import input_t, output_t

import aiter.fused_moe as _fmoe_mod
from aiter import ActivationType, QuantType
from aiter.fused_moe import fused_moe, get_2stage_cfgs

# ── Kernel name constants ──────────────────────────────────────────────────────
# gemm1 kernels (all FP4X2 a4w4, silu, bf16 output)
_K1_32  = "moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
_K1_32w = "moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
_K1_64  = "moe_ck2stages_gemm1_256x64x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
_K1_128 = "moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"

# gemm2 kernels
_K2_32  = "moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
_K2_64  = "moe_ck2stages_gemm2_256x64x128x128_1x4_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
_K2_64s = "moe_ck2stages_gemm2_64x64x128x128_1x1_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
_K2_128 = "moe_ck2stages_gemm2_256x128x128x128_1x4_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"

# ── Config injection ───────────────────────────────────────────────────────────
_ACT = "ActivationType.Silu"
_DT  = "torch.bfloat16"
_QDA = "torch.float4_e2m1fn_x2"
_QDW = "torch.float4_e2m1fn_x2"
_QT  = "QuantType.per_1x32"


def _inject_all_configs():
    entries = {}

    # E=257 d=256 bs=16,128: cktile ksplit=4 (v51 confirmed: -4.8µs bs=16, -17µs bs=128)
    for bs in (16, 128):
        entries[(256, bs, 7168, 256, 257, 9, _ACT, _DT, _QDA, _QDW, _QT, True, False)] = {
            "block_m": 16, "ksplit": 4, "kernelName1": _K1_32, "kernelName2": _K2_32, "run_1stage": False,
        }
    # E=257 d=256 bs=512: CK path (cktile regressed +22µs in v35)
    entries[(256, 512, 7168, 256, 257, 9, _ACT, _DT, _QDA, _QDW, _QT, True, False)] = {
        "block_m": 32, "ksplit": 0, "kernelName1": _K1_32, "kernelName2": _K2_32, "run_1stage": False,
    }

    # E=33 d=512, bs=16: cktile ksplit=4 (v51 confirmed: -3µs)
    entries[(256, 16, 7168, 512, 33, 9, _ACT, _DT, _QDA, _QDW, _QT, True, False)] = {
        "block_m": 16, "ksplit": 4, "kernelName1": _K1_32, "kernelName2": _K2_32, "run_1stage": False,
    }
    # E=33 d=512, bs=128: cktile ksplit=2 (ksplit=4 regressed +6µs in v51)
    entries[(256, 128, 7168, 512, 33, 9, _ACT, _DT, _QDA, _QDW, _QT, True, False)] = {
        "block_m": 16, "ksplit": 2, "kernelName1": _K1_32, "kernelName2": _K2_32, "run_1stage": False,
    }
    # E=33 d=512, bs=512: bm=64 with 64-thread sparse gemm2 (-21µs vs bm=128)
    entries[(256, 512, 7168, 512, 33, 9, _ACT, _DT, _QDA, _QDW, _QT, True, False)] = {
        "block_m": 64, "ksplit": 0, "kernelName1": _K1_64, "kernelName2": _K2_64s, "run_1stage": False,
    }

    # E=33 d=2048, bs=512: bm=64, 256-thread gemm2
    entries[(256, 512, 7168, 2048, 33, 9, _ACT, _DT, _QDA, _QDW, _QT, True, False)] = {
        "block_m": 64, "ksplit": 0, "kernelName1": _K1_64, "kernelName2": _K2_64, "run_1stage": False,
    }

    if _fmoe_mod.cfg_2stages is None:
        _fmoe_mod.cfg_2stages = {}
    _fmoe_mod.cfg_2stages.update(entries)
    get_2stage_cfgs.cache_clear()


_inject_all_configs()


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
    _ = gate_up_weight, down_weight, gate_up_weight_scale, down_weight_scale
    # block_m=32 helps E=33d512 bs=128 (K_gemm2=512 > KPerBlock=256 → 2 k-iters)
    # block_m=32 hurts E=257 bs=128 (cktile: +25µs; CK: +6µs — both worse with bm=32)
    # block_m=32 hurts E=33d512 bs=16 (bm=32 doubles gemm1 CTA time when <1 wave)
    block_size_M = 32 if (hidden_states.shape[0] == 128 and config.get('d_expert', 0) == 512) else None
    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=block_size_M,
    )