submission 678641

#!POPCORN leaderboard amd-moe-mxfp4
#!POPCORN gpu MI355X

import os
import sys
import torch
from task import input_t, output_t

# System-level env vars (set before any aiter imports trigger caching)
os.environ["HIP_FORCE_DEV_KERNARG"] = "1"
os.environ["AITER_USE_NT"] = "1"

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

# Import separated quant + sort functions
from aiter.ops.quant import per_1x32_f4_quant_hip as _quant_hip
from aiter.utility import fp4_utils as _fp4_utils

# Monkey-patch: replace fused quant+sort with separated HIP quant + Triton sort
def _split_quant_moe_sort(x, sorted_ids, num_valid_ids, token_num, topk, block_size=32, scaling_mode="even"):
    x_fp4, scale = _quant_hip(x)
    sorted_scale = _fp4_utils.moe_mxfp4_sort(
        scale, sorted_ids=sorted_ids, num_valid_ids=num_valid_ids,
        token_num=token_num, block_size=block_size,
    )
    return x_fp4, sorted_scale

_fm.fused_dynamic_mxfp4_quant_moe_sort = _split_quant_moe_sort
print("[moe] Monkey-patched fused_quant+sort → split HIP quant + Triton sort", file=sys.stderr)

# Register t16x256x128_atomic in _KERNEL_PARAMS
try:
    from aiter.ops.flydsl.moe_kernels import _KERNEL_PARAMS
    _t16_name = "flydsl_moe2_afp4_wfp4_bf16_t16x256x128_atomic"
    if _t16_name not in _KERNEL_PARAMS:
        _KERNEL_PARAMS[_t16_name] = {
            "stage": 2, "a_dtype": "fp4", "b_dtype": "fp4", "out_dtype": "bf16",
            "tile_m": 16, "tile_n": 256, "tile_k": 128,
            "mode": "atomic", "MPerBlock": 16,
        }
        print(f"[moe] Registered {_t16_name}", file=sys.stderr)
except Exception as e:
    print(f"[moe] _KERNEL_PARAMS registration failed: {e}", file=sys.stderr)

# CSV injection: shapes 5+6+7 with FlyDSL t16 stage2
# Shape 5 (NEW): M=128, E=33, inter=512, block_m=64 — fp4 2-stage instead of BYPASS
_tune_file = _fm.AITER_CONFIGS.AITER_CONFIG_FMOE_FILE
try:
    with open(_tune_file, "a") as f:
        # Shape 5: M=128, E=33, inter=512, block_m=64
        f.write("256,128,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%,0,flydsl_moe2_afp4_wfp4_bf16_t16x256x128_atomic,0.0%,0,0,0,0\n")
        # Shape 6: M=512, E=33, inter=512, block_m=128 (v67: -6.7% from 120→112µs)
        f.write("256,512,7168,512,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%,0,flydsl_moe2_afp4_wfp4_bf16_t16x256x128_atomic,0.0%,0,0,0,0\n")
        # Shape 7: M=512, E=33, inter=2048, block_m=64
        f.write("256,512,7168,2048,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%,0,flydsl_moe2_afp4_wfp4_bf16_t16x256x128_atomic,0.0%,0,0,0,0\n")
    print("[moe] CSV: shapes 5+7 bm=64, shape 6 bm=128, all FlyDSL t16", file=sys.stderr)
except Exception as e:
    print(f"[moe] CSV append failed: {e}", file=sys.stderr)

# Pre-compile FlyDSL stage2 MLIR for both inter_dim variants
try:
    from aiter.ops.flydsl.moe_kernels import _get_compiled_stage2
    _get_compiled_stage2(
        model_dim=7168, inter_dim=512, experts=33, topk=9,
        tile_m=16, tile_n=256, tile_k=128,
        doweight=False, a_dtype="fp4", b_dtype="fp4", out_dtype="bf16",
        accumulate=True,
    )
    _get_compiled_stage2(
        model_dim=7168, inter_dim=2048, experts=33, topk=9,
        tile_m=16, tile_n=256, tile_k=128,
        doweight=False, a_dtype="fp4", b_dtype="fp4", out_dtype="bf16",
        accumulate=True,
    )
    print("[moe] FlyDSL t16x256x128 MLIR pre-compiled for shapes 5,6,7", file=sys.stderr)
except Exception as e:
    print(f"[moe] FlyDSL t16 MLIR pre-compile failed: {e}", file=sys.stderr)

