submission 743515

import torch
from task import input_t, output_t

import aiter
from aiter import ActivationType, QuantType, dtypes
from aiter.fused_moe import (
    moe_sorting,
    get_inter_dim,
    get_2stage_cfgs,
    get_padded_M,
)
from aiter.ops.triton.quant.fused_mxfp4_quant import fused_dynamic_mxfp4_quant_moe_sort
import aiter.fused_moe as _fm


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

    M            = hidden_states.shape[0]
    E            = gate_up_weight_shuffled.shape[0]
    top_k        = topk_ids.shape[1]
    dtype        = hidden_states.dtype
    device       = hidden_states.device
    d_hidden     = config["d_hidden"]
    d_hidden_pad = config["d_hidden_pad"]
    d_expert_pad = config["d_expert_pad"]
    hidden_pad   = d_hidden_pad - d_hidden
    inter_pad    = d_expert_pad - config["d_expert"]

    _, model_dim, inter_dim = get_inter_dim(
        gate_up_weight_shuffled.shape,
        down_weight_shuffled.shape,
    )

    metadata = get_2stage_cfgs(
        get_padded_M(M), model_dim, inter_dim, E, top_k,
        dtype, dtypes.fp4x2, dtypes.fp4x2,
        QuantType.per_1x32, True, ActivationType.Silu,
        False, hidden_pad, inter_pad, True,
    )
    block_m = int(metadata.block_m)

    sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, moe_buf = moe_sorting(
        topk_ids, topk_weights, E, model_dim, dtype, block_m,
    )

    a1, a1_scale = fused_dynamic_mxfp4_quant_moe_sort(
        hidden_states,
        sorted_ids=sorted_ids, num_valid_ids=num_valid_ids,
        token_num=M, topk=1, block_size=block_m,
    )

    a2 = torch.empty((M, top_k, inter_dim), dtype=dtype, device=device)
    a2 = metadata.stage1(
        a1, gate_up_weight_shuffled, down_weight_shuffled,
        sorted_ids, sorted_expert_ids, num_valid_ids, a2, top_k,
        block_m=block_m, a1_scale=a1_scale,
        w1_scale=gate_up_weight_scale_shuffled, sorted_weights=None,
    )

    a2_flat = a2.view(-1, inter_dim)
    a2_q, a2_scale = fused_dynamic_mxfp4_quant_moe_sort(
        a2_flat, sorted_ids=sorted_ids, num_valid_ids=num_valid_ids,
        token_num=M, topk=top_k, block_size=block_m,
    )
    a2_q = a2_q.view(M, top_k, -1)

    metadata.stage2(
        a2_q, gate_up_weight_shuffled, down_weight_shuffled,
        sorted_ids, sorted_expert_ids, num_valid_ids, moe_buf, top_k,
        block_m=block_m, w2_scale=down_weight_scale_shuffled,
        a2_scale=a2_scale, sorted_weights=sorted_weights,
    )

    if hidden_pad > 0:
        return moe_buf[:, :d_hidden].contiguous()
    return moe_buf