submission 515387

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

"""Current best config mix plus separate HIP quant + scale sort for wide stage2 bs512 paths only."""

import csv
import importlib.util
import importlib
import os
from pathlib import Path

from task import input_t, output_t
from aiter import ActivationType, QuantType, dtypes, get_hip_quant
from aiter.fused_moe import fused_moe
from aiter.utility import fp4_utils


_FIELDNAMES = [
    "cu_num",
    "token",
    "model_dim",
    "inter_dim",
    "expert",
    "topk",
    "act_type",
    "dtype",
    "q_dtype_a",
    "q_dtype_w",
    "q_type",
    "use_g1u1",
    "doweight_stage1",
    "block_m",
    "ksplit",
    "us1",
    "kernelName1",
    "err1",
    "us2",
    "kernelName2",
    "err2",
    "us",
    "run_1stage",
    "tflops",
    "bw",
    "_tag",
]


def _write_row(
    writer: csv.DictWriter,
    *,
    cu_num: str,
    token: str,
    inter_dim: str,
    expert: str,
    ksplit: str,
    block_m: str,
    kernel1: str = "Null",
    kernel2: str = "Null",
    us1: str = "0.0",
    us2: str = "0.0",
    us: str = "0.0",
    err2: str = "0.0%",
) -> None:
    writer.writerow(
        {
            "cu_num": cu_num,
            "token": token,
            "model_dim": "7168",
            "inter_dim": inter_dim,
            "expert": expert,
            "topk": "9",
            "act_type": "ActivationType.Silu",
            "dtype": "torch.bfloat16",
            "q_dtype_a": "torch.float4_e2m1fn_x2",
            "q_dtype_w": "torch.float4_e2m1fn_x2",
            "q_type": "QuantType.per_1x32",
            "use_g1u1": "1",
            "doweight_stage1": "0",
            "block_m": block_m,
            "ksplit": ksplit,
            "us1": us1,
            "kernelName1": kernel1,
            "err1": "0.0%",
            "us2": us2,
            "kernelName2": kernel2,
            "err2": err2,
            "us": us,
            "run_1stage": "0",
            "tflops": "0.0",
            "bw": "0.0",
            "_tag": "",
        }
    )


def _prepare_env() -> None:
    if importlib.util.find_spec("aiter") is None:
        return
    cfg_path = Path("/tmp/gpumode_amd_moe_mxfp4_cfg_257k2_33k2_16128_blockmwide_stage2_sepqsort.csv")
    with cfg_path.open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=_FIELDNAMES)
        writer.writeheader()
        for cu_num in ("256", "288", "304"):
            for token in ("16", "128"):
                _write_row(
                    writer,
                    cu_num=cu_num,
                    token=token,
                    inter_dim="256",
                    expert="257",
                    ksplit="2",
                    block_m="32",
                )
            _write_row(
                writer,
                cu_num=cu_num,
                token="512",
                inter_dim="256",
                expert="257",
                ksplit="0",
                block_m="32",
                kernel1="moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16",
                kernel2="moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16",
                us1="97.693",
                us2="70.2081",
                us="167.9011",
                err2="1.3%",
            )
            for token in ("16", "128"):
                _write_row(
                    writer,
                    cu_num=cu_num,
                    token=token,
                    inter_dim="512",
                    expert="33",
                    ksplit="2",
                    block_m="32",
                )
    os.environ["AITER_CONFIG_FMOE"] = str(cfg_path)
    os.environ.setdefault("AITER_LOG_LEVEL", "ERROR")


def _block_size(config: dict) -> int | None:
    if config["n_routed_experts"] != 32 or config["n_shared_experts"] != 1:
        return None
    if config["bs"] == 512 and config["d_expert"] == 512:
        return 128
    if config["bs"] == 512 and config["d_expert"] == 2048:
        return 64
    return None


_prepare_env()

_fused_moe_mod = importlib.import_module("aiter.fused_moe")
_orig_fused_qsort = _fused_moe_mod.fused_dynamic_mxfp4_quant_moe_sort
_hip_quant_fp4 = get_hip_quant(QuantType.per_1x32)


def _patched_fused_dynamic_mxfp4_quant_moe_sort(
    x,
    sorted_ids,
    num_valid_ids,
    token_num,
    topk,
    block_size=32,
    scaling_mode="even",
):
    if token_num == 512 and topk > 1 and block_size in (64, 128):
        x_q, x_scale = _hip_quant_fp4(
            x,
            scale=None,
            quant_dtype=dtypes.fp4x2,
            num_rows=None,
            num_rows_factor=topk,
        )
        x_scale_sorted = fp4_utils.moe_mxfp4_sort(
            x_scale[: token_num * topk, :].view(token_num, topk, -1),
            sorted_ids=sorted_ids,
            num_valid_ids=num_valid_ids,
            token_num=token_num,
            block_size=block_size,
        )
        return x_q, x_scale_sorted
    return _orig_fused_qsort(
        x,
        sorted_ids=sorted_ids,
        num_valid_ids=num_valid_ids,
        token_num=token_num,
        topk=topk,
        block_size=block_size,
        scaling_mode=scaling_mode,
    )


_fused_moe_mod.fused_dynamic_mxfp4_quant_moe_sort = _patched_fused_dynamic_mxfp4_quant_moe_sort


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

    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=_block_size(config),
        hidden_pad=config["d_hidden_pad"] - config["d_hidden"],
        intermediate_pad=config["d_expert_pad"] - config["d_expert"],
    )