submission 695763

"""
V98: V97 + Try FlyDSL stage2 for E=257 bs=512 (instead of CK 64WG).
DSV3 CSV shows E=256 gets flydsl_moe2_afp4_wfp4_bf16_t32x256x256_atomic.
E=257 gets CK stage2. Let's test if FlyDSL is faster for E=257 too.
Also test CK_256WG for stage1 while block_m=64 for E=257 bs=512.
Plus dispatch_policy=2 from V97.
"""
#!POPCORN leaderboard amd-moe-mxfp4
#!POPCORN gpu MI355X
import torch
import os
import aiter
from aiter import ActivationType, QuantType, dtypes
from aiter.fused_moe import (
    fused_moe, get_2stage_cfgs, get_inter_dim, get_padded_M,
    fused_dynamic_mxfp4_quant_moe_sort,
)
from task import input_t, output_t

# Stage1 kernels
CK_S1_256WG_64 = "moe_ck2stages_gemm1_256x64x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
CK_S1_64WG_32 = "moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"

# Stage2 kernels
FLY_S2_32x256A = "flydsl_moe2_afp4_wfp4_bf16_t32x256x256_atomic"
FLY_S2_64x256R = "flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce"
FLY_S2_64x128R = "flydsl_moe2_afp4_wfp4_bf16_t64x128x256_reduce"
CK_S2_64WG_32 = "moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"

_CSV_HEADER = "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"

def _row(cu, tok, mdim, idim, E, topk, bm, k1, k2, ks=0):
    return f"{cu},{tok},{mdim},{idim},{E},{topk},ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,True,False,{bm},{ks},0,{k1},0,0,{k2},0,0,False,0,0"

def _build_custom_csv():
    rows = [_CSV_HEADER]
    # E=33 shapes (same as V92/V97)
    rows.append(_row(256, 512, 7168, 512, 33, 9, 64, CK_S1_256WG_64, FLY_S2_64x256R))
    rows.append(_row(256, 512, 7168, 2048, 33, 9, 64, CK_S1_256WG_64, FLY_S2_64x128R))
    # E=257 bs=512: Try FlyDSL atomic stage2 with block_m=32
    rows.append(_row(256, 512, 7168, 256, 257, 9, 32, CK_S1_64WG_32, FLY_S2_32x256A))
    return "\n".join(rows) + "\n"

_CKTILE_SHAPES = {(16, 257), (128, 257), (16, 33), (128, 33)}
_CKTILE_KSPLIT = {
    (16, 257): 7,
    (128, 257): 4,
    (16, 33): 2,
    (128, 33): 2,
}

_cache = {}
_initialized = False
_sort_fn = None
_DISPATCH_POLICY = 2