# GPU binary warmup for BOTH inter_dim variants
# Shape 5 runs before shape 6 in benchmark, so we must warmup inter_dim=512 here
try:
    from aiter.ops.flydsl.moe_kernels import flydsl_moe_stage2
    _dev = "cuda"
    _E = 33
    _model_dim = 7168
    _topk = 9
    _tile_m = 16

    for _inter_dim in (512, 2048):
        _inter_dim_packed = _inter_dim // 2
        _dummy_inter = torch.zeros((_tile_m, _topk, _inter_dim_packed), dtype=torch.uint8, device=_dev)
        _dummy_w2 = torch.zeros((_E, _model_dim, _inter_dim_packed), dtype=torch.uint8, device=_dev)
        _dummy_w2_scale = torch.ones((_E, _model_dim, _inter_dim // 32), dtype=torch.uint8, device=_dev)
        _dummy_a2_scale = torch.ones((_tile_m * _topk, _inter_dim // 32), dtype=torch.uint8, device=_dev)
        _dummy_sorted_ids = torch.arange(_tile_m, dtype=torch.int32, device=_dev)
        _dummy_expert_ids = torch.zeros(1, dtype=torch.int32, device=_dev)
        _dummy_num_valid = torch.tensor([_tile_m] + [0] * (_E - 1), dtype=torch.int32, device=_dev)

        flydsl_moe_stage2(
            _dummy_inter, _dummy_w2,
            _dummy_sorted_ids, _dummy_expert_ids, _dummy_num_valid,
            topk=_topk, tile_m=_tile_m, tile_n=256, tile_k=128,
            a_dtype="fp4", b_dtype="fp4", out_dtype="bf16", mode="atomic",
            w2_scale=_dummy_w2_scale, a2_scale=_dummy_a2_scale,
        )
        del _dummy_inter, _dummy_w2, _dummy_w2_scale, _dummy_a2_scale
        del _dummy_sorted_ids, _dummy_expert_ids, _dummy_num_valid

    torch.cuda.empty_cache()
    print("[moe] FlyDSL t16x256x128 GPU warmup complete (inter=512,2048)", file=sys.stderr)
except Exception as e:
    print(f"[moe] FlyDSL t16 warmup failed: {e}", file=sys.stderr)

# Warmup the HIP quant kernel
try:
    _dummy_hs = torch.randn(512, 7168, dtype=torch.bfloat16, device="cuda")
    _dummy_fp4, _dummy_scale = _quant_hip(_dummy_hs)
    del _dummy_hs, _dummy_fp4, _dummy_scale
    torch.cuda.empty_cache()
    print("[moe] HIP quant kernel warmed up", file=sys.stderr)
except Exception as e:
    print(f"[moe] HIP quant warmup failed: {e}", file=sys.stderr)

_PAD_CACHE = {}
_LAST_MODE = [None]
_environ = os.environ


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

    M = topk_ids.shape[0]
    E = gate_up_weight_shuffled.shape[0]

    cfg_id = id(config)
    cached = _PAD_CACHE.get(cfg_id)
    if cached is not None:
        hidden_pad, intermediate_pad = cached
    else:
        hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
        intermediate_pad = config["d_expert_pad"] - config["d_expert"]
        _PAD_CACHE[cfg_id] = (hidden_pad, intermediate_pad)

    # Per-shape mode selection:
    # - M<=16: always BYPASS (shapes 1,4)
    # - M<=128 + E>64: BYPASS (shape 2 with E=257)
    # - M<=128 + E<=64: 2-stage fp4 with CSV (shape 5 with E=33)
    # - M>128: 2-stage fp4 with CSV (shapes 3,6,7)
    if M <= 16 or (M <= 128 and E > 64):
        mode = 1  # BYPASS
    else:
        mode = 0  # 2-stage

    if mode != _LAST_MODE[0]:
        _LAST_MODE[0] = mode
        if mode:
            _environ["AITER_KSPLIT"] = "2"
            _environ["AITER_BYPASS_TUNE_CONFIG"] = "1"
        else:
            _environ["AITER_KSPLIT"] = "0"
            _environ["AITER_BYPASS_TUNE_CONFIG"] = "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,
        w1_scale=gate_up_weight_scale_shuffled,
        w2_scale=down_weight_scale_shuffled,
        hidden_pad=hidden_pad, intermediate_pad=intermediate_pad,
    )