Diff Coverage

    tokens: torch.Tensor,
    indices: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Load the optional NPU fused permutation only when requested."""
    from hyper_parallel.components.functional.moe_token_permute import (  # pylint: disable=C0415
        moe_token_permute,
    )

    return moe_token_permute(tokens, indices)


def _npu_moe_token_unpermute(
    permuted_tokens: torch.Tensor,

    sorted_indices: torch.Tensor,
    probabilities: torch.Tensor,
) -> torch.Tensor:
    """Load the optional NPU fused unpermutation only when requested."""
    from hyper_parallel.components.functional.moe_token_unpermute import (  # pylint: disable=C0415
        moe_token_unpermute,
    )

    return moe_token_unpermute(permuted_tokens, sorted_indices, probabilities)


def _npu_token_dispatch_enabled(hidden_states: torch.Tensor, enabled: bool) -> bool:
    """Use NPU routing kernels only for NPU tensors when enabled."""