submission 656144

#!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-49: Import-time cfg injection + shape 7 block_m=64 ──
# Pre-load cfg_2stages by calling get_2stage_cfgs at import time,
# then inject our tuned entries before any benchmark call.
# Shapes 1,2,4,5: block_m=32, ksplit=4, 256x32 tile (EXP-32)
# Shape 7: block_m=64, ksplit=0, 256x64 tile (EXP-49)

def _preload_and_inject():
    """Trigger cfg_2stages loading and inject tuned configs at import time."""
    try:
        # Call get_2stage_cfgs with shape 1 args to trigger CSV loading
        # This only loads config (no JIT), so it's fast
        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

        # Trigger loading with shape 1 args (will load CSV + cache result)
        _fm.get_2stage_cfgs(
            16, 7168, 256, 257, 9,       # token, model_dim, inter_dim, expert, topk
            dtype_t, q_a_t, q_w_t, q_type_t,  # dtype, q_dtype_a, q_dtype_w, q_type
            True, act, False,              # use_g1u1, activation, doweight_stage1
            0, 0, True                     # hidden_pad, intermediate_pad, is_shuffled
        )

        # Now cfg_2stages is loaded. Inject our entries.
        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 = 'moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16'
        kn2 = '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_64 = 'moe_ck2stages_gemm2_256x64x128x128_1x4_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16'

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

        # Shapes 1,2 (E=257, bs=16/128) — CKTile ksplit=4
        cfg[(cu, 16, 7168, 256, 257, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
            make_entry(32, 4, kn1, kn2)
        cfg[(cu, 128, 7168, 256, 257, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
            make_entry(32, 4, kn1, kn2)

        # Shapes 4,5 (E=33, bs=16/128) — CKTile ksplit=4
        cfg[(cu, 16, 7168, 512, 33, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
            make_entry(32, 4, kn1, kn2)
        cfg[(cu, 128, 7168, 512, 33, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
            make_entry(32, 4, kn1, kn2)

        # Shapes 3,6,7: NOT injected — default CK path optimal
        # Except: Shape 7 (bs=512, E=33, d=2048) — try block_m=64 (default=128)
        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_64)

        # Clear ALL caches so next calls use our injected entries
        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()

    except Exception as e:
        pass

# Execute injection at import time
_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"]

    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=hidden_pad, intermediate_pad=intermediate_pad,
    )