submission 643890

# MoE v36 — Try 1WG_M32 stage1 for E=33 bs=16 (sparse)
#
# v30 uses default CKTile for E=33 bs=16 (ksplit=2).
# v34 tries 1WG_M16 for E=257. v36 tries 1WG_M32 for E=33 bs=16.
#
# E=33 bs=16: 16 tokens / ~9 experts = ~4.4 tokens per expert
# With block_m=32, most tiles have 1-2 valid tokens (wasteful).
# A 1WG stage1 kernel with M32 has lower launch overhead than 4WG,
# which could help for this very sparse case.
#
# Also try: block_m=16 for E=33 bs=16 (instead of 32)
# With 4.4 tokens/expert, block_m=16 wastes less than block_m=32.
#
# Test:       popcorn submit --gpu MI355X --leaderboard moe-mxfp4 --mode test moe_v36.py
# Benchmark:  popcorn submit --gpu MI355X --leaderboard moe-mxfp4 --mode benchmark moe_v36.py

import os
import sys
import functools
import torch
from task import input_t, output_t
from aiter import ActivationType, QuantType
from aiter.fused_moe import fused_moe
import aiter.fused_moe as _fused_moe_module

# ── FlyDSL registration (identical to v19) ──
try:
    import aiter.ops.flydsl.moe_kernels as _flydsl_moe_kernels
    for tile_m, tile_n in [(32, 128), (32, 256), (16, 256), (16, 128), (64, 128), (64, 256)]:
        name = f"flydsl_moe2_afp4_wfp4_bf16_t{tile_m}x{tile_n}x128_atomic"
        _flydsl_moe_kernels._KERNEL_PARAMS[name] = {
            "stage": 2, "a_dtype": "fp4", "b_dtype": "fp4", "out_dtype": "bf16",
            "tile_m": tile_m, "tile_n": tile_n, "tile_k": 128,
            "mode": "atomic", "MPerBlock": tile_m,
        }
except ImportError:
    pass

def _key(token, inter_dim, expert):
    return (
        256, token, 7168, inter_dim, expert, 9,
        "ActivationType.Silu", "torch.bfloat16",
        "torch.float4_e2m1fn_x2", "torch.float4_e2m1fn_x2",
        "QuantType.per_1x32", True, False,
    )

# ── Kernel names ──
_4WG_M128 = "moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
_4WG_M64  = "moe_ck2stages_gemm1_256x64x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
_4WG_M32  = "moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
_1WG_M32  = "moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
_1WG_M16  = "moe_ck2stages_gemm1_256x16x128x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
_FLY_16x128 = "flydsl_moe2_afp4_wfp4_bf16_t16x128x128_atomic"

_CUSTOM_CONFIGS = {}

# ═══════════════════════════════════════════════════════════════
# E=33 shapes
# ═══════════════════════════════════════════════════════════════

# bs=16: CHANGED — block_m=16 + ksplit=2 (smaller tiles for 4.4 tokens/expert)
_CUSTOM_CONFIGS[_key(16, 512, 33)] = {
    "block_m": 16, "ksplit": 2,  # CHANGED: block_m 32→16
    "kernelName1": "", "kernelName2": "",
    "run_1stage": False,
    "use_non_temporal_load": True,  # sparse — bypass cache
}

# bs=128: CHANGED from v19 — block_m=32 instead of 64 (v29: 91.4→89.6, 2% win)
_CUSTOM_CONFIGS[_key(128, 512, 33)] = {
    "block_m": 32, "ksplit": 0,  # CHANGED: block_m 64→32
    "kernelName1": _4WG_M128, "kernelName2": _FLY_16x128,
    "run_1stage": False,
}

# bs=512 d=512: EXACT v19
_CUSTOM_CONFIGS[_key(512, 512, 33)] = {
    "block_m": 64, "ksplit": 0,
    "kernelName1": _4WG_M128, "kernelName2": _FLY_16x128,
    "run_1stage": False,
}

# bs=512 d=2048: EXACT v19 (FLY_64x128 was 24% worse in v29)
_CUSTOM_CONFIGS[_key(512, 2048, 33)] = {
    "block_m": 64, "ksplit": 0,
    "kernelName1": _4WG_M128, "kernelName2": _FLY_16x128,
    "run_1stage": False,
}

# ═══════════════════════════════════════════════════════════════
# E=257 shapes
# ═══════════════════════════════════════════════════════════════

# bs=16: EXACT v19
_CUSTOM_CONFIGS[_key(16, 256, 257)] = {
    "block_m": 16, "ksplit": 2,
    "kernelName1": "", "kernelName2": "",
    "run_1stage": False,
    "use_non_temporal_load": True,
}

# bs=128: EXACT v19
_CUSTOM_CONFIGS[_key(128, 256, 257)] = {
    "block_m": 16, "ksplit": 2,
    "kernelName1": "", "kernelName2": "",
    "run_1stage": False,
    "use_non_temporal_load": True,
}

# bs=512: CHANGED from v19 — 4WG_M64 instead of 4WG_M32 (v29: 208→180, 13% win!)
_CUSTOM_CONFIGS[_key(512, 256, 257)] = {
    "block_m": 32, "ksplit": 0,
    "kernelName1": _4WG_M64,  # CHANGED: M32→M64
    "kernelName2": _FLY_16x128,
    "run_1stage": False,
    "use_non_temporal_load": True,
}

# ── Config injection (identical to v19) ──
_injected = False
def _inject_configs():
    global _injected
    if _injected:
        return
    _injected = True
    if _fused_moe_module.cfg_2stages is None:
        import pandas as pd
        from aiter.jit.core import AITER_CONFIGS
        tune_file = AITER_CONFIGS.AITER_CONFIG_FMOE_FILE
        if os.path.exists(tune_file):
            _INDEX_COLS = ["cu_num", "token", "model_dim", "inter_dim", "expert", "topk",
                           "act_type", "dtype", "q_dtype_a", "q_dtype_w", "q_type", "use_g1u1", "doweight_stage1"]
            df = pd.read_csv(tune_file)
            if "_tag" in df.columns:
                df = df[df["_tag"].fillna("") == ""]
            _fused_moe_module.cfg_2stages = df.set_index(_INDEX_COLS).to_dict("index")
        else:
            _fused_moe_module.cfg_2stages = {}
    _fused_moe_module.cfg_2stages.update(_CUSTOM_CONFIGS)

    _original = _fused_moe_module.get_2stage_cfgs
    @functools.lru_cache(maxsize=2048)
    def _patched(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 = _original(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)
        from aiter.jit.utils.chip_info import get_cu_num
        cu_num = get_cu_num()
        keys = (cu_num, token, model_dim, inter_dim, expert, topk,
                str(activation), str(dtype), str(q_dtype_a), str(q_dtype_w),
                str(q_type), use_g1u1, doweight_stage1)
        cfg = _fused_moe_module.cfg_2stages.get(keys)
        if cfg and cfg.get("use_non_temporal_load") is not None:
            nt = cfg["use_non_temporal_load"]
            old_s1 = metadata.stage1
            kw = getattr(old_s1, 'keywords', None) or {}
            if hasattr(old_s1, 'func') and 'non_temporal_load' in kw:
                new_kw = dict(kw)
                new_kw['non_temporal_load'] = nt
                metadata = _fused_moe_module.MOEMetadata(
                    functools.partial(old_s1.func, **new_kw),
                    metadata.stage2, metadata.block_m, metadata.ksplit,
                    metadata.run_1stage, metadata.has_bias, nt)
        return metadata
    _fused_moe_module.get_2stage_cfgs = _patched


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

    _inject_configs()

    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