submission 726454

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

import csv
import importlib.util
import os
from pathlib import Path
from typing import Iterable, Optional

os.environ.setdefault("TRITON_CACHE_DIR", "/tmp/triton_cache")
os.environ.setdefault("AITER_USE_OPUS_MOE_SORTING", "1")
os.environ.setdefault("AITER_BYPASS_TUNE_CONFIG", "0")
os.environ.setdefault("AITER_ONLINE_TUNE", "0")

_INDEX_COLS = (
    "cu_num", "token", "model_dim", "inter_dim", "expert", "topk",
    "act_type", "dtype", "q_dtype_a", "q_dtype_w", "q_type",
    "use_g1u1", "doweight_stage1",
)

# v015k: Best combined version (v015j + sepqsort)
#
# Optimizations stacked:
#   1. CSV merge (dsv3_fp4 + model configs) for 257-expert shapes
#   2. Tuner-found CK kernel configs for B5/B6
#   3. CKTile ksplit=2 for decode/thin shapes (B1/B2/B4) — skips quant entirely
#   4. sepqsort: separate quant() + sort() replaces fused Triton kernel (B3/B5/B6/B7)
#   5. Opus sorting + is_shuffled for CKTile path
#   6. B7: default heuristic + block_m=64
#
# Expected geomean: ~139μs (vs v011 baseline 181μs = -23%)

_KN1_SMALL = "moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
_KN1_LARGE = "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"

_CUSTOM_CONFIGS = {
    # CKTile: ksplit=2 for thin/decode shapes (a16w4, no quant overhead)
    (16,  7168,  256, 257, 9): (16, 2, "cktile", "cktile"),      # B1: CKTile -27%
    (128, 7168,  256, 257, 9): (16, 2, "cktile", "cktile"),      # B2: CKTile -8.6%
    (16,  7168,  512,  33, 9): (16, 2, "cktile", "cktile"),      # B4: CKTile -27%
    # CK 2-stage: tuner configs for medium/thick shapes (a4w4, fp4 saves BW)
    (128, 7168,  512,  33, 9): (32, 0, _KN1_SMALL, _KN2_SMALL),  # B5: CK tuner
    (512, 7168,  512,  33, 9): (32, 0, _KN1_LARGE, _KN2_SMALL),  # B6: CK tuner -15%
    # B3 (257,256,512): no custom entry → dsv3_fp4 CSV (CK 2-stage)
    # B7 (33,2048,512): no custom entry → default heuristic + block_m=64
}


def _as_float(v):
    try:
        return float(v)
    except Exception:
        return float("inf")


def _iter_existing(paths: Iterable[Path]) -> Iterable[Path]:
    for p in paths:
        if p.exists() and p.is_file():
            yield p


def _make_custom_row(token, model_dim, inter_dim, expert, topk,
                     block_m, ksplit, kn1, kn2):
    return {
        "cu_num": "256", "token": str(token), "model_dim": str(model_dim),
        "inter_dim": str(inter_dim), "expert": str(expert), "topk": str(topk),
        "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": str(block_m), "ksplit": str(ksplit),
        "kernelName1": kn1, "kernelName2": kn2,
        "us1": "0", "err1": "0", "us2": "0", "err2": "0",
        "us": "1", "run_1stage": "0", "tflops": "0", "bw": "0",
    }


def _prepare_writable_fmoe_config() -> Optional[str]:
    if os.environ.get("AITER_CONFIG_FMOE"):
        return os.environ["AITER_CONFIG_FMOE"]
    try:
        spec = importlib.util.find_spec("aiter")
        if spec is None:
            return None
        if spec.origin:
            pkg_dir = Path(spec.origin).resolve().parent
        else:
            locations = list(spec.submodule_search_locations or [])
            if not locations:
                return None
            pkg_dir = Path(locations[0]).resolve()

        cfg_dir = pkg_dir / "configs"
        model_cfg_dir = cfg_dir / "model_configs"
        candidates = list(
            _iter_existing(
                [cfg_dir / "tuned_fmoe.csv"]
                + (sorted(p for p in model_cfg_dir.glob("*tuned_fmoe*.csv")
                          if "untuned" not in p.name)
                   if model_cfg_dir.exists() else [])
            )
        )
        if not candidates:
            return None

        out_dir = Path("/tmp/aiter_configs")
        out_dir.mkdir(parents=True, exist_ok=True)
        out_path = out_dir / "tuned_fmoe_v015k.csv"

        fieldnames: list[str] = []
        seen_fields: set[str] = set()
        rows: list[dict[str, str]] = []

        for csv_path in candidates:
            with csv_path.open("r", newline="") as f:
                reader = csv.DictReader(f)
                if not reader.fieldnames:
                    continue
                for name in reader.fieldnames:
                    if name not in seen_fields:
                        seen_fields.add(name)
                        fieldnames.append(name)
                for row in reader:
                    rows.append({k: "" if v is None else v for k, v in row.items()})

        if not rows:
            return None

        for key, (bm, ks, kn1, kn2) in _CUSTOM_CONFIGS.items():
            rows.append(_make_custom_row(*key, bm, ks, kn1, kn2))

        has_tag = "_tag" in seen_fields
        best: dict[tuple, dict[str, str]] = {}
        for row in rows:
            key = tuple(row.get(c, "") for c in _INDEX_COLS)
            if has_tag:
                key = key + ((row.get("_tag", "") or ""),)
            prev = best.get(key)
            if prev is None or _as_float(row.get("us", "inf")) < _as_float(prev.get("us", "inf")):
                best[key] = row

        merged = sorted(best.values(), key=lambda r: _as_float(r.get("us", "inf")))

        with out_path.open("w", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore")
            writer.writeheader()
            for row in merged:
                writer.writerow({name: row.get(name, "") for name in fieldnames})

        os.environ["AITER_CONFIG_FMOE"] = str(out_path)
        return str(out_path)
    except Exception:
        return None


_PREPARED = _prepare_writable_fmoe_config()

import torch
from aiter import ActivationType, QuantType, dtypes
from aiter import get_hip_quant as _get_quant
from aiter.utility import fp4_utils
from aiter.fused_moe import fused_moe
import aiter.fused_moe as _fm

# --- sepqsort monkey-patch ---
# Replace fused quant+sort with separate quant() → sort().
# The fused Triton kernel may have overhead for small batches.
_quant_func = _get_quant(QuantType.per_1x32)

def _sep_quant_sort(x, sorted_ids, num_valid_ids, token_num, topk, block_size):
    a, scale = _quant_func(x, scale=None, quant_dtype=dtypes.fp4x2, num_rows=None)
    if topk > 1:
        scale = fp4_utils.moe_mxfp4_sort(
            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,
        )
    else:
        scale = fp4_utils.moe_mxfp4_sort(
            scale, sorted_ids=sorted_ids, num_valid_ids=num_valid_ids,
            token_num=token_num, block_size=block_size,
        )
    return a, scale

_fm.fused_dynamic_mxfp4_quant_moe_sort = _sep_quant_sort


@torch.no_grad()
def custom_kernel(data):
    (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

    del gate_up_weight, down_weight, gate_up_weight_scale, down_weight_scale

    hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
    intermediate_pad = config["d_expert_pad"] - config["d_expert"]

    # Required for CKTile path (B1, B2, B4): ksplit>1 + is_shuffled triggers a16w4.
    # Harmless for CK 2-stage shapes (B3, B5, B6, B7) since their ksplit=0.
    gate_up_weight_shuffled.is_shuffled = True

    # B7: no CSV entry → default heuristic. Override block_m=64 (proven best).
    e_total = config["n_routed_experts"] + config["n_shared_experts"]
    shape_key = (e_total, config["d_expert"], config["bs"])
    block_m = {(33, 2048, 512): 64}.get(shape_key)

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