submission 531025

import os
import sys
import functools
from pathlib import Path

import pandas as pd
import torch
from aiter import ActivationType, QuantType
from task import input_t, output_t


DEBUG_FMOE = os.getenv("JOSU_DEBUG_FMOE") == "1"
FORCE_REBUILD_FMOE = DEBUG_FMOE or os.getenv("JOSU_REBUILD_FMOE") == "1"
CFG_VERSION = "v53_safe_hybrid"


def _bootstrap_cfg() -> str:
    out_path = Path(f"/tmp/josu_cfg_hybrid_under150_fmoe_{CFG_VERSION}.csv")
    if out_path.exists() and out_path.stat().st_size > 0 and not FORCE_REBUILD_FMOE:
        return str(out_path)

    source_paths = [
        Path("/home/runner/aiter/aiter/configs/tuned_fmoe.csv"),
        Path(
            "/home/runner/aiter/aiter/configs/model_configs/"
            "a8w8_blockscale_tuned_fmoe_qwen3_235b.csv"
        ),
        Path("/home/runner/aiter/aiter/configs/model_configs/dsv3_fp4_tuned_fmoe.csv"),
    ]
    frames = [pd.read_csv(path) for path in source_paths if path.exists()]
    merged = pd.concat(frames, ignore_index=True)

    untuned_path = Path("/home/runner/aiter/aiter/configs/untuned_fmoe.csv")
    keys = pd.read_csv(untuned_path, nrows=0).columns.tolist()
    if "cu_num" not in keys:
        keys.append("cu_num")

    defaults = {column: "" for column in merged.columns}
    defaults.update(
        {
            "cu_num": 256,
            "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": 32,
            "ksplit": 0,
            "us1": 0.0,
            "kernelName1": "",
            "err1": "0.0%",
            "us2": 0.0,
            "kernelName2": "",
            "err2": "0.0%",
            "us": 0.0,
            "run_1stage": 0,
            "tflops": 0.0,
            "bw": 0.0,
            "_tag": "",
        }
    )

    K1_SMALL = "moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
    K1_MED = "moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
    K1_LARGE = "moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
    K2_SMALL = "moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
    K2_LARGE = "moe_ck2stages_gemm2_256x128x128x128_1x4_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"

    # Safe hybrid: use ksplit=4 for E=33 bs=16 instead of ksplit=7 (avoids intermittent correctness failures)
    # ksplit=7 for E=257 bs=16 only (larger expert count seems stable with ksplit=7)
    online_tuned_rows = [
        # E=257 bs=16: ksplit=7 (reliable for E=257)
        {"token": 16, "model_dim": 7168, "inter_dim": 256, "expert": 257, "topk": 9,
         "block_m": 32, "ksplit": 7,
         "kernelName1": K1_SMALL, "kernelName2": K2_SMALL, "run_1stage": 0},
        # E=257 bs=128: ksplit=4
        {"token": 128, "model_dim": 7168, "inter_dim": 256, "expert": 257, "topk": 9,
         "block_m": 32, "ksplit": 4,
         "kernelName1": K1_SMALL, "kernelName2": K2_SMALL, "run_1stage": 0},
        # E=257 bs=512: ksplit=0
        {"token": 512, "model_dim": 7168, "inter_dim": 256, "expert": 257, "topk": 9,
         "block_m": 32, "ksplit": 0,
         "kernelName1": K1_SMALL, "kernelName2": K2_SMALL, "run_1stage": 0},

        # E=33 bs=16: ksplit=4 (SAFE - ksplit=7 has intermittent failures for E=33)
        {"token": 16, "model_dim": 7168, "inter_dim": 512, "expert": 33, "topk": 9,
         "block_m": 32, "ksplit": 4,
         "kernelName1": K1_SMALL, "kernelName2": K2_SMALL, "run_1stage": 0},
        # E=33 bs=128: ksplit=4 + K1_MED
        {"token": 128, "model_dim": 7168, "inter_dim": 512, "expert": 33, "topk": 9,
         "block_m": 32, "ksplit": 4,
         "kernelName1": K1_MED, "kernelName2": K2_SMALL, "run_1stage": 0},
        # E=33 bs=512: ksplit=0 + NT
        {"token": 512, "model_dim": 7168, "inter_dim": 512, "expert": 33, "topk": 9,
         "block_m": 32, "ksplit": 0,
         "kernelName1": K1_MED, "kernelName2": K2_SMALL, "run_1stage": 0},

        # E=33 d=2048
        {"token": 512, "model_dim": 7168, "inter_dim": 2048, "expert": 33, "topk": 9,
         "block_m": 128, "ksplit": 0,
         "kernelName1": K1_LARGE, "kernelName2": K2_LARGE, "run_1stage": 0},
    ]

    rows = []
    for row in online_tuned_rows:
        seeded = defaults.copy()
        seeded.update(row)
        seeded["us"] = -1.0
        rows.append(seeded)

    merged = pd.concat([merged, pd.DataFrame(rows)], ignore_index=True)
    merged = (
        merged.sort_values("us")
        .drop_duplicates(subset=keys, keep="first")
        .reset_index(drop=True)
    )

    out_path.parent.mkdir(parents=True, exist_ok=True)
    merged.to_csv(out_path, index=False)
    return str(out_path)


os.environ["AITER_CONFIG_FMOE"] = _bootstrap_cfg()

from aiter.fused_moe import fused_moe
import aiter.fused_moe as _fm

_orig_ck_moe_stage1 = _fm.ck_moe_stage1

@functools.wraps(_orig_ck_moe_stage1)
def _nt_ck_moe_stage1(*args, **kwargs):
    kwargs["use_non_temporal_load"] = True
    return _orig_ck_moe_stage1(*args, **kwargs)

_fm.ck_moe_stage1 = _nt_ck_moe_stage1

if hasattr(_fm, "ck_moe_stage2"):
    _orig_ck_moe_stage2 = _fm.ck_moe_stage2

    @functools.wraps(_orig_ck_moe_stage2)
    def _nt_ck_moe_stage2(*args, **kwargs):
        kwargs["use_non_temporal_load"] = True
        return _orig_ck_moe_stage2(*args, **kwargs)

    _fm.ck_moe_stage2 = _nt_ck_moe_stage2

print("[v53] safe hybrid: ksplit=4 for E=33 bs=16 (avoids ksplit=7 flakiness) + NT", file=sys.stderr)


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