submission 722866

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

def custom_kernel(data):
    (
        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.size(0)
    hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
    intermediate_pad = config["d_expert_pad"] - config["d_expert"]
    routed_top_k = config["total_top_k"] - config["n_shared_experts"]

    # Small batches (including pre‑check) – no sorting, direct AITER call
    if M <= 8:
        return 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,
        )

    # Large batches – sort tokens by first expert for better L2 cache locality
    if routed_top_k > 0:
        first_expert = topk_ids[:, 0]
        sorted_idx = torch.argsort(first_expert)
        inv_sorted_idx = torch.argsort(sorted_idx)
        hidden_sorted = hidden_states[sorted_idx]
        weights_sorted = topk_weights[sorted_idx]
        ids_sorted = topk_ids[sorted_idx]
    else:
        hidden_sorted = hidden_states
        weights_sorted = topk_weights
        ids_sorted = topk_ids
        inv_sorted_idx = torch.arange(M, device=hidden_states.device)

    output_sorted = fused_moe(
        hidden_sorted,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        weights_sorted,
        ids_sorted,
        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,
    )
    return output_sorted[inv_sorted_idx]