submission 754739

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

"""
v193: v186 + Triton quant kernel BLK_Mx=64 (smaller tiles, more parallelism).
"""

import os
import functools
import torch
import triton
from task import input_t, output_t

import aiter
from aiter import ActivationType, QuantType, dtypes
from aiter.fused_moe import get_2stage_cfgs, get_padded_M, get_inter_dim
import aiter.fused_moe as _fused_moe_module
import aiter.ops.flydsl.moe_kernels as _flydsl_moe_kernels
from aiter.ops.triton._triton_kernels.quant.fused_mxfp4_quant import (
    _fused_dynamic_mxfp4_quant_moe_sort_kernel,
)

_flydsl_moe_kernels._KERNEL_PARAMS["flydsl_moe2_afp4_wfp4_bf16_t16x128x128_atomic"] = {
    "stage": 2, "a_dtype": "fp4", "b_dtype": "fp4", "out_dtype": "bf16",
    "tile_m": 16, "tile_n": 128, "tile_k": 128, "mode": "atomic", "MPerBlock": 16,
}

_CUSTOM_CONFIGS = {}

def _make_key(token, inter_dim, expert):
    return (
        256, token, 7168, inter_dim, expert, 9,
        "ActivationType.Silu", "torch.bfloat16",
        "torch.float4_e2m1fn_x2", "torch.float4_e2m1fn_x2",
        "QuantType.per_1x32", True, False,
    )

_4WG = "moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
_FLY = "flydsl_moe2_afp4_wfp4_bf16_t16x128x128_atomic"

# E=257: bs=16/128 ksplit=2 (sparse, bf16), bs=512 ksplit=0 (fp4 with shared 4WG_M128)
_CUSTOM_CONFIGS[_make_key(16, 256, 257)] = {
    "block_m": 16, "ksplit": 2, "kernelName1": "", "kernelName2": "", "run_1stage": False,
}
_CUSTOM_CONFIGS[_make_key(128, 256, 257)] = {
    "block_m": 32, "ksplit": 0,
    "kernelName1": _4WG, "kernelName2": _FLY, "run_1stage": False,
}
_CUSTOM_CONFIGS[_make_key(512, 256, 257)] = {
    "block_m": 32, "ksplit": 0,
    "kernelName1": _4WG, "kernelName2": _FLY, "run_1stage": False,
}

# E=33: bs=16 ksplit=2, bs=128+ ksplit=0 (all use same 4WG_M128)
_CUSTOM_CONFIGS[_make_key(16, 512, 33)] = {
    "block_m": 32, "ksplit": 2, "kernelName1": "", "kernelName2": "", "run_1stage": False,
}
_CUSTOM_CONFIGS[_make_key(128, 512, 33)] = {
    "block_m": 32, "ksplit": 0,
    "kernelName1": _4WG, "kernelName2": _FLY, "run_1stage": False,
}
_CUSTOM_CONFIGS[_make_key(512, 512, 33)] = {
    "block_m": 64, "ksplit": 0,
    "kernelName1": _4WG, "kernelName2": _FLY, "run_1stage": False,
}  # block_m=64 wins here (dense experts, small K=512)
_CUSTOM_CONFIGS[_make_key(512, 2048, 33)] = {
    "block_m": 32, "ksplit": 0,
    "kernelName1": _4WG, "kernelName2": _FLY, "run_1stage": False,
}  # block_m=32 wins (large K=2048 amortizes block overhead)

# --- Buffer caches ---
_buf = {}
_qbuf = {}

def _get_sorting_bufs(M, E, topk, model_dim, block_m, device):
    key = (M, E, topk, model_dim, block_m)
    if key not in _buf:
        max_pad = M * topk + E * block_m - topk
        max_blk = (max_pad + block_m - 1) // block_m
        _buf[key] = {
            "sid": torch.empty(max_pad, dtype=dtypes.i32, device=device),
            "sw": torch.empty(max_pad, dtype=dtypes.fp32, device=device),
            "se": torch.empty(max_blk, dtype=dtypes.i32, device=device),
            "nv": torch.empty(2, dtype=dtypes.i32, device=device),
            "out": torch.empty((M, model_dim), dtype=torch.bfloat16, device=device),
            "a2": torch.empty((M, topk, 0), dtype=torch.bfloat16, device=device),
        }
    return _buf[key]

def _get_a2(M, topk, inter_dim, device):
    key = ("a2", M, topk, inter_dim)
    if key not in _buf:
        _buf[key] = torch.empty((M, topk, inter_dim), dtype=torch.bfloat16, device=device)
    return _buf[key]

