submission 664132

#!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
# The fused Triton kernel does quant+sort in one dispatch (~58µs for M=512)
# Separated path: HIP quant (~30µs) + Triton sort (~15µs) = potentially faster
def _split_quant_moe_sort(x, sorted_ids, num_valid_ids, token_num, topk, block_size=32, scaling_mode="even"):
    # Step 1: HIP quant (bf16 → fp4x2 + scale)
    x_fp4, scale = _quant_hip(x)
    # Step 2: Sort scale by MoE expert assignment
    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 (not in default stage2 tile_ms)
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} in _KERNEL_PARAMS", file=sys.stderr)
except Exception as e:
    print(f"[moe] _KERNEL_PARAMS registration failed: {e}", file=sys.stderr)

# CSV injection: shapes 6+7 with FlyDSL t16 stage2, block_m=64 for both
_tune_file = _fm.AITER_CONFIGS.AITER_CONFIG_FMOE_FILE
try:
    with open(_tune_file, "a") as f:
        # Shape 6: E=33, inter=512, block_m=64
        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,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 7: 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 appended: shapes 6+7 FlyDSL t16, block_m=64", file=sys.stderr)
except Exception as e:
    print(f"[moe] CSV append failed: {e}", file=sys.stderr)

# Pre-compile FlyDSL stage2 MLIR for both shapes
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 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: trigger @flyc.jit for t16x256x128 (shape 7, largest)
try:
    from aiter.ops.flydsl.moe_kernels import flydsl_moe_stage2
    _dev = "cuda"
    _E = 33
    _model_dim = 7168
    _inter_dim = 2048
    _inter_dim_packed = 1024
    _tile_m = 16
    _topk = 9

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

    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)

    # Only toggle env vars when M mode changes
    mode = M <= 128
    if mode != _LAST_M_MODE[0]:
        _LAST_M_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,
    )