submission 712180

#!POPCORN leaderboard amd-moe-mxfp4
#!POPCORN gpu MI355X

import torch
from typing import Dict
from task import input_t, output_t

from aiter import ActivationType, QuantType, dtypes
from aiter.fused_moe import fused_moe
import aiter.fused_moe as _fm

# ── EXP-202: FlyDSL stage1/stage2 (best perf) + hipcc+ctypes infrastructure ──

import subprocess, ctypes, tempfile, pathlib, base64 as _b64

# Base64-encoded identity kernel (placeholder for future MFMA kernel)
_HIP_KERNEL = _b64.b64decode(
    "CiNpbmNsdWRlIDxoaXAvaGlwX3J1bnRpbWUuaD4KI2luY2x1ZGUgPGhpcC9oaXBfYmYxNi5oPgoK"
    "ZXh0ZXJuICJDIiBfX2dsb2JhbF9fIHZvaWQgYmYxNl9pZGVudGl0eShjb25zdCBfX2hpcF9iZmxv"
    "YXQxNiogX19yZXN0cmljdF9fIHNyYywKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAg"
    "ICAgICAgICAgX19oaXBfYmZsb2F0MTYqIF9fcmVzdHJpY3RfXyBkc3QsCiAgICAgICAgICAgICAg"
    "ICAgICAgICAgICAgICAgICAgICAgICAgICAgIGludCBOKSB7CiAgICBpbnQgaWR4ID0gYmxvY2tJ"
    "ZHgueCAqIGJsb2NrRGltLnggKyB0aHJlYWRJZHgueDsKICAgIGlmIChpZHggPCBOKSBkc3RbaWR4"
    "XSA9IHNyY1tpZHhdOwp9CgpleHRlcm4gIkMiIHZvaWQgbGF1bmNoX2JmMTZfY29weSh2b2lkKiBz"
    "cmNfcHRyLCB2b2lkKiBkc3RfcHRyLCBpbnQgTiwgdm9pZCogc3RybSkgewogICAgY29uc3QgaW50"
    "IHRocmVhZHMgPSAyNTY7CiAgICBjb25zdCBpbnQgYmxvY2tzID0gKE4gKyB0aHJlYWRzIC0gMSkg"
    "LyB0aHJlYWRzOwogICAgYmYxNl9pZGVudGl0eTw8PGJsb2NrcywgdGhyZWFkcywgMCwgKGhpcFN0"
    "cmVhbV90KXN0cm0+Pj4oCiAgICAgICAgcmVpbnRlcnByZXRfY2FzdDxjb25zdCBfX2hpcF9iZmxv"
    "YXQxNio+KHNyY19wdHIpLAogICAgICAgIHJlaW50ZXJwcmV0X2Nhc3Q8X19oaXBfYmZsb2F0MTYq"
    "Pihkc3RfcHRyKSwKICAgICAgICBOCiAgICApOwp9Cg=="
).decode()

_hip_lib = None
_hip_fn = None
try:
    _tmpdir = tempfile.mkdtemp(prefix="hip_e202_")
    _src = pathlib.Path(_tmpdir) / "kernel.hip"
    _so = pathlib.Path(_tmpdir) / "kernel.so"
    _src.write_text(_HIP_KERNEL)
    _t0 = __import__("time").time()
    _proc = subprocess.run(
        ["hipcc", "-shared", "-fPIC", "--offload-arch=gfx950", "-O3",
         "-o", str(_so), str(_src)],
        capture_output=True, text=True, timeout=120,
    )
    _dt = __import__("time").time() - _t0
    if _proc.returncode == 0 and _so.exists():
        _hip_lib = ctypes.CDLL(str(_so))
        _hip_fn = _hip_lib.launch_bf16_copy
        _hip_fn.restype = None
        _hip_fn.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p]
        print(f"[EXP-202] hipcc compiled OK in {_dt:.1f}s", flush=True)
    else:
        print(f"[EXP-202] hipcc FAILED: {_proc.stderr[:500]}", flush=True)
