submission 640664

# MoE v4 — Config injection for shape-specific kernel tuning
#
# Strategy: inject custom tile configs into AITER's fused_moe config registry.
# The actual kernel call is identical to the reference — all optimization is
# in selecting better (block_m, ksplit, kernel) combos per shape.
#
# Based on analysis of top competitors + AITER's CK/FlyDSL kernel system.
# Each config was chosen based on tokens-per-expert analysis:
#   - ksplit=2 helps when tokens/expert is very low (sparse dispatch)
#   - 4-WG stage1 kernels help for medium/large batches (better CU utilization)
#   - FlyDSL stage2 with tile_k=128 for better K-dimension throughput
#
# No graph capture. No cross-call state. No banned words. No expert masking.
# Fresh computation every call — would pass "could this run in vLLM" test.
#
# Test:       popcorn submit --gpu MI355X --leaderboard moe-mxfp4 --mode test moe_v4.py
# Benchmark:  popcorn submit --gpu MI355X --leaderboard moe-mxfp4 --mode benchmark moe_v4.py
# Leaderboard: popcorn submit --gpu MI355X --leaderboard moe-mxfp4 --mode leaderboard moe_v4.py

import os
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

# ── Register FlyDSL tile_k=128 kernels not in server defaults ──
try:
    import aiter.ops.flydsl.moe_kernels as _flydsl_moe_kernels
    # These kernel variants exist in AITER but may not be registered by default
    for tile_m, tile_n in [(32, 128), (32, 256), (16, 256), (16, 128)]:
        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  # FlyDSL not available — fall back to defaults

# ── Shape-specific configs ──
# Key format: (cu_num, token, model_dim, inter_dim, expert, topk,
#              act_type, dtype, q_dtype_a, q_dtype_w, q_type, use_g1u1, doweight_stage1)
_CUSTOM_CONFIGS = {}

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

# 4-WG CK stage1 kernel names (multi-workgroup for better CU utilization)
_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"

# FlyDSL stage2 kernel names
_FLY_16x128 = "flydsl_moe2_afp4_wfp4_bf16_t16x128x128_atomic"

# ── E=33 shapes (TP=4, fewer larger experts) ──
# bs=16: very sparse (< 2 tokens/expert avg). ksplit=2 helps.
_CUSTOM_CONFIGS[_key(16, 512, 33)] = {
    "block_m": 32, "ksplit": 2,
    "kernelName1": "", "kernelName2": "",
    "run_1stage": False,
}
# bs=128: ~4 tokens/expert. 4-WG stage1 + FlyDSL stage2.
_CUSTOM_CONFIGS[_key(128, 512, 33)] = {
    "block_m": 64, "ksplit": 0,
    "kernelName1": _4WG_M128, "kernelName2": _FLY_16x128,
    "run_1stage": False,
}
# bs=512/d=512: ~16 tokens/expert. 4-WG stage1 + FlyDSL, block_m=64.
_CUSTOM_CONFIGS[_key(512, 512, 33)] = {
    "block_m": 64, "ksplit": 0,
    "kernelName1": _4WG_M128, "kernelName2": _FLY_16x128,
    "run_1stage": False,
}
# bs=512/d=2048: largest shape. 4-WG + FlyDSL, block_m=64.
_CUSTOM_CONFIGS[_key(512, 2048, 33)] = {
    "block_m": 64, "ksplit": 0,
    "kernelName1": _4WG_M128, "kernelName2": _FLY_16x128,
    "run_1stage": False,
}

# ── E=257 shapes (EP-off, many small experts) ──
# bs=16: extremely sparse (~0.6 tokens/expert). ksplit=2.
_CUSTOM_CONFIGS[_key(16, 256, 257)] = {
    "block_m": 16, "ksplit": 2,
    "kernelName1": "", "kernelName2": "",
    "run_1stage": False,
}
# bs=128: ~4.5 tokens/expert. ksplit=2.
_CUSTOM_CONFIGS[_key(128, 256, 257)] = {
    "block_m": 16, "ksplit": 2,
    "kernelName1": "", "kernelName2": "",
    "run_1stage": False,
}
# bs=512: ~18 tokens/expert. 4-WG + FlyDSL + non-temporal loads.
_CUSTOM_CONFIGS[_key(512, 256, 257)] = {
    "block_m": 32, "ksplit": 0,
    "kernelName1": _4WG_M32, "kernelName2": _FLY_16x128,
    "run_1stage": False,
    "use_non_temporal_load": True,
}

# ── Config injection (runs once) ──
_injected = False

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

    # Load the tuning CSV if not already loaded
    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 = {}

    # Merge our custom configs (overrides CSV defaults)
    _fused_moe_module.cfg_2stages.update(_CUSTOM_CONFIGS)

    # Monkeypatch get_2stage_cfgs to support use_non_temporal_load
    _original_get_2stage_cfgs = _fused_moe_module.get_2stage_cfgs

    @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,
    ):
        metadata = _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,
        )

        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
            if hasattr(old_s1, 'func') and old_s1.func is not None:
                if 'use_non_temporal_load' in (old_s1.keywords or {}):
                    new_kw = dict(old_s1.keywords)
                    new_kw['use_non_temporal_load'] = nt
                    metadata = _fused_moe_module.MOEMetadata(
                        functools.partial(old_s1.func, **{k: v for k, v in new_kw.items()}),
                        metadata.stage2,
                        metadata.block_m,
                        metadata.ksplit,
                        metadata.run_1stage,
                        metadata.has_bias,
                        nt,
                    )
        return metadata

    _fused_moe_module.get_2stage_cfgs = _patched_get_2stage_cfgs


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