Diff Coverage

def _local_expert_major_forward(experts, dispatched_states, local_expert_counts):
    """Run grouped experts for an already expert-major token stream."""
    grouped_forward = getattr(experts, "forward_expert_major", None)
    if not callable(grouped_forward):
        raise TypeError(
            f"{type(experts).__name__}: NPU token dispatch requires "
            "experts.forward_expert_major"
        )
    return grouped_forward(dispatched_states, local_expert_counts)

            )
    if apply_gate is not None and use_grouped_gemm:
        raise ValueError("custom expert gate activation is not supported by grouped GEMM")
    if use_npu_moe_token_dispatch and not use_grouped_gemm:
        raise ValueError("NPU token dispatch requires grouped GEMM")
    if use_npu_moe_token_dispatch and not callable(
        getattr(module.experts, "forward_expert_major", None)
    ):
        raise TypeError(
            f"{type(module.experts).__name__}: NPU token dispatch requires "
            "experts.forward_expert_major"
        )
    # ``_ep_*`` is the contract the local compute path reads back off the module's

    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 _reorder_variable_chunks(
    tensor: torch.Tensor,

    chunk_count: int,
) -> torch.Tensor:
    """Validate and run the optional NPU token dispatcher."""
    if chunk_count > 1:
        raise ValueError("NPU token dispatch does not support HP_EP_DISPATCH_CHUNKS > 1")
    if _resolve_capacity_factor(module) is not None:
        raise ValueError("NPU token dispatch does not support expert capacity limits")
    topk_indices, topk_weights = routing
    return _npu_token_dispatch_forward(
        module,
        hidden_states,