submission 749771

from task import input_t, output_t

import csv
import importlib.util
import os


_MOD = None
_FMOE_OVERRIDE_PATH = "/tmp/amd_moe_mxfp4_fmoe_override.csv"
_FMOE_INDEX_FIELDS = (
    "cu_num",
    "token",
    "model_dim",
    "inter_dim",
    "expert",
    "topk",
    "act_type",
    "dtype",
    "q_dtype_a",
    "q_dtype_w",
    "q_type",
    "use_g1u1",
    "doweight_stage1",
)
_BASE_OVERRIDE_ROW = {
    "cu_num": "256",
    "model_dim": "7168",
    "inter_dim": "256",
    "expert": "257",
    "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",
    "ksplit": "0",
    "us1": "0",
    "err1": "0",
    "us2": "0",
    "err2": "0",
    "us": "0",
    "run_1stage": "0",
    "tflops": "0",
    "bw": "0",
    "_tag": "",
}
_KERNEL_64_32 = (
    "moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_"
    "Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
)
_KERNEL_256_32 = (
    "moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_"
    "Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
)
_KERNEL_256_64 = (
    "moe_ck2stages_gemm1_256x64x128x128_1x4_MulABScaleShuffled_v3_"
    "Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
)
_STAGE2_64_32 = (
    "moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_"
    "v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
)
_STAGE2_256_32 = (
    "moe_ck2stages_gemm2_256x32x128x128_1x4_MulABScaleExpertWeightShuffled_"
    "v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
)
_STAGE2_64_64 = (
    "moe_ck2stages_gemm2_64x64x128x128_1x1_MulABScaleExpertWeightShuffled_"
    "v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
)
_FAMILY_7168_256_ROWS = (
    ("1", "32", _KERNEL_64_32, _STAGE2_64_32),
    ("2", "32", _KERNEL_64_32, _STAGE2_64_32),
    ("4", "32", _KERNEL_64_32, _STAGE2_64_32),
    ("8", "32", _KERNEL_64_32, _STAGE2_64_32),
    ("16", "32", _KERNEL_64_32, _STAGE2_64_32),
    ("32", "32", _KERNEL_64_32, _STAGE2_64_32),
    ("64", "32", _KERNEL_64_32, _STAGE2_64_32),
    ("128", "32", _KERNEL_256_32, _STAGE2_64_32),
    ("256", "32", _KERNEL_256_32, _STAGE2_64_32),
    ("512", "32", _KERNEL_64_32, _STAGE2_64_32),
    ("1024", "32", _KERNEL_64_32, _STAGE2_64_32),
    ("2048", "64", _KERNEL_256_64, _STAGE2_64_64),
)
_BASE_OVERRIDE_ROW_4096_1024 = {
    **_BASE_OVERRIDE_ROW,
    "model_dim": "4096",
    "inter_dim": "1024",
}
_FAMILY_4096_1024_ROWS = (
    ("1", "32", _KERNEL_64_32, _STAGE2_256_32),
    ("2", "32", _KERNEL_64_32, _STAGE2_64_32),
    ("4", "32", _KERNEL_64_32, _STAGE2_64_32),
    ("8", "32", _KERNEL_64_32, _STAGE2_64_32),
    ("16", "32", _KERNEL_64_32, _STAGE2_64_32),
    ("32", "32", _KERNEL_64_32, _STAGE2_64_32),
    ("64", "32", _KERNEL_64_32, _STAGE2_64_32),
    ("128", "32", _KERNEL_64_32, _STAGE2_64_32),
    ("256", "32", _KERNEL_64_32, _STAGE2_64_32),
    ("512", "32", _KERNEL_64_32, _STAGE2_64_32),
    ("1024", "32", _KERNEL_64_32, _STAGE2_64_32),
    ("2048", "32", _KERNEL_64_32, _STAGE2_64_32),
    ("4096", "64", _KERNEL_256_64, _STAGE2_64_64),
    ("8192", "64", _KERNEL_256_64, _STAGE2_64_64),
    ("16384", "32", _KERNEL_256_32, _STAGE2_64_32),
)


def _family_rows(base_row, family_rows):
    rows = []
    for token, block_m, kernel1, kernel2 in family_rows:
        row = dict(base_row)
        row.update(
            {
                "token": token,
                "block_m": block_m,
                "kernelName1": kernel1,
                "kernelName2": kernel2,
            }
        )
        rows.append(row)
    return rows


def _override_rows():
    return _family_rows(_BASE_OVERRIDE_ROW, _FAMILY_7168_256_ROWS) + _family_rows(
        _BASE_OVERRIDE_ROW_4096_1024, _FAMILY_4096_1024_ROWS
    )

def _aiter_config_candidates():
    spec = importlib.util.find_spec("aiter")
    if spec is None or spec.origin is None:
        return ()
    pkg_dir = os.path.dirname(spec.origin)
    config_dir = os.path.join(pkg_dir, "configs")
    return (
        os.path.join(config_dir, "tuned_fmoe.csv"),
        os.path.join(
            config_dir,
            "model_configs",
            "a8w8_blockscale_tuned_fmoe_qwen3_235b.csv",
        ),
        os.path.join(config_dir, "model_configs", "dsv3_fp4_tuned_fmoe.csv"),
    )


def _configure_aiter_env():
    fieldnames = []
    row_order = []
    rows_by_key = {}
    for path in _aiter_config_candidates():
        if not os.path.exists(path):
            continue
        with open(path, "r", encoding="ascii", newline="") as f:
            reader = csv.DictReader(f)
            if reader.fieldnames is None:
                continue
            for fieldname in reader.fieldnames:
                if fieldname not in fieldnames:
                    fieldnames.append(fieldname)
            for row in reader:
                key = tuple(row.get(name, "") for name in _FMOE_INDEX_FIELDS)
                if key not in rows_by_key:
                    row_order.append(key)
                rows_by_key[key] = row

    if not fieldnames:
        fieldnames = list(_BASE_OVERRIDE_ROW.keys())
    else:
        for fieldname in _BASE_OVERRIDE_ROW:
            if fieldname not in fieldnames:
                fieldnames.append(fieldname)

    for override_row in _override_rows():
        override_key = tuple(override_row[name] for name in _FMOE_INDEX_FIELDS)
        if override_key not in rows_by_key:
            row_order.append(override_key)
        rows_by_key[override_key] = override_row
    with open(_FMOE_OVERRIDE_PATH, "w", encoding="ascii", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        for key in row_order:
            row = rows_by_key[key]
            writer.writerow({name: row.get(name, "") for name in fieldnames})
    os.environ["AITER_CONFIG_FMOE"] = _FMOE_OVERRIDE_PATH


def _ensure_mod():
    global _MOD
    if _MOD is not None:
        return _MOD

    _configure_aiter_env()

    import torch
    from aiter import ActivationType, QuantType
    from aiter.fused_moe import fused_moe

    _MOD = {
        "torch": torch,
        "ActivationType": ActivationType,
        "QuantType": QuantType,
        "fused_moe": fused_moe,
    }
    return _MOD

def _run_fused_2stage(data: input_t) -> output_t:
    mod = _ensure_mod()
    (
        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 mod["fused_moe"](
        hidden_states,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        topk_weights,
        topk_ids,
        expert_mask=None,
        activation=mod["ActivationType"].Silu,
        quant_type=mod["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,
    )


def _run_direct(data: input_t) -> output_t:
    return _run_fused_2stage(data)


def custom_kernel(data: input_t) -> output_t:
    return _run_direct(data)