submission 703448

"""
Solution 180: s179 + FlyDSL stage2 for E=257/bs=512.

E=257/bs=512 currently uses dsv3 default (CK stage1 + CK stage2).
Try: CK stage1 + FlyDSL stage2 reduce (which helps E=33 shapes).
Also try FlyDSL stage2 t32x256x256_reduce for E=33/bs=16 (instead of t32x256x256 in s170).

Changes vs s179:
1. E=257/bs=512: Add explicit entry with FlyDSL stage2 t64x256x256_reduce
2. E=33/bs=512/d=2048: Try t64x128x256_reduce (tile_n=128, less parallelism but different cache pattern)
"""
import os
import tempfile

os.environ["PYTORCH_ROCM_ARCH"] = "gfx950"

import aiter
import aiter.fused_moe as _fmoe_mod

_aiter_root = os.path.dirname(os.path.dirname(os.path.abspath(aiter.__file__)))
_default_csv = os.path.join(_aiter_root, "aiter", "configs", "tuned_fmoe.csv")
_model_csv = os.path.join(_aiter_root, "aiter", "configs", "model_configs", "dsv3_fp4_tuned_fmoe.csv")

CUSTOM_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,_tag
256,16,7168,256,257,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,2,0,x,0,0,x,0,0,0,0,0,
256,128,7168,256,257,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,4,0,x,0,0,x,0,0,0,0,0,
256,512,7168,256,257,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,0,moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0,0,flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce,0,0,0,0,0,
256,16,7168,512,33,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,2,0,x,0,0,flydsl_moe2_afp4_wfp4_bf16_t32x256x256_reduce,0,0,0,0,0,
256,128,7168,512,33,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,64,2,0,x,0,0,flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce,0,0,0,0,0,
256,512,7168,512,33,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,64,0,0,moe_ck2stages_gemm1_256x64x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0,0,flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce,0,0,0,0,0,
256,512,7168,2048,33,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,128,0,0,moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0,0,flydsl_moe2_afp4_wfp4_bf16_t64x128x256_reduce,0,0,0,0,0,
"""
# Key changes:
# - E=257/bs=512: FlyDSL stage2 t64x256x256_reduce (was CK stage2)
# - E=33/d=2048: FlyDSL stage2 t64x128x256_reduce (was t64x256x256_reduce)
# - E=33/d=512: CK stage1 256x64x128x128 (was 256x64x128x128 same)

_custom_csv = os.path.join(tempfile.gettempdir(), "custom_fmoe_180.csv")
with open(_custom_csv, "w") as f:
    f.write(CUSTOM_ENTRIES)

_paths = [p for p in [_default_csv, _model_csv] if os.path.exists(p)]
_paths.append(_custom_csv)
os.environ["AITER_CONFIG_FMOE"] = ":".join(_paths)

import torch
from task import input_t, output_t
from aiter import ActivationType, QuantType
from aiter.fused_moe import fused_moe, use_nt, get_ksplit


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"]
    E = config["n_routed_experts"] + config["n_shared_experts"]
    bs = config["bs"]

    _fmoe_mod._USE_OPUS_MOE_SORTING = (E > 33)

    if E <= 33:
        os.environ["AITER_USE_NT"] = "0"
    else:
        os.environ["AITER_USE_NT"] = "-1"

    os.environ.pop("AITER_KSPLIT", None)

    use_nt.cache_clear()
    get_ksplit.cache_clear()

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