submission 689727

#!POPCORN leaderboard amd-moe-mxfp4
#!POPCORN gpu MI355X
"""
v141: Try 'atomic' mode for FlyDSL stage2 on S2 and S6.
- S2: t32x256x256 atomic (back to tile_n=256 since 128 was worse, try atomic)
- S5: t64x256x256 reduce (keep winning config from v139)
- S6: t64x256x256 atomic (try atomic for large K=2048)
Also compile both atomic and reduce for each so we can compare.
"""
import os, functools, torch
os.environ["PYTORCH_ROCM_ARCH"] = "gfx950"
os.environ["AITER_BYPASS_TUNE_CONFIG"] = "1"

from task import input_t, output_t
from aiter import ActivationType, QuantType

import aiter.fused_moe as _fm
import aiter

from aiter.ops.flydsl.moe_kernels import compile_flydsl_moe_stage2
from aiter.ops.flydsl import flydsl_moe_stage2

# Pre-compile FlyDSL stage2 variants
# S2: E=257, d=256 — atomic mode
compile_flydsl_moe_stage2(
    model_dim=7168, inter_dim=256, experts=257, topk=9,
    tile_m=32, tile_n=256, tile_k=256,
    doweight_stage2=True, a_dtype='fp4', b_dtype='fp4', out_dtype='bf16',
    accumulate=True,
)
# S5: E=33, d=512 — reduce (proven winner)
compile_flydsl_moe_stage2(
    model_dim=7168, inter_dim=512, experts=33, topk=9,
    tile_m=64, tile_n=256, tile_k=256,
    doweight_stage2=True, a_dtype='fp4', b_dtype='fp4', out_dtype='bf16',
    accumulate=True,
)
# S6: E=33, d=2048 — atomic mode
compile_flydsl_moe_stage2(
    model_dim=7168, inter_dim=2048, experts=33, topk=9,
    tile_m=64, tile_n=256, tile_k=256,
    doweight_stage2=True, a_dtype='fp4', b_dtype='fp4', out_dtype='bf16',
    accumulate=True,
)
print("FlyDSL stage2 pre-compiled", flush=True)


# FlyDSL stage2 wrappers
def _make_flydsl_stage2(tile_m, tile_n, tile_k=256, mode='reduce'):
    def _flydsl_s2(inter_states, w1, w2, sorted_token_ids, sorted_expert_ids,
                    num_valid_ids, out, topk, kernelName=None, w2_scale=None,
                    a2_scale=None, block_m=32, sorted_weights=None, **kwargs):
        flydsl_moe_stage2(
            inter_states=inter_states, w2=w2,
            sorted_token_ids=sorted_token_ids,
            sorted_expert_ids=sorted_expert_ids,
            num_valid_ids=num_valid_ids,
            out=out, topk=topk,
            tile_m=tile_m, tile_n=tile_n, tile_k=tile_k,
            a_dtype='fp4', b_dtype='fp4', out_dtype='bf16',
            mode=mode,
            w2_scale=w2_scale, a2_scale=a2_scale,
            sorted_weights=sorted_weights,
        )
    return _flydsl_s2

# Specific wrappers
_flydsl_s2_S2 = _make_flydsl_stage2(32, 256, 256, 'atomic')   # S2: atomic
_flydsl_s2_S5 = _make_flydsl_stage2(64, 256, 256, 'reduce')   # S5: reduce (proven)
_flydsl_s2_S6 = _make_flydsl_stage2(64, 256, 256, 'atomic')   # S6: atomic


# Standard ksplit/bm config
@functools.lru_cache(maxsize=2048)
def _custom_get_ksplit(token, topk, expert, inter_dim, model_dim):
    if token <= 128:
        return 2
    return 0

_fm.get_ksplit = _custom_get_ksplit

_orig_get_block_size_M = getattr(_fm.get_block_size_M, '__wrapped__', _fm.get_block_size_M)

@functools.lru_cache(maxsize=2048)
def _custom_get_block_size_M(token, topk, expert, inter_dim):
    if token <= 128:
        return 16
    if expert > 128:
        return 32
    if inter_dim > 1024:
        return 64
    return _orig_get_block_size_M(token, topk, expert, inter_dim)

_fm.get_block_size_M = _custom_get_block_size_M


# Monkey-patch: inject FlyDSL stage2 per-shape
_orig_get_2stage_cfgs = _fm.get_2stage_cfgs

def _patched_get_2stage_cfgs(*args, **kwargs):
    cfg = _orig_get_2stage_cfgs(*args, **kwargs)
    if cfg is None:
        return cfg

    if cfg.ksplit == 0:
        # args[0]=token, args[2]=inter_dim, args[3]=expert
        token = args[0]
        inter_dim = args[2]
        expert = args[3]

        if expert > 128:  # E=257 (S2)
            cfg.stage2 = _flydsl_s2_S2
        elif inter_dim > 1024:  # E=33, d=2048 (S6)
            cfg.stage2 = _flydsl_s2_S6
        else:  # E=33, d=512 (S5)
            cfg.stage2 = _flydsl_s2_S5

    return cfg

_fm.get_2stage_cfgs = _patched_get_2stage_cfgs

from aiter.fused_moe import fused_moe


def custom_kernel(data: input_t) -> output_t:
    (hs, rw1, rw2, rw1s, rw2s, w1, w2, w1s, w2s, tw, ti, cfg) = data

    hp = cfg.get("d_hidden_pad", cfg["d_hidden"]) - cfg["d_hidden"]
    ip = cfg.get("d_expert_pad", cfg["d_expert"]) - cfg["d_expert"]

    return fused_moe(
        hs, w1, w2, tw, ti,
        expert_mask=None, activation=ActivationType.Silu,
        quant_type=QuantType.per_1x32, doweight_stage1=False,
        w1_scale=w1s, w2_scale=w2s, a1_scale=None, a2_scale=None,
        hidden_pad=hp, intermediate_pad=ip,
    )