submission 716987

"""s64 + t32x256x256 for d512 stage2 (3µs improvement)."""
import os
os.environ["AITER_KSPLIT"]="2"
from task import input_t, output_t
import torch,functools,sys
from aiter import ActivationType,QuantType,dtypes
from aiter.fused_moe import fused_moe
from aiter.ops.triton.quant import dynamic_mxfp4_quant
_p=False;_c=False;_w=False
def _warm():
    global _w
    if _w:return
    _w=True
    try:dynamic_mxfp4_quant(torch.randn(4,64,dtype=torch.bfloat16,device='cuda'));torch.cuda.synchronize()
    except:pass
def _patch():
    global _p
    if _p:return
    _p=True
    try:
        import aiter.fused_moe as fm
        from aiter.utility import fp4_utils
        @functools.lru_cache(maxsize=2048)
        def sk(t,tp,e,i,m):
            if t>=512 and i>=512:return 0
            if e<=33 and(m%2==0)and((m//2)%256==0):return 2
            return 0
        fm.get_ksplit=sk
        quant_func=fm.get_quant(QuantType.per_1x32)
        def sqs(data,sorted_ids,num_valid_ids,token_num,topk,block_size):
            kw={};
            if topk>1:kw['num_rows_factor']=topk
            q_data,q_scale=quant_func(data,scale=None,quant_dtype=dtypes.fp4x2,num_rows=None,**kw)
            if topk>1:q_scale=fp4_utils.moe_mxfp4_sort(q_scale[:token_num*topk,:].view(token_num,topk,-1),sorted_ids=sorted_ids,num_valid_ids=num_valid_ids,token_num=token_num,block_size=block_size)
            else:q_scale=fp4_utils.moe_mxfp4_sort(q_scale,sorted_ids=sorted_ids,num_valid_ids=num_valid_ids,token_num=token_num,block_size=block_size)
            return q_data,q_scale
        fm.fused_dynamic_mxfp4_quant_moe_sort=sqs
    except:pass
def _inject():
    global _c
    if _c:return
    _c=True
    try:
        import aiter.fused_moe as fm
        if fm.cfg_2stages is None:return
        d={'block_m':32,'ksplit':2,'kernelName1':'','kernelName2':'','run_1stage':0,'us':0,'tflops':0,'bw':0}
        for t in[16,128]:
            for i in[512,2048]:
                fm.cfg_2stages[(256,t,7168,i,33,9,'ActivationType.Silu','torch.bfloat16','torch.float4_e2m1fn_x2','torch.float4_e2m1fn_x2','QuantType.per_1x32',True,False)]=d.copy()
        # d512/bs512: 256x32 + FlyDSL t32x256x256_reduce (3µs faster than t32x128x256)
        fm.cfg_2stages[(256,512,7168,512,33,9,'ActivationType.Silu','torch.bfloat16','torch.float4_e2m1fn_x2','torch.float4_e2m1fn_x2','QuantType.per_1x32',True,False)]={'block_m':32,'ksplit':0,
            'kernelName1':'moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16',
            'kernelName2':'flydsl_moe2_afp4_wfp4_bf16_t32x256x256_reduce','run_1stage':0,'us':0,'tflops':0,'bw':0}
        fm.cfg_2stages[(256,512,7168,2048,33,9,'ActivationType.Silu','torch.bfloat16','torch.float4_e2m1fn_x2','torch.float4_e2m1fn_x2','QuantType.per_1x32',True,False)]={'block_m':128,'ksplit':0,'kernelName1':'','kernelName2':'flydsl_moe2_afp4_wfp4_bf16_t32x128x256_reduce','run_1stage':0,'us':0,'tflops':0,'bw':0}
        fm.cfg_2stages[(256,512,7168,256,257,9,'ActivationType.Silu','torch.bfloat16','torch.float4_e2m1fn_x2','torch.float4_e2m1fn_x2','QuantType.per_1x32',True,False)]={
            'block_m':32,'ksplit':0,
            'kernelName1':'moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16',
            'kernelName2':'moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16',
            'run_1stage':0,'us':0,'tflops':0,'bw':0}
        fm.cfg_2stages[(256,128,7168,256,257,9,'ActivationType.Silu','torch.bfloat16','torch.float4_e2m1fn_x2','torch.float4_e2m1fn_x2','QuantType.per_1x32',True,False)]={
            'block_m':32,'ksplit':0,
            'kernelName1':'moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16',
            'kernelName2':'moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16',
            'run_1stage':0,'us':0,'tflops':0,'bw':0}
    except:pass
_c=False
def custom_kernel(data:input_t)->output_t:
    _warm()
    _patch()
    h,_,_,_,_,w1s,w2s,w1ss,w2ss,tw,ti,cfg=data
    r=fused_moe(h,w1s,w2s,tw,ti,expert_mask=None,activation=ActivationType.Silu,quant_type=QuantType.per_1x32,doweight_stage1=False,w1_scale=w1ss,w2_scale=w2ss,a1_scale=None,a2_scale=None,hidden_pad=cfg["d_hidden_pad"]-cfg["d_hidden"],intermediate_pad=cfg["d_expert_pad"]-cfg["d_expert"])
    _inject()
    return r