except Exception as _e:
    print(f"[EXP-202] compile error: {_e}", flush=True)

import os, functools

def _patch_flydsl_stage1():
    """Patch compute_f8f6f4_tile signature to fix compilation bug."""
    target = "/home/runner/aiter/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py"
    if not os.path.exists(target):
        return False
    with open(target) as f:
        content = f.read()
    old_sig = (
        "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,"
    )
    new_sig = (
        "b_tile_in,\n"
        "                        lds_base,\n"
        "                        *,\n"
        "                        a0_prefetch=None,\n"
        "                        a_scale=None,\n"
        "                        b_scale=None,"
    )
    if old_sig in content:
        content = content.replace(old_sig, new_sig, 1)
        with open(target, 'w') as f:
            f.write(content)
        print(f"[EXP-202] Patched compute_f8f6f4_tile signature", flush=True)
    else:
        print(f"[EXP-202] Signature already patched or different", flush=True)
    return True

_patch_flydsl_stage1()

# FlyDSL stage1 wrapper matching ck_moe_stage1 interface
def _flydsl_stage1_wrapper(
    hidden_states, w1, w2,
    sorted_token_ids, sorted_expert_ids, num_valid_ids,
    out, topk,
    block_m=64,
    a1_scale=None, w1_scale=None,
    kernelName='', sorted_weights=None,
    tile_m=64, tile_n=256, tile_k=128,
    **_kwargs,
):
    """Bridge between ck_moe_stage1 interface and flydsl_moe_stage1."""
    from aiter.ops.flydsl import flydsl_moe_stage1
    result = flydsl_moe_stage1(
        a=hidden_states,
        w1=w1,
        sorted_token_ids=sorted_token_ids,
        sorted_expert_ids=sorted_expert_ids,
        num_valid_ids=num_valid_ids,
        out=out,
        topk=topk,
        tile_m=tile_m,
        tile_n=tile_n,
        tile_k=tile_k,
        a_dtype='fp4',
        b_dtype='fp4',
        out_dtype='bf16',
        w1_scale=w1_scale,
        a1_scale=a1_scale,
        sorted_weights=sorted_weights,
    )
    return result


def _preload_and_inject():
    try:
        cu = _fm.get_cu_count() if hasattr(_fm, 'get_cu_count') else 256
        act = ActivationType.Silu
        dtype_t = torch.bfloat16
        q_a_t = dtypes.fp4x2
        q_w_t = dtypes.fp4x2
        q_type_t = QuantType.per_1x32

        _fm.get_2stage_cfgs(
            16, 7168, 256, 257, 9,
            dtype_t, q_a_t, q_w_t, q_type_t,
            True, act, False,
            0, 0, True
        )

        cfg = _fm.cfg_2stages
        if cfg is None:
            return

        act_s = str(act)
        dtype_s = str(dtype_t)
        q_a = str(q_a_t)
        q_w = str(q_w_t)
        q_type = str(q_type_t)

        kn1_small = 'moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16'
        kn2_small = 'moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16'
        kn1_32 = 'moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16'
        kn2_32 = 'moe_ck2stages_gemm2_256x32x128x128_1x4_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16'
        kn1_64 = 'moe_ck2stages_gemm1_256x64x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16'
        kn2_fly64_reduce = 'flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce'
        kn2_fly32_atomic_tn128 = 'flydsl_moe2_afp4_wfp4_bf16_t32x128x256_atomic'

        def make_entry(block_m, ksplit, kn1, kn2):
            return {
                'block_m': block_m, 'ksplit': ksplit,
                'kernelName1': kn1, 'kernelName2': kn2,
                'run_1stage': 0,
                'us': 0.0, 'us1': 0.0, 'us2': 0.0,
                'err1': '0', 'err2': '0',
                'tflops': 0.0, 'bw': 0.0, '_tag': float('nan'),
            }

        # Shape 1-3: E=257 decode shapes
        cfg[(cu, 16, 7168, 256, 257, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
            make_entry(32, 4, kn1_32, kn2_32)
        cfg[(cu, 128, 7168, 256, 257, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
            make_entry(32, 4, kn1_32, kn2_32)
        cfg[(cu, 512, 7168, 256, 257, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
            make_entry(32, 0, kn1_small, kn2_small)
        # Shape 4-5: E=33 decode shapes
        cfg[(cu, 16, 7168, 512, 33, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
            make_entry(32, 2, kn1_32, kn2_32)
        cfg[(cu, 128, 7168, 512, 33, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
            make_entry(32, 2, kn1_32, kn2_32)
        # Shape 6: CK stage1 + FlyDSL stage2 reduce
        cfg[(cu, 512, 7168, 512, 33, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
            make_entry(64, 0, kn1_64, kn2_fly64_reduce)
        # Shape 7: CK stage1 + FlyDSL stage2 atomic
        cfg[(cu, 512, 7168, 2048, 33, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
            make_entry(64, 0, kn1_64, kn2_fly32_atomic_tn128)

        if hasattr(_fm.get_2stage_cfgs, 'cache_clear'):
            _fm.get_2stage_cfgs.cache_clear()
        for fn_name in ['get_block_size_M', 'use_nt', 'get_ksplit']:
            fn = getattr(_fm, fn_name, None)
            if fn and hasattr(fn, 'cache_clear'):
                fn.cache_clear()

        # ── Monkeypatch get_2stage_cfgs to use FlyDSL stage1 for S6/S7 ──
        _monkeypatch_flydsl_stage1()

    except Exception as e:
        import traceback; traceback.print_exc()

def _monkeypatch_flydsl_stage1():
    """Replace metadata.stage1 with FlyDSL for target shapes (E=33, inter_dim>=512)."""
    try:
        orig_func = _fm.get_2stage_cfgs.__wrapped__
    except AttributeError:
        try:
            orig_func = _fm.get_2stage_cfgs
        except:
            return

    MOEMetadata = _fm.MOEMetadata if hasattr(_fm, 'MOEMetadata') else None

    @functools.lru_cache(maxsize=None)
    def _patched_get_2stage_cfgs(*args, **kwargs):
        result = orig_func(*args, **kwargs)
        if len(args) >= 5:
            token, model_dim, inter_dim, expert, topk = args[:5]
            # S6/S7: E=33, large batch, tile_m=64
            if expert == 33 and inter_dim >= 512 and token >= 256:
                tile_m, tile_n, tile_k = 64, 256, 128
                new_stage1 = functools.partial(
                    _flydsl_stage1_wrapper,
                    tile_m=tile_m, tile_n=tile_n, tile_k=tile_k,
                )
                if MOEMetadata is not None:
                    result = MOEMetadata(new_stage1, result.stage2, result.block_m, result.ksplit)
                else:
                    try:
                        result = type(result)(new_stage1, result.stage2, result.block_m, result.ksplit)
                    except:
                        pass
            # S2/S3: E=257, bs>=16, tile_m=32
            elif expert == 257 and token >= 16:
                tile_m, tile_n, tile_k = 32, 256, 128
                new_stage1 = functools.partial(
                    _flydsl_stage1_wrapper,
                    tile_m=tile_m, tile_n=tile_n, tile_k=tile_k,
                )
                if MOEMetadata is not None:
                    result = MOEMetadata(new_stage1, result.stage2, result.block_m, result.ksplit)
                else:
                    try:
                        result = type(result)(new_stage1, result.stage2, result.block_m, result.ksplit)
                    except:
                        pass
        return result

    _fm.get_2stage_cfgs = _patched_get_2stage_cfgs

_preload_and_inject()


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

    try:
        bsm = 32 if hidden_states.shape[0] == 128 and gate_up_weight_shuffled.shape[0] == 33 else None
    except Exception:
        bsm = 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,
        block_size_M=bsm,
        hidden_pad=hidden_pad, intermediate_pad=intermediate_pad,
    )