# SPDX-License-Identifier: MIT
import torch
from typing import Optional
from ..jit.core import compile_ops

MD_NAME = "module_moe_sorting"


@compile_ops("module_moe_sorting")
def moe_sorting_fwd(
    topk_ids: torch.Tensor,
    topk_weights: torch.Tensor,
    sorted_token_ids: torch.Tensor,
    sorted_weights: torch.Tensor,
    sorted_expert_ids: torch.Tensor,
    tokens_positions_per_expert: torch.Tensor,
    num_valid_ids: torch.Tensor,
    moe_buf: Optional[torch.Tensor],
    num_experts: int,
    unit_size: int,
    local_expert_mask: Optional[torch.Tensor] = None,
) ->None: ...