def _quant_prealloc(x, sorted_ids, num_valid_ids, token_num, topk, block_m, device):
    M, N = x.shape
    QBS = 32
    BLK_Mx = 64
    BLK_M, BLK_N = 32, 8
    BLK_M_u32, BLK_N_u32 = 16, 4

    scaleN = triton.cdiv(N, QBS)
    M_o = sorted_ids.shape[0]

    qk = (M, N, M_o, topk)
    if qk not in _qbuf:
        _qbuf[qk] = {
            "fp4": torch.empty((M, N // 2), dtype=torch.uint8, device=device),
            "bs": torch.empty(
                (triton.cdiv(M_o, BLK_M), triton.cdiv(scaleN, BLK_N),
                 BLK_N_u32, BLK_M_u32, 4),
                dtype=torch.uint8, device=device,
            ),
        }
    qb = _qbuf[qk]

    num_pid = triton.cdiv(M, BLK_Mx) * scaleN + triton.cdiv(
        M_o, BLK_M
    ) * triton.cdiv(scaleN, BLK_N)

    _fused_dynamic_mxfp4_quant_moe_sort_kernel[(num_pid,)](
        x, qb["fp4"], sorted_ids, num_valid_ids, qb["bs"],
        M, N, scaleN,
        *x.stride(), *qb["fp4"].stride(), *qb["bs"].stride(),
        token_num=token_num, M_i=M, N_i=scaleN,
        MXFP4_QUANT_BLOCK_SIZE=QBS, BLOCK_SIZE_Mx=BLK_Mx,
        BLOCK_SIZE_M=BLK_M // 2, BLOCK_SIZE_N=BLK_N // 2,
        TOPK=topk,
    )

    return (
        qb["fp4"].view(dtypes.fp4x2),
        qb["bs"].view(dtypes.fp8_e8m0).view(-1, scaleN),
    )


_injected = False

def _inject():
    global _injected
    if _injected:
        return
    _injected = True
    if _fused_moe_module.cfg_2stages is None:
        import pandas as pd
        from aiter.jit.core import AITER_CONFIGS
        tf = AITER_CONFIGS.AITER_CONFIG_FMOE_FILE
        if os.path.exists(tf):
            _IDX = [
                "cu_num", "token", "model_dim", "inter_dim", "expert", "topk",
                "act_type", "dtype", "q_dtype_a", "q_dtype_w", "q_type",
                "use_g1u1", "doweight_stage1",
            ]
            df = pd.read_csv(tf)
            if "_tag" in df.columns:
                df = df[df["_tag"].fillna("") == ""]
            _fused_moe_module.cfg_2stages = df.set_index(_IDX).to_dict("index")
        else:
            _fused_moe_module.cfg_2stages = {}
    _fused_moe_module.cfg_2stages.update(_CUSTOM_CONFIGS)


def custom_kernel(data: input_t) -> output_t:
    (
        hidden_states, _w1r, _w2r, _w1sr, _w2sr,
        w1, w2, w1s, w2s,
        topk_weights, topk_ids, config,
    ) = data

    _inject()

    M = hidden_states.shape[0]
    topk = topk_ids.shape[1]
    device = hidden_states.device
    dhp = config["d_hidden_pad"]
    dep = config["d_expert_pad"]
    hidden_pad = dhp - config["d_hidden"]
    intermediate_pad = dep - config["d_expert"]

    E, model_dim, inter_dim = get_inter_dim(w1.shape, w2.shape)
    padded_M = get_padded_M(M)

    metadata = get_2stage_cfgs(
        padded_M, model_dim, inter_dim, E, topk,
        torch.bfloat16, dtypes.fp4x2, dtypes.fp4x2,
        QuantType.per_1x32, True, ActivationType.Silu,
        False, hidden_pad, intermediate_pad, True,
    )
    block_m = int(metadata.block_m)

    b = _get_sorting_bufs(M, E, topk, model_dim, block_m, device)
    aiter.moe_sorting_fwd(
        topk_ids, topk_weights,
        b["sid"], b["sw"], b["se"], b["nv"], b["out"],
        E, block_m, None, None, 0,
    )

    w1sv = w1s.view(dtypes.fp8_e8m0)
    w2sv = w2s.view(dtypes.fp8_e8m0)
    a2_buf = _get_a2(M, topk, inter_dim, device)

    if metadata.ksplit > 1:
        a1 = hidden_states.to(torch.bfloat16)
        a2 = metadata.stage1(
            a1, w1, w2, b["sid"], b["se"], b["nv"], a2_buf, topk,
            block_m=block_m, a1_scale=None, w1_scale=w1sv, sorted_weights=None,
        )
        metadata.stage2(
            a2, w1, w2, b["sid"], b["se"], b["nv"], b["out"], topk,
            w2_scale=w2sv, a2_scale=None, block_m=block_m, sorted_weights=b["sw"],
        )
    else:
        a1, a1s = _quant_prealloc(
            hidden_states, b["sid"], b["nv"], M, 1, block_m, device,
        )
        a2 = metadata.stage1(
            a1, w1, w2, b["sid"], b["se"], b["nv"], a2_buf, topk,
            block_m=block_m, a1_scale=a1s, w1_scale=w1sv, sorted_weights=None,
        )
        a2_flat = a2.view(-1, inter_dim)
        a2q, a2s = _quant_prealloc(
            a2_flat, b["sid"], b["nv"], M, topk, block_m, device,
        )
        a2q = a2q.view(M, topk, -1)
        metadata.stage2(
            a2q, w1, w2, b["sid"], b["se"], b["nv"], b["out"], topk,
            w2_scale=w2sv, a2_scale=a2s, block_m=block_m, sorted_weights=b["sw"],
        )

    return b["out"]