submission 654867

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

# v69: Ksplit hack ONLY (no OPUS sorting) — isolate ksplit effect from v56
# v56 showed S3 -5µs, S4 -5µs improvement from ksplit=2 but S0/S1/S6 regressed
# Hypothesis: OPUS sorting regresses ne=257 shapes, ksplit helps ne=33 small shapes
# This version: env vars + ksplit hack only

import os
os.environ["HIP_FORCE_DEV_KERNARG"] = "1"
os.environ["AMD_DIRECT_DISPATCH"] = "1"
os.environ["GPU_MAX_HW_QUEUES"] = "2"
os.environ["HSA_NO_SCRATCH_RECLAIM"] = "1"

from task import input_t, output_t

_CU=256;_TM=2*_CU;_TK=256;_EM=33;_KC=(2,4,8)
def _bm(t): return 32 if t<=32 else (64 if t<=64 else 128)
def _et(t,k,e,d):
    p=max(1,(t*k)//e);m=max(1,(p+_bm(p)-1)//_bm(p));return e*m*((2*d+127)//128)
def _vk(m): return tuple(k for k in _KC if m%k==0 and (m//k)%_TK==0)
def _pk(m,t):
    if t>=_TM: return 0
    v=_vk(m)
    if not v: return 0
    for k in v:
        if t*k>=_TM: return k
    return v[-1] if t<_TM//2 else 0
try:
    import torch
    from aiter import ActivationType, QuantType
    import aiter.fused_moe as _fmoe
    _ACT=ActivationType.Silu;_QUANT=QuantType.per_1x32;_fused_moe=_fmoe.fused_moe
    _o=getattr(_fmoe,"get_ksplit",None)
    if callable(_o):
        def _pk2(t,k,e,d,m):
            r=_o(t,k,e,d,m)
            if r!=0 or e<=0 or e>_EM: return r
            b=_pk(m,_et(t,k,e,d));return b if b>0 else r
        _fmoe.get_ksplit=_pk2
    _gc=getattr(_fmoe,"get_2stage_cfgs",None)
    if _gc and hasattr(_gc,"cache_clear"): _gc.cache_clear()
    if _gc:
        dt=torch.bfloat16;qd=torch.float4_e2m1fn_x2
        for a,b,c,d,e in [(16,7168,256,257,9),(128,7168,256,257,9),(512,7168,256,257,9),(16,7168,512,33,9),(128,7168,512,33,9),(512,7168,512,33,9),(512,7168,2048,33,9)]:
            try: _gc(a,b,c,d,e,dt,qd,qd,_QUANT,True,_ACT,False,0,0,True)
            except: pass
    _INIT_OK=True
except:
    _INIT_OK=False;_fused_moe=None;_ACT=None;_QUANT=None

def custom_kernel(data: input_t) -> output_t:
    if not _INIT_OK:
        from aiter import ActivationType, QuantType
        from aiter.fused_moe import fused_moe
        h,_,_,_,_,w1,w2,s1,s2,tw,ti,c=data
        return fused_moe(h,w1,w2,tw,ti,expert_mask=None,activation=ActivationType.Silu,quant_type=QuantType.per_1x32,doweight_stage1=False,w1_scale=s1,w2_scale=s2,a1_scale=None,a2_scale=None,hidden_pad=c["d_hidden_pad"]-c["d_hidden"],intermediate_pad=c["d_expert_pad"]-c["d_expert"])
    h,_,_,_,_,w1,w2,s1,s2,tw,ti,c=data
    return _fused_moe(h,w1,w2,tw,ti,expert_mask=None,activation=_ACT,quant_type=_QUANT,doweight_stage1=False,w1_scale=s1,w2_scale=s2,a1_scale=None,a2_scale=None,hidden_pad=c["d_hidden_pad"]-c["d_hidden"],intermediate_pad=c["d_expert_pad"]-c["d_expert"])