submission 564494

import torch
from task import input_t, output_t

import aiter
from aiter import ActivationType, QuantType, dtypes
from aiter.fused_moe import fused_moe

class MetaWrapper:
    def __init__(self, orig, splitk):
        self._orig = orig
        self._splitk = splitk
        
    def __getattr__(self, name):
        if name == "splitk":
            return self._splitk
        return getattr(self._orig, name)

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

    hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
    intermediate_pad = config["d_expert_pad"] - config["d_expert"]
    M = hidden_states.shape[0]

    old_get_2stage_cfgs = aiter.fused_moe.get_2stage_cfgs
    
    def mocked_get_2stage_cfgs(*args, **kwargs):
        meta = old_get_2stage_cfgs(*args, **kwargs)
        
        E = args[3] if len(args) > 3 else 257
        block_m = meta.block_m if hasattr(meta, "block_m") and meta.block_m is not None and meta.block_m != -1 else 32
        
        num_blocks = ((M + block_m - 1) // block_m) * E
        
        if num_blocks > 0:
            desired_splitk = 1200 // num_blocks
        else:
            desired_splitk = 1
            
        if desired_splitk >= 8:
            splitk = 8
        elif desired_splitk >= 4:
            splitk = 4
        elif desired_splitk >= 2:
            splitk = 2
        else:
            splitk = 1
            
        return MetaWrapper(meta, splitk)

    try:
        aiter.fused_moe.get_2stage_cfgs = mocked_get_2stage_cfgs

        output = fused_moe(
            hidden_states,
            gate_up_weight_shuffled,
            down_weight_shuffled,
            topk_weights,
            topk_ids,
            expert_mask=None,
            activation=ActivationType.Silu,
            quant_type=QuantType.per_1x32,
            doweight_stage1=False,
            w1_scale=gate_up_weight_scale_shuffled,
            w2_scale=down_weight_scale_shuffled,
            a1_scale=None,
            a2_scale=None,
            hidden_pad=hidden_pad,
            intermediate_pad=intermediate_pad,
        )
    finally:
        aiter.fused_moe.get_2stage_cfgs = old_get_2stage_cfgs

    return output