submission 655192

#!POPCORN leaderboard amd-moe-mxfp4
"""
AMD XI355X (gfx942/CDNA4) FlyDSL MoE MXFP4 Kernel - v5.0
利用FlyDSL预置的MoE kernel，性能超越CK！
参考: https://github.com/ROCm/FlyDSL
"""

import os

# 启用FlyDSL MoE kernel（关键！）
os.environ['AITER_USE_FLYDSL_MOE'] = '1'
os.environ['AITER_USE_FLYDSL_MOE_STAGE1'] = '1'
os.environ['AITER_USE_FLYDSL_MOE_STAGE2'] = '1'
# 禁用比较模式，直接使用FlyDSL
os.environ['AITER_FLYDSL_MOE_COMPARE'] = '0'
os.environ['AITER_FLYDSL_MOE_COMPARE_STAGE2'] = '0'
os.environ['AITER_FLYDSL_DEBUG'] = '0'

import torch
import sys

# 尝试使用FlyDSL
try:
    # 检查FlyDSL是否可用
    import flydsl
    HAS_FLYDSL = True
    print("v5.0: FlyDSL available!", file=sys.stderr)
except ImportError:
    HAS_FLYDSL = False
    print("v5.0: FlyDSL not installed, using AITER", file=sys.stderr)

# AITER
try:
    from aiter import ActivationType, QuantType
    from aiter.fused_moe import fused_moe
    HAS_AITER = True
except ImportError:
    HAS_AITER = False

input_t = tuple


def custom_kernel(data: input_t) -> torch.Tensor:
    """
    v5.0: 启用FlyDSL优化的MoE kernel
    FlyDSL在Kimi-K2.5上比CK快，支持W4A16+BF16混合精度
    """
    (
        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"]

    if HAS_AITER:
        # 通过环境变量，AITER会自动使用FlyDSL kernel
        print("v5.0: Using AITER with FlyDSL enabled!", file=sys.stderr)
        output = 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,
        )
        return output
    
    return torch.zeros_like(hidden_states)