def custom_kernel(data: input_t) -> output_t:
    global _initialized, _sort_fn

    (
        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

    M = hidden_states.shape[0]
    d_hidden = config["d_hidden"]
    d_hidden_pad = config["d_hidden_pad"]
    d_expert = config["d_expert"]
    d_expert_pad = config["d_expert_pad"]
    E = config["n_routed_experts"] + config["n_shared_experts"]
    topk = config["total_top_k"]
    hidden_pad = d_hidden_pad - d_hidden
    intermediate_pad = d_expert_pad - d_expert

    w1s = gate_up_weight_scale_shuffled.view(dtypes.fp8_e8m0)
    w2s = down_weight_scale_shuffled.view(dtypes.fp8_e8m0)

    if not _initialized:
        csv_path = "/tmp/custom_tuned_fmoe_v98.csv"
        with open(csv_path, 'w') as f:
            f.write(_build_custom_csv())
        aiter_root = os.path.dirname(os.path.abspath(aiter.__file__))
        default_csv = os.path.join(aiter_root, "configs", "tuned_fmoe.csv")
        dsv3_csv = os.path.join(aiter_root, "configs", "model_configs", "dsv3_fp4_tuned_fmoe.csv")
        paths = []
        if os.path.exists(default_csv):
            paths.append(default_csv)
        if os.path.exists(dsv3_csv):
            paths.append(dsv3_csv)
        paths.append(csv_path)
        os.environ["AITER_CONFIG_FMOE"] = ":".join(paths)
        os.environ["AITER_KSPLIT"] = "0"
        get_2stage_cfgs.cache_clear()
        _sort_fn = getattr(aiter, 'moe_sorting_opus_fwd', aiter.moe_sorting_fwd)
        _initialized = True

    shape_key = (M, E, d_expert)
    is_cktile = (M, E) in _CKTILE_SHAPES

    if shape_key not in _cache:
        if is_cktile:
            ks = _CKTILE_KSPLIT.get((M, E), 2)
            os.environ["AITER_KSPLIT"] = str(ks)
            os.environ["AITER_BYPASS_TUNE_CONFIG"] = "1"
        else:
            os.environ["AITER_KSPLIT"] = "0"
            os.environ["AITER_BYPASS_TUNE_CONFIG"] = "0"
        get_2stage_cfgs.cache_clear()

        result = fused_moe(
            hidden_states, gate_up_weight_shuffled, down_weight_shuffled,
            topk_weights, topk_ids,
            activation=ActivationType.Silu, quant_type=QuantType.per_1x32,
            w1_scale=w1s, w2_scale=w2s,
            hidden_pad=hidden_pad, intermediate_pad=intermediate_pad,
        )

        E2, model_dim, inter_dim = get_inter_dim(
            gate_up_weight_shuffled.shape, down_weight_shuffled.shape
        )
        is_shuffled = getattr(gate_up_weight_shuffled, "is_shuffled", False)
        metadata = get_2stage_cfgs(
            get_padded_M(M), model_dim, inter_dim, E, topk,
            torch.bfloat16, dtypes.fp4x2, dtypes.fp4x2,
            QuantType.per_1x32, True, ActivationType.Silu,
            False, hidden_pad, intermediate_pad, is_shuffled,
        )
        block_m = metadata.block_m

        device = hidden_states.device
        max_tok = M * topk + E * block_m
        max_mb = (max_tok + block_m - 1) // block_m

        sorted_ids = torch.empty(max_tok, dtype=torch.int32, device=device)
        sorted_weights = torch.empty(max_tok, dtype=torch.float32, device=device)
        sorted_expert_ids = torch.empty(max_mb, dtype=torch.int32, device=device)
        num_valid_ids = torch.empty(2, dtype=torch.int32, device=device)
        moe_buf = torch.zeros((M, d_hidden_pad), dtype=torch.bfloat16, device=device)

        entry = {
            'block_m': block_m,
            'metadata': metadata,
            'sorted_ids': sorted_ids,
            'sorted_weights': sorted_weights,
            'sorted_expert_ids': sorted_expert_ids,
            'num_valid_ids': num_valid_ids,
            'moe_buf': moe_buf,
            'is_cktile': is_cktile,
            'inter_dim': inter_dim,
        }

        if not is_cktile:
            entry['a2_buf'] = torch.empty(
                (M, topk, inter_dim),
                dtype=torch.bfloat16,
                device=device,
            )

        _cache[shape_key] = entry
        return result

    c = _cache[shape_key]

    _sort_fn(topk_ids, topk_weights, c['sorted_ids'], c['sorted_weights'],
             c['sorted_expert_ids'], c['num_valid_ids'], c['moe_buf'],
             E, c['block_m'], None, None, _DISPATCH_POLICY)

    if c['is_cktile']:
        a2 = c['metadata'].stage1(
            hidden_states,
            gate_up_weight_shuffled, down_weight_shuffled,
            c['sorted_ids'], c['sorted_expert_ids'], c['num_valid_ids'],
            None, topk,
            block_m=c['block_m'], a1_scale=None, w1_scale=w1s,
            sorted_weights=None,
        )
        c['metadata'].stage2(
            a2,
            gate_up_weight_shuffled, down_weight_shuffled,
            c['sorted_ids'], c['sorted_expert_ids'], c['num_valid_ids'],
            c['moe_buf'], topk,
            w2_scale=w2s, a2_scale=None, block_m=c['block_m'],
            sorted_weights=c['sorted_weights'],
        )
    else:
        a1_fp4, a1_scale_sorted = fused_dynamic_mxfp4_quant_moe_sort(
            hidden_states,
            sorted_ids=c['sorted_ids'],
            num_valid_ids=c['num_valid_ids'],
            token_num=M,
            topk=1,
            block_size=c['block_m'],
        )
        c['metadata'].stage1(
            a1_fp4,
            gate_up_weight_shuffled, down_weight_shuffled,
            c['sorted_ids'], c['sorted_expert_ids'], c['num_valid_ids'],
            c['a2_buf'], topk,
            block_m=c['block_m'],
            a1_scale=a1_scale_sorted,
            w1_scale=w1s,
            sorted_weights=None,
        )
        a2_flat = c['a2_buf'].view(-1, c['inter_dim'])
        a2_fp4, a2_scale_sorted = fused_dynamic_mxfp4_quant_moe_sort(
            a2_flat,
            sorted_ids=c['sorted_ids'],
            num_valid_ids=c['num_valid_ids'],
            token_num=M,
            topk=topk,
            block_size=c['block_m'],
        )
        a2_fp4_3d = a2_fp4.view(M, topk, -1)
        c['metadata'].stage2(
            a2_fp4_3d,
            gate_up_weight_shuffled, down_weight_shuffled,
            c['sorted_ids'], c['sorted_expert_ids'], c['num_valid_ids'],
            c['moe_buf'], topk,
            w2_scale=w2s,
            a2_scale=a2_scale_sorted,
            block_m=c['block_m'],
            sorted_weights=c['sorted_weights'],
        )

    return c['moe_buf']