submission 723222

#!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
from aiter.fused_moe import fused_moe
import aiter.fused_moe as _fmoe_mod

# CK kernel names
_KN1_LARGE = "moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
_KN1_MED = "moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
_KN2_SMALL = "moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"

_injected = False


def _inject_configs():
    """Inject overrides into cfg_2stages and clear lru_cache."""
    global _injected
    if _injected:
        return
    _injected = True

    cfg_dict = _fmoe_mod.cfg_2stages
    if cfg_dict is None:
        return

    common = (
        "ActivationType.Silu", "torch.bfloat16",
        "torch.float4_e2m1fn_x2", "torch.float4_e2m1fn_x2",
        "QuantType.per_1x32", True, False,
    )

    # E=257 bs=16: CKTile ksplit=4 → split-K path with block_m=16
    # Colleague: ksplit=4 → 89.6µs vs ksplit=2 → 91µs
    cfg_dict[(256, 16, 7168, 256, 257, 9) + common] = {
        "block_m": 16, "ksplit": 4,
        "kernelName1": "", "kernelName2": "", "run_1stage": 0,
    }

    # E=33 bs=16/d=512: CKTile ksplit=2 (test if it helps small batch E=33)
    cfg_dict[(256, 16, 7168, 512, 33, 9) + common] = {
        "block_m": 16, "ksplit": 2,
        "kernelName1": "", "kernelName2": "", "run_1stage": 0,
    }

    # E=257 bs=128/512: LARGE+SMALL block_m=32
    e257 = {"block_m": 32, "ksplit": 0, "kernelName1": _KN1_LARGE,
            "kernelName2": _KN2_SMALL, "run_1stage": 0}
    cfg_dict[(256, 128, 7168, 256, 257, 9) + common] = e257.copy()
    cfg_dict[(256, 512, 7168, 256, 257, 9) + common] = e257.copy()

    # E=33 d=512 bs=128/512: MED+SMALL block_m=32
    e33 = {"block_m": 32, "ksplit": 0, "kernelName1": _KN1_MED,
           "kernelName2": _KN2_SMALL, "run_1stage": 0}
    cfg_dict[(256, 128, 7168, 512, 33, 9) + common] = e33.copy()
    cfg_dict[(256, 512, 7168, 512, 33, 9) + common] = e33.copy()

    # E=33 d=2048 bs=512: block_m=64 auto
    cfg_dict[(256, 512, 7168, 2048, 33, 9) + common] = {
        "block_m": 64, "ksplit": 0,
        "kernelName1": "", "kernelName2": "", "run_1stage": 0,
    }

    # CRITICAL: Clear lru_cache so ALL shapes (including already-cached ones)
    # pick up the new configs on next call
    _fmoe_mod.get_2stage_cfgs.cache_clear()


_call_count = 0


def custom_kernel(data: input_t) -> output_t:
    global _call_count

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

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

    # Inject after first call (once cfg_2stages is loaded by fused_moe)
    # cache_clear ensures ALL shapes use new configs on subsequent calls
    _call_count += 1
    if _call_count == 1:
        _inject_configs()

    return output