submission 755039

"""DSV3 MoE MXFP4 submission backed by AITER's fused_moe library kernel.

The two notable behaviors layered on top of the default fused_moe call:

1. AITER's tuned MoE config CSVs are merged with two extra rows that pin
   block_m=32 (E=33, d_expert=512) and block_m=128 (E=33, d_expert=2048),
   selecting the CK 2-stage kernels that benchmark best on the dense E=33
   tail. The merged CSV is exported via AITER_CONFIG_FMOE so AITER's
   internal dispatch picks it up automatically.

2. For sparse small-M shapes, the merged CSV also forces ksplit=2 on every
   E=257 row with token<=128, and the per-call dispatch toggles
   AITER_KSPLIT="2" whenever estimated_m_per_expert<48 so the runtime
   override matches the CSV hint for the same shape family.
"""

import io
import os

import pandas as pd
from task import input_t, output_t

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


_EXISTING_CSVS = [
    "/home/runner/aiter/aiter/configs/tuned_fmoe.csv",
    "/home/runner/aiter/aiter/configs/model_configs/dsv3_fp4_tuned_fmoe.csv",
    "/home/runner/aiter/aiter/configs/model_configs/a8w8_blockscale_tuned_fmoe_qwen3_235b.csv",
]

_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",
]

# Pinned dense E=33 entries: block_m and CK kernel names that benchmark best
# on the d_expert=512 and d_expert=2048 tails of the DSV3 MoE workload.
_CUSTOM_E33_ENTRIES = """\
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
256,512,7168,512,33,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,True,False,32,0,0,moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0,0,moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16,0,0,False,0,0
256,512,7168,2048,33,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,True,False,128,0,0,moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0,0,moe_ck2stages_gemm2_256x128x128x128_1x4_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16,0,0,False,0,0
"""

# Merge AITER's shipped tuned CSVs with the pinned E=33 rows above. The
# resulting file is what AITER's dispatch reads via AITER_CONFIG_FMOE.
_csv_dataframes = []
for _csv_path in _EXISTING_CSVS:
    try:
        _csv_dataframes.append(pd.read_csv(_csv_path, on_bad_lines="skip"))
    except Exception:
        pass

_csv_dataframes.append(pd.read_csv(io.StringIO(_CUSTOM_E33_ENTRIES)))
_merged_fmoe_config = pd.concat(_csv_dataframes, ignore_index=True)

# Sparse E=257 small-M shapes benchmark best with ksplit=2; rewrite every
# matching CSV row before the dedup so the override survives drop_duplicates.
_sparse_e257_csv_mask = (
    (_merged_fmoe_config["expert"] == 257)
    & (_merged_fmoe_config["token"] <= 128)
)
_merged_fmoe_config.loc[_sparse_e257_csv_mask, "ksplit"] = 2
_merged_fmoe_config = _merged_fmoe_config.drop_duplicates(
    subset=_INDEX_COLS,
    keep="last",
)

_merged_csv_path = "/tmp/merged_fmoe_ksplit_e257.csv"
_merged_fmoe_config.to_csv(_merged_csv_path, index=False)
os.environ["AITER_CONFIG_FMOE"] = _merged_csv_path


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

    num_experts = config["n_routed_experts"] + config.get("n_shared_experts", 0)
    top_k = config["total_top_k"]
    estimated_m_per_expert = (config["bs"] * top_k + num_experts - 1) // num_experts

    # Runtime ksplit hint mirrors the CSV override: dense small-M (every
    # token routed to <48 experts on average) wins from ksplit=2, larger
    # batches use AITER's default scheduler.
    if estimated_m_per_expert < 48:
        os.environ["AITER_KSPLIT"] = "2"
    else:
        os.environ["AITER_KSPLIT"] = "0"

    return fused_moe(
        hidden_states,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        topk_weights,
        topk_ids,
        activation=ActivationType.Silu,
        quant_type=QuantType.per_1x32,
        doweight_stage1=False,
        w1_scale=gate_up_weight_scale_shuffled,
        w2_scale=down_weight_scale_shuffled,
        hidden_pad=hidden_pad,
        intermediate_pad=intermediate_pad,
